From caa5b2b5f797d0419cf228906a4d13956917a1a8 Mon Sep 17 00:00:00 2001
From: Thomas-HuWei <wei.hu@zilliz.com>
Date: Fri, 12 Jul 2024 16:03:04 +0800
Subject: [PATCH] make indexType and metricType enum

---
 .../api/table/catalog/VectorIndex.java        | 64 ++++++++++++++++++-
 .../milvus/catalog/MilvusCatalog.java         |  4 +-
 2 files changed, 64 insertions(+), 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
index 368704231c5..5d6dd1beaae 100644
--- 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/seatunnel-connectors-v2/connector-milvus/src/main/java/org/apache/seatunnel/connectors/seatunnel/milvus/catalog/MilvusCatalog.java b/seatunnel-connectors-v2/connector-milvus/src/main/java/org/apache/seatunnel/connectors/seatunnel/milvus/catalog/MilvusCatalog.java
index 21253f71ff5..dcca41320c0 100644
--- a/seatunnel-connectors-v2/connector-milvus/src/main/java/org/apache/seatunnel/connectors/seatunnel/milvus/catalog/MilvusCatalog.java
+++ b/seatunnel-connectors-v2/connector-milvus/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);