com.aerospike.vector.client.Conversions Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of avs-client-java Show documentation
This project includes the Java client for Aerospike Vector Search for high-performance data interactions.
The newest version!
package com.aerospike.vector.client;

import com.aerospike.vector.client.proto.*;
import com.google.common.base.Preconditions;
import com.google.protobuf.ByteString;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Utility class to convert data suitable for AVS. User would rarely need to invoke them directly.
 */
public class Conversions {
    private static final Logger log = LoggerFactory.getLogger(Conversions.class);

    private Conversions() {}

    /**
     * Constructs a record key using Aerospike namespace, set and key. User may not need to these
     * function directly Key must be of type String, byte[], Integer or Long.
     *
     * @param namespace Aerospike namespace.
     * @param set Aerospike set.
     * @param key key, must be of type String, byte[], Integer ot Long.
     * @return key {@link Key}.
     */
    public static Key buildKey(String namespace, @Nullable String set, Object key) {

        Preconditions.checkArgument(
                namespace != null && !namespace.isBlank(), "namespace must not be empty.");
        Preconditions.checkArgument(key != null, "key must not be null.");

        Key.Builder builder = Key.newBuilder().setNamespace(namespace);
        if (set != null) {
            builder.setSet(set);
        }

        if (key instanceof String) {
            builder.setStringValue((String) key);
        } else if (key instanceof byte[]) {
            builder.setBytesValue(ByteString.copyFrom((byte[]) key));
        } else if (key instanceof Integer) {
            builder.setIntValue((Integer) key);
        } else if (key instanceof Long) {
            builder.setLongValue((Long) key);
        } else {
            throw new IllegalArgumentException("Unsupported key type: " + key.getClass().getName());
        }

        return builder.build();
    }

    /**
     * Constructs a {@link Field} using bin name and value
     *
     * @param name bin name in the aerospike.
     * @param value corresponding value in the bin.
     * @return {@link Field}.
     */
    public static com.aerospike.vector.client.proto.Field buildField(
            String name, @Nullable Object value) {
        Preconditions.checkArgument(
                name != null && !name.isBlank(), "field name must not be empty.");

        return com.aerospike.vector.client.proto.Field.newBuilder()
                .setName(name)
                .setValue(toVectorDbValue(value))
                .build();
    }

    private static Value toVectorDbValue(@Nullable Object value) {
        if (value == null) {
            return null;
        }
        log.debug("buildValue value: {} class type: {}.", value, value.getClass().getName());

        Value.Builder builder = Value.newBuilder();
        if (value instanceof String) {
            builder.setStringValue((String) value);
        } else if (value instanceof byte[]) {
            builder.setBytesValue(ByteString.copyFrom((byte[]) value));
        } else if (value instanceof Integer) {
            builder.setIntValue((Integer) value);
        } else if (value instanceof Long) {
            builder.setLongValue((Long) value);
        } else if (value instanceof Float) {
            builder.setFloatValue((Float) value);
        } else if (value instanceof Double) {
            builder.setDoubleValue((Double) value);
        } else if (value instanceof float[]) {
            builder.setVectorValue(buildVectorValue(value));
        } else if (value instanceof double[]) {
            builder.setVectorValue(buildVectorValue(value));
        } else if (value instanceof boolean[]) {
            builder.setVectorValue(buildVectorValue(value));
        } else if (value instanceof List) {
            builder.setListValue(buildListValue((List) value));
        } else if (value instanceof Map) {
            builder.setMapValue(buildMapValue((Map) value));
        } else {
            throw new IllegalArgumentException(
                    "Unsupported value type: " + value.getClass().getName());
        }
        return builder.build();
    }

    private static com.aerospike.vector.client.proto.List buildListValue(List list) {
        Preconditions.checkArgument(list != null, "list must not be null.");
        return com.aerospike.vector.client.proto.List.newBuilder()
                .addAllEntries(
                        list.stream()
                                .filter(Objects::nonNull)
                                .map(Conversions::toVectorDbValue)
                                .collect(Collectors.toList()))
                .build();
    }

    private static com.aerospike.vector.client.proto.Map buildMapValue(Map map) {
        Preconditions.checkArgument(map != null, "map must not be null.");
        return com.aerospike.vector.client.proto.Map.newBuilder()
                .addAllEntries(
                        map.entrySet().stream()
                                .map(
                                        entry ->
                                                MapEntry.newBuilder()
                                                        .setKey(toMapKey(entry.getKey()))
                                                        .setValue(toVectorDbValue(entry.getValue()))
                                                        .build())
                                .collect(Collectors.toList()))
                .build();
    }

    /**
     * Constructs a {@link Vector} from an array of booleans or floats, used for vector search. This
     * method supports creating vectors from either float or boolean arrays by detecting the type of
     * the input array. Depending on the type, it constructs the appropriate {@link Vector} using
     * either {@link FloatData} or {@link BoolData}.
     *
     * @param array an array of floats or booleans to be converted into a {@link Vector}.
     * @return a {@link Vector} containing the data from the provided array.
     * @throws RuntimeException if the array type is neither float[] nor boolean[].
     */
    public static Vector buildVectorValue(Object array) {
        Preconditions.checkArgument(array != null, "array must not be null.");
        if (array instanceof float[] floats) {

            Vector.Builder vectorBuilder = Vector.newBuilder();
            FloatData.Builder floatDataBuilder = FloatData.newBuilder();

            for (float value : floats) {
                floatDataBuilder.addValue(value);
            }
            vectorBuilder.setFloatData(floatDataBuilder);
            return vectorBuilder.build();
        } else if (array instanceof boolean[] bools) {

            Vector.Builder vectorBuilder = Vector.newBuilder();
            BoolData.Builder booleanDataBuilder = BoolData.newBuilder();

            for (boolean value : bools) {
                booleanDataBuilder.addValue(value);
            }
            vectorBuilder.setBoolData(booleanDataBuilder);
            return vectorBuilder.build();
        } else {
            log.error("Found unexpected type :{}.", array.getClass().getName());
            throw new RuntimeException("Found unexpected type.");
        }
    }

    private static MapKey toMapKey(Object key) {
        Preconditions.checkArgument(key != null, "map key must not be null.");

        MapKey.Builder builder = MapKey.newBuilder();
        if (key instanceof String) {
            builder.setStringValue((String) key);
        } else if (key instanceof Integer) {
            builder.setIntValue((Integer) key);
        } else if (key instanceof Long) {
            builder.setLongValue((Long) key);
        } else if (key instanceof Float) {
            builder.setFloatValue((Float) key);
        } else if (key instanceof Double) {
            builder.setDoubleValue((Double) key);
        } else {
            throw new IllegalArgumentException(
                    "Unsupported map key type: " + key.getClass().getName());
        }
        return builder.build();
    }

    /**
     * Converts a Protobuf MapKey object to a Java object.
     *
     * @param mapKey The Protobuf MapKey object to convert.
     * @return The Java object representing the key.
     */
    public static Object fromMapKey(MapKey mapKey) {
        Preconditions.checkArgument(mapKey != null, "MapKey must not be null.");

        if (mapKey.hasStringValue()) {
            return mapKey.getStringValue();
        } else if (mapKey.hasIntValue()) {
            return mapKey.getIntValue();
        } else if (mapKey.hasLongValue()) {
            return mapKey.getLongValue();
        } else if (mapKey.hasFloatValue()) {
            return mapKey.getFloatValue();
        } else if (mapKey.hasDoubleValue()) {
            return mapKey.getDoubleValue();
        } else {
            throw new IllegalStateException("Unsupported map key type.");
        }
    }

    /**
     * Converts the data extracted from Vector DB {@link Value} to standard java types.
     *
     * @param value {@link Value} types.
     * @return java types extracted from {@link Value} types.
     */
    public static Object fromVectorDBValue(@Nullable Value value) {

        if (value == null) {
            log.info("Found null Vector value.");
            return null;
        }

        Object local = null;
        if (value.hasStringValue()) {
            local = value.getStringValue();
        } else if (value.hasBytesValue()) {
            local = value.getBytesValue().toByteArray();
        } else if (value.hasIntValue()) {
            local = value.getIntValue();
        } else if (value.hasLongValue()) {
            local = value.getLongValue();
        } else if (value.hasFloatValue()) {
            local = value.getFloatValue();
        } else if (value.hasDoubleValue()) {
            local = value.getDoubleValue();
        } else if (value.hasMapValue()) {
            local =
                    value.getMapValue().getEntriesList().stream()
                            .collect(
                                    Collectors.toMap(
                                            entry -> fromMapKey(entry.getKey()),
                                            entry -> fromVectorDBValue(entry.getValue())));
        } else if (value.hasListValue()) {
            local =
                    value.getListValue().getEntriesList().stream()
                            .map(Conversions::fromVectorDBValue)
                            .collect(Collectors.toList());
        } else {
            if (value.hasVectorValue()) {
                if (value.getVectorValue().hasFloatData()) {
                    local =
                            value.getVectorValue().getFloatData().getValueList().stream()
                                    .collect(Collectors.toList());
                } else {
                    local = new ArrayList<>(value.getVectorValue().getBoolData().getValueList());
                }

            } else {
                log.warn("Found null in Value: {} object.", value);
            }
        }
        return local;
    }
}