Skip to content

Commit 4710870

Browse files
JingsongLiJingsongLi
committed
[vector] Align JNI wrapper with vector index API
1 parent e3ccf56 commit 4710870

21 files changed

Lines changed: 97 additions & 171 deletions

‎paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/IvfFlatVectorGlobalIndexerFactory.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,6 @@
1818

1919
package org.apache.paimon.vector.index;
2020

21-
import org.apache.paimon.index.ivfpq.IndexType;
22-
2321
/** Factory for the {@code ivf-flat} vector index identifier. */
2422
public class IvfFlatVectorGlobalIndexerFactory extends VectorGlobalIndexerFactory {
2523

@@ -31,7 +29,7 @@ public String identifier() {
3129
}
3230

3331
@Override
34-
protected IndexType indexType() {
35-
return IndexType.IVF_FLAT;
32+
protected VectorIndexType indexType() {
33+
return VectorIndexType.IVF_FLAT;
3634
}
3735
}

‎paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/IvfHnswFlatVectorGlobalIndexerFactory.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,6 @@
1818

1919
package org.apache.paimon.vector.index;
2020

21-
import org.apache.paimon.index.ivfpq.IndexType;
22-
2321
/** Factory for the {@code ivf-hnsw-flat} vector index identifier. */
2422
public class IvfHnswFlatVectorGlobalIndexerFactory extends VectorGlobalIndexerFactory {
2523

@@ -31,7 +29,7 @@ public String identifier() {
3129
}
3230

3331
@Override
34-
protected IndexType indexType() {
35-
return IndexType.IVF_HNSW_FLAT;
32+
protected VectorIndexType indexType() {
33+
return VectorIndexType.IVF_HNSW_FLAT;
3634
}
3735
}

‎paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/IvfHnswSqVectorGlobalIndexerFactory.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,6 @@
1818

1919
package org.apache.paimon.vector.index;
2020

21-
import org.apache.paimon.index.ivfpq.IndexType;
22-
2321
/** Factory for the {@code ivf-hnsw-sq} vector index identifier. */
2422
public class IvfHnswSqVectorGlobalIndexerFactory extends VectorGlobalIndexerFactory {
2523

@@ -31,7 +29,7 @@ public String identifier() {
3129
}
3230

3331
@Override
34-
protected IndexType indexType() {
35-
return IndexType.IVF_HNSW_SQ;
32+
protected VectorIndexType indexType() {
33+
return VectorIndexType.IVF_HNSW_SQ;
3634
}
3735
}

‎paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/IvfPqAlgorithmVectorGlobalIndexerFactory.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,6 @@
1818

1919
package org.apache.paimon.vector.index;
2020

21-
import org.apache.paimon.index.ivfpq.IndexType;
22-
2321
/** Factory for the {@code ivf-pq} vector index identifier. */
2422
public class IvfPqAlgorithmVectorGlobalIndexerFactory extends VectorGlobalIndexerFactory {
2523

@@ -31,7 +29,7 @@ public String identifier() {
3129
}
3230

3331
@Override
34-
protected IndexType indexType() {
35-
return IndexType.IVF_PQ;
32+
protected VectorIndexType indexType() {
33+
return VectorIndexType.IVF_PQ;
3634
}
3735
}

‎paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/VectorGlobalIndexReader.java‎

Lines changed: 13 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,10 @@
2424
import org.apache.paimon.globalindex.GlobalIndexResult;
2525
import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
2626
import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
27-
import org.apache.paimon.index.ivfpq.Metric;
28-
import org.apache.paimon.index.ivfpq.VectorIndexInput;
29-
import org.apache.paimon.index.ivfpq.VectorIndexMetadata;
30-
import org.apache.paimon.index.ivfpq.VectorIndexReader;
31-
import org.apache.paimon.index.ivfpq.VectorSearchResult;
27+
import org.apache.paimon.index.vector.VectorIndexInput;
28+
import org.apache.paimon.index.vector.VectorIndexMetadata;
29+
import org.apache.paimon.index.vector.VectorIndexReader;
30+
import org.apache.paimon.index.vector.VectorSearchResult;
3231
import org.apache.paimon.predicate.FieldRef;
3332
import org.apache.paimon.predicate.VectorSearch;
3433
import org.apache.paimon.types.ArrayType;
@@ -101,7 +100,7 @@ private ScoredGlobalIndexResult search(VectorSearch vectorSearch) throws IOExcep
101100
float[] queryVector = vectorSearch.vector().clone();
102101
int limit = vectorSearch.limit();
103102
int nprobe = indexMeta.nprobe();
104-
Metric metric = nativeMeta.metric();
103+
String metric = nativeMeta.metric();
105104

106105
RoaringNavigableMap64 includeRowIds = vectorSearch.includeRowIds();
107106
VectorSearchResult result;
@@ -158,17 +157,15 @@ private ScoredGlobalIndexResult search(VectorSearch vectorSearch) throws IOExcep
158157
});
159158
}
160159

161-
private static float convertDistanceToScore(float distance, Metric metric) {
162-
switch (metric) {
163-
case L2:
164-
return 1.0f / (1.0f + distance);
165-
case COSINE:
166-
return 1.0f - distance;
167-
case INNER_PRODUCT:
168-
return distance;
169-
default:
170-
throw new IllegalArgumentException("Unknown metric: " + metric);
160+
private static float convertDistanceToScore(float distance, String metric) {
161+
if ("l2".equals(metric)) {
162+
return 1.0f / (1.0f + distance);
163+
} else if ("cosine".equals(metric)) {
164+
return 1.0f - distance;
165+
} else if ("inner_product".equals(metric)) {
166+
return distance;
171167
}
168+
throw new IllegalArgumentException("Unknown metric: " + metric);
172169
}
173170

174171
private void validateSearchVector(Object vector) {

‎paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/VectorGlobalIndexWriter.java‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,7 @@
2424
import org.apache.paimon.globalindex.GlobalIndexSingletonWriter;
2525
import org.apache.paimon.globalindex.ResultEntry;
2626
import org.apache.paimon.globalindex.io.GlobalIndexFileWriter;
27-
import org.apache.paimon.index.ivfpq.IndexType;
28-
import org.apache.paimon.index.ivfpq.VectorIndexWriter;
27+
import org.apache.paimon.index.vector.VectorIndexWriter;
2928
import org.apache.paimon.options.Options;
3029
import org.apache.paimon.types.ArrayType;
3130
import org.apache.paimon.types.DataType;
@@ -79,7 +78,7 @@ public class VectorGlobalIndexWriter implements GlobalIndexSingletonWriter, Clos
7978
private static final int DEFAULT_ADD_BATCH_SIZE = 10000;
8079

8180
private final GlobalIndexFileWriter fileWriter;
82-
private final IndexType indexType;
81+
private final VectorIndexType indexType;
8382
private final String identifier;
8483
private final int dim;
8584
private final String metric;
@@ -109,7 +108,7 @@ public VectorGlobalIndexWriter(
109108
GlobalIndexFileWriter fileWriter,
110109
DataType fieldType,
111110
Options options,
112-
IndexType indexType,
111+
VectorIndexType indexType,
113112
String identifier) {
114113
this.fileWriter = fileWriter;
115114
this.indexType = indexType;
@@ -414,7 +413,7 @@ private void addVectorsFromTempFile(VectorIndexWriter writer) throws IOException
414413

415414
private Map<String, String> nativeOptions(int effectiveNlist) {
416415
Map<String, String> nativeOptions = new LinkedHashMap<>();
417-
nativeOptions.put("index.type", nativeIndexType(indexType));
416+
nativeOptions.put("index.type", indexType.nativeName());
418417
nativeOptions.put("dimension", String.valueOf(dim));
419418
nativeOptions.put("nlist", String.valueOf(effectiveNlist));
420419
nativeOptions.put("metric", metric);
@@ -445,7 +444,7 @@ private Map<String, String> metadata() {
445444
}
446445

447446
private void validateOptions() {
448-
if (indexType == IndexType.IVF_PQ && dim % pqM != 0) {
447+
if (indexType == VectorIndexType.IVF_PQ && dim % pqM != 0) {
449448
throw new IllegalArgumentException(
450449
String.format("vector.pq.m (%d) must divide vector dimension (%d)", pqM, dim));
451450
}
@@ -503,21 +502,6 @@ private static double doubleOption(Options options, String key, double defaultVa
503502
return options.getDouble(OPTION_PREFIX + key, defaultValue);
504503
}
505504

506-
private static String nativeIndexType(IndexType indexType) {
507-
switch (indexType) {
508-
case IVF_FLAT:
509-
return "ivf_flat";
510-
case IVF_PQ:
511-
return "ivf_pq";
512-
case IVF_HNSW_FLAT:
513-
return "ivf_hnsw_flat";
514-
case IVF_HNSW_SQ:
515-
return "ivf_hnsw_sq";
516-
default:
517-
throw new IllegalArgumentException("Unsupported vector index type: " + indexType);
518-
}
519-
}
520-
521505
private static void ensureAvailable(ByteBuffer readBuf, FileChannel channel, int minBytes)
522506
throws IOException {
523507
int zeroReadCount = 0;

‎paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/VectorGlobalIndexer.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import org.apache.paimon.globalindex.GlobalIndexer;
2525
import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
2626
import org.apache.paimon.globalindex.io.GlobalIndexFileWriter;
27-
import org.apache.paimon.index.ivfpq.IndexType;
2827
import org.apache.paimon.options.Options;
2928
import org.apache.paimon.types.DataType;
3029

@@ -37,11 +36,11 @@ public class VectorGlobalIndexer implements GlobalIndexer {
3736

3837
private final DataType fieldType;
3938
private final Options options;
40-
private final IndexType indexType;
39+
private final VectorIndexType indexType;
4140
private final String identifier;
4241

4342
public VectorGlobalIndexer(
44-
DataType fieldType, Options options, IndexType indexType, String identifier) {
43+
DataType fieldType, Options options, VectorIndexType indexType, String identifier) {
4544
this.fieldType = fieldType;
4645
this.options = options;
4746
this.indexType = Objects.requireNonNull(indexType, "indexType must not be null");

‎paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/VectorGlobalIndexerFactory.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,13 @@
2020

2121
import org.apache.paimon.globalindex.GlobalIndexer;
2222
import org.apache.paimon.globalindex.GlobalIndexerFactory;
23-
import org.apache.paimon.index.ivfpq.IndexType;
2423
import org.apache.paimon.options.Options;
2524
import org.apache.paimon.types.DataField;
2625

2726
/** Factory for creating vector indexes backed by paimon-vector-index. */
2827
public abstract class VectorGlobalIndexerFactory implements GlobalIndexerFactory {
2928

30-
protected abstract IndexType indexType();
29+
protected abstract VectorIndexType indexType();
3130

3231
@Override
3332
public GlobalIndexer create(DataField field, Options options) {

‎paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/VectorIndexType.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.paimon.vector.index;
20+
21+
/** Vector index type used by the native paimon-vector-index writer. */
22+
public enum VectorIndexType {
23+
IVF_FLAT("ivf_flat"),
24+
IVF_PQ("ivf_pq"),
25+
IVF_HNSW_FLAT("ivf_hnsw_flat"),
26+
IVF_HNSW_SQ("ivf_hnsw_sq");
27+
28+
private final String nativeName;
29+
30+
VectorIndexType(String nativeName) {
31+
this.nativeName = nativeName;
32+
}
33+
34+
public String nativeName() {
35+
return nativeName;
36+
}
37+
}

‎paimon-vector/paimon-vector-index/src/test/java/org/apache/paimon/vector/index/VectorGlobalIndexTest.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,7 @@
2727
import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
2828
import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
2929
import org.apache.paimon.globalindex.io.GlobalIndexFileWriter;
30-
import org.apache.paimon.index.ivfpq.IndexType;
31-
import org.apache.paimon.index.ivfpq.NativeLoader;
30+
import org.apache.paimon.index.vector.NativeLoader;
3231
import org.apache.paimon.options.Options;
3332
import org.apache.paimon.predicate.VectorSearch;
3433
import org.apache.paimon.types.ArrayType;
@@ -334,7 +333,8 @@ public void testViaIndexer() throws IOException {
334333
};
335334

336335
VectorGlobalIndexer indexer =
337-
new VectorGlobalIndexer(vectorType, options, IndexType.IVF_PQ, IVF_PQ_IDENTIFIER);
336+
new VectorGlobalIndexer(
337+
vectorType, options, VectorIndexType.IVF_PQ, IVF_PQ_IDENTIFIER);
338338

339339
GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
340340
VectorGlobalIndexWriter writer = (VectorGlobalIndexWriter) indexer.createWriter(fileWriter);
@@ -357,7 +357,7 @@ public void testViaIndexer() throws IOException {
357357
private VectorGlobalIndexWriter createIvfPqWriter(
358358
GlobalIndexFileWriter fileWriter, DataType fieldType, Options options) {
359359
return new VectorGlobalIndexWriter(
360-
fileWriter, fieldType, options, IndexType.IVF_PQ, IVF_PQ_IDENTIFIER);
360+
fileWriter, fieldType, options, VectorIndexType.IVF_PQ, IVF_PQ_IDENTIFIER);
361361
}
362362

363363
private Options createDefaultOptions(int dimension) {

0 commit comments

Comments
 (0)