make indexType and metricType enum

zilliztech · Jul 12, 2024 · caa5b2b · caa5b2b
1 parent 6a867dd
commit caa5b2b
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 4 deletions.
diff --git a/seatunnel-api/src/main/java/org/apache/seatunnel/api/table/catalog/VectorIndex.java b/seatunnel-api/src/main/java/org/apache/seatunnel/api/table/catalog/VectorIndex.java
@@ -31,14 +31,22 @@ public class VectorIndex extends ConstraintKey.ConstraintKeyColumn implements Se
     private final String indexName;
 
     /** Vector indexType, such as IVF_FLAT, HNSW, DISKANN */
-    private final String indexType;
+    private final IndexType indexType;
 
     /** Vector index metricType, such as L2, IP, COSINE */
-    private final String metricType;
+    private final MetricType metricType;
 
     public VectorIndex(String indexName, String columnName, String indexType, String metricType) {
         super(columnName, null);
         this.indexName = indexName;
+        this.indexType = IndexType.of(indexType);
+        this.metricType = MetricType.of(metricType);
+    }
+
+    public VectorIndex(
+            String indexName, String columnName, IndexType indexType, MetricType metricType) {
+        super(columnName, null);
+        this.indexName = indexName;
         this.indexType = indexType;
         this.metricType = metricType;
     }
@@ -47,4 +55,56 @@ public VectorIndex(String indexName, String columnName, String indexType, String
     public ConstraintKey.ConstraintKeyColumn copy() {
         return new VectorIndex(indexName, getColumnName(), indexType, metricType);
     }
+
+    public enum IndexType {
+        FLAT,
+        IVF_FLAT,
+        IVF_SQ8,
+        IVF_PQ,
+        HNSW,
+        DISKANN,
+        AUTOINDEX,
+        SCANN,
+
+        // GPU indexes only for float vectors
+        GPU_IVF_FLAT,
+        GPU_IVF_PQ,
+        GPU_BRUTE_FORCE,
+        GPU_CAGRA,
+
+        // Only supported for binary vectors
+        BIN_FLAT,
+        BIN_IVF_FLAT,
+
+        // Only for varchar type field
+        TRIE,
+        // Only for scalar type field
+        STL_SORT, // only for numeric type field
+        INVERTED, // works for all scalar fields except JSON type field
+
+        // Only for sparse vectors
+        SPARSE_INVERTED_INDEX,
+        SPARSE_WAND,
+        ;
+
+        public static IndexType of(String name) {
+            return valueOf(name.toUpperCase());
+        }
+    }
+
+    public enum MetricType {
+        // Only for float vectors
+        L2,
+        IP,
+        COSINE,
+
+        // Only for binary vectors
+        HAMMING,
+        JACCARD,
+        ;
+
+        public static MetricType of(String name) {
+            return valueOf(name.toUpperCase());
+        }
+    }
 }
diff --git a/...src/main/java/org/apache/seatunnel/connectors/seatunnel/milvus/catalog/MilvusCatalog.java b/...src/main/java/org/apache/seatunnel/connectors/seatunnel/milvus/catalog/MilvusCatalog.java
@@ -217,8 +217,8 @@ private void createIndexInternal(
                             .withCollectionName(tablePath.getTableName())
                             .withFieldName(index.getColumnName())
                             .withIndexName(index.getIndexName())
-                            .withIndexType(IndexType.valueOf(index.getIndexType()))
-                            .withMetricType(MetricType.valueOf(index.getMetricType()))
+                            .withIndexType(IndexType.valueOf(index.getIndexType().name()))
+                            .withMetricType(MetricType.valueOf(index.getMetricType().name()))
                             .build();
 
             R<RpcStatus> response = client.createIndex(createIndexParam);