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);