[feat][sdk] Implement BruteForce index.

wt0530 · Dec 8, 2023 · b648e90 · b648e90
1 parent faa349c
commit b648e90
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 2 deletions.
diff --git a/java/dingo-sdk/src/main/java/io/dingodb/sdk/common/index/BruteForceParam.java b/java/dingo-sdk/src/main/java/io/dingodb/sdk/common/index/BruteForceParam.java
@@ -0,0 +1,18 @@
+package io.dingodb.sdk.common.index;
+
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+import lombok.ToString;
+
+@Getter
+@ToString
+@EqualsAndHashCode
+@NoArgsConstructor
+@AllArgsConstructor
+public class BruteForceParam {
+
+    private Integer dimension;
+    private VectorIndexParameter.MetricType metricType;
+}
diff --git a/java/dingo-sdk/src/main/java/io/dingodb/sdk/common/index/VectorIndexParameter.java b/java/dingo-sdk/src/main/java/io/dingodb/sdk/common/index/VectorIndexParameter.java
@@ -33,6 +33,7 @@ public class VectorIndexParameter {
     private IvfPqParam ivfPqParam;
     private HnswParam hnswParam;
     private DiskAnnParam diskAnnParam;
+    private BruteForceParam bruteForceParam;
 
     public VectorIndexParameter(VectorIndexType vectorIndexType, FlatParam flatParam) {
         this.vectorIndexType = vectorIndexType;
@@ -59,13 +60,19 @@ public VectorIndexParameter(VectorIndexType vectorIndexType, DiskAnnParam diskAn
         this.diskAnnParam = diskAnnParam;
     }
 
+    public VectorIndexParameter(VectorIndexType vectorIndexType, BruteForceParam bruteForceParam) {
+        this.vectorIndexType = vectorIndexType;
+        this.bruteForceParam = bruteForceParam;
+    }
+
     public enum VectorIndexType {
         VECTOR_INDEX_TYPE_NONE,
         VECTOR_INDEX_TYPE_FLAT,
         VECTOR_INDEX_TYPE_IVF_FLAT,
         VECTOR_INDEX_TYPE_IVF_PQ,
         VECTOR_INDEX_TYPE_HNSW,
-        VECTOR_INDEX_TYPE_DISKANN
+        VECTOR_INDEX_TYPE_DISKANN,
+        VECTOR_INDEX_TYPE_BRUTEFORCE
     }
 
     public enum MetricType {

diff --git a/java/dingo-sdk/src/main/java/io/dingodb/sdk/common/utils/EntityConversion.java b/java/dingo-sdk/src/main/java/io/dingodb/sdk/common/utils/EntityConversion.java
@@ -40,6 +40,7 @@
 import io.dingodb.sdk.common.codec.CodecUtils;
 import io.dingodb.sdk.common.codec.DingoKeyValueCodec;
 import io.dingodb.sdk.common.codec.KeyValueCodec;
+import io.dingodb.sdk.common.index.BruteForceParam;
 import io.dingodb.sdk.common.index.DiskAnnParam;
 import io.dingodb.sdk.common.index.FlatParam;
 import io.dingodb.sdk.common.index.HnswParam;
@@ -464,7 +465,14 @@ public static IndexParameter mapping(Common.IndexParameter parameter) {
                                     new DiskAnnParam(
                                             annParam.getDimension(),
                                             VectorIndexParameter.MetricType.valueOf(annParam.getMetricType().name()))));
-
+                case VECTOR_INDEX_TYPE_BRUTEFORCE:
+                    Common.CreateBruteForceParam bruteForceParam = vectorParam.getBruteforceParameter();
+                    return new IndexParameter(
+                            IndexParameter.IndexType.valueOf(parameter.getIndexType().name()),
+                            new VectorIndexParameter(VectorIndexParameter.VectorIndexType.VECTOR_INDEX_TYPE_BRUTEFORCE,
+                                    new BruteForceParam(
+                                            bruteForceParam.getDimension(),
+                                            VectorIndexParameter.MetricType.valueOf(bruteForceParam.getMetricType().name()))));
                 default:
                     throw new IllegalStateException("Unexpected value: " + vectorParam.getVectorIndexType());
             }
@@ -530,6 +538,12 @@ public static Common.IndexParameter mapping(IndexParameter parameter) {
                             .setMetricType(Common.MetricType.valueOf(diskAnnParam.getMetricType().name()))
                             .build());
                     break;
+                case VECTOR_INDEX_TYPE_BRUTEFORCE:
+                    BruteForceParam bruteForceParam = vectorParameter.getBruteForceParam();
+                    build.setBruteforceParameter(Common.CreateBruteForceParam.newBuilder()
+                            .setDimension(bruteForceParam.getDimension())
+                            .setMetricType(Common.MetricType.valueOf(bruteForceParam.getMetricType().name()))
+                            .build());
             }
             builder.setVectorIndexParameter(build.build());
         } else {

diff --git a/java/dingo-sdk/src/main/java/io/dingodb/sdk/common/vector/VectorSearchParameter.java b/java/dingo-sdk/src/main/java/io/dingodb/sdk/common/vector/VectorSearchParameter.java
@@ -38,6 +38,25 @@ public class VectorSearchParameter {
     private Coprocessor coprocessor;
     private List<Long> vectorIds;
 
+    private boolean useBruteForce;
+
+    @Deprecated
+    public VectorSearchParameter(
+            Integer topN,
+            boolean withoutVectorData,
+            boolean withoutScalarData,
+            List<String> selectedKeys,
+            boolean withoutTableData,
+            Search search,
+            VectorFilter vectorFilter,
+            VectorFilterType vectorFilterType,
+            Coprocessor coprocessor,
+            List<Long> vectorIds
+    ) {
+        this(topN, withoutVectorData, withoutScalarData, selectedKeys, withoutTableData, search,
+                vectorFilter, vectorFilterType, coprocessor, vectorIds, false);
+    }
+
     public enum VectorFilter {
         SCALAR_FILTER,
         TABLE_FILTER,