From f1b82e9369f6a9aeed025886f749ec0fcef88caf Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 28 Nov 2025 17:27:51 +0800 Subject: [PATCH 01/11] init --- paimon-vector/pom.xml | 69 +++++ .../vector/VectorGlobalIndexReader.java | 293 ++++++++++++++++++ .../vector/VectorGlobalIndexWriter.java | 202 ++++++++++++ .../paimon/vector/VectorGlobalIndexer.java | 54 ++++ .../vector/VectorGlobalIndexerFactory.java | 40 +++ ...he.paimon.globalindex.GlobalIndexerFactory | 16 + .../paimon/vector/VectorGlobalIndexTest.java | 57 ++++ pom.xml | 1 + 8 files changed, 732 insertions(+) create mode 100644 paimon-vector/pom.xml create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexer.java create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexerFactory.java create mode 100644 paimon-vector/src/main/resources/META-INF/services/org.apache.paimon.globalindex.GlobalIndexerFactory create mode 100644 paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java diff --git a/paimon-vector/pom.xml b/paimon-vector/pom.xml new file mode 100644 index 000000000000..499aa0d04c3e --- /dev/null +++ b/paimon-vector/pom.xml @@ -0,0 +1,69 @@ + + + + 4.0.0 + + + paimon-parent + org.apache.paimon + 1.4-SNAPSHOT + + + paimon-vector + Paimon : Vector + + + 3.0.6 + 11 + + + + + org.apache.paimon + paimon-common + ${project.version} + + + + io.github.jbellis + jvector + ${jvector.version} + + + + + + org.apache.paimon + paimon-common + ${project.version} + test-jar + test + + + + org.apache.paimon + paimon-test-utils + ${project.version} + test + + + diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java new file mode 100644 index 000000000000..7d252f073901 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -0,0 +1,293 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.paimon.fs.SeekableInputStream; +import org.apache.paimon.globalindex.GlobalIndexIOMeta; +import org.apache.paimon.globalindex.GlobalIndexReader; +import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.globalindex.io.GlobalIndexFileReader; +import org.apache.paimon.options.Options; +import org.apache.paimon.predicate.FieldRef; +import org.apache.paimon.types.DataType; + +import java.io.ByteArrayInputStream; +import java.io.DataInputStream; +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Optional; +import java.util.Set; + +/** + * Vector global index reader. + * + *

This is a framework implementation that provides the structure for vector search. In + * production, integrate with JVector or other vector search engines. + */ +public class VectorGlobalIndexReader implements GlobalIndexReader { + + private final List shards; + + public VectorGlobalIndexReader( + GlobalIndexFileReader fileReader, + List files, + DataType fieldType, + Options options) + throws IOException { + this.shards = new ArrayList<>(); + loadShards(fileReader, files); + } + + private void loadShards(GlobalIndexFileReader fileReader, List files) + throws IOException { + for (GlobalIndexIOMeta meta : files) { + try (SeekableInputStream in = fileReader.getInputStream(meta.fileName())) { + byte[] indexBytes = new byte[(int) meta.fileSize()]; + int totalRead = 0; + while (totalRead < indexBytes.length) { + int read = in.read(indexBytes, totalRead, indexBytes.length - totalRead); + if (read == -1) { + throw new IOException("Unexpected end of stream"); + } + totalRead += read; + } + + IndexShard shard = deserializeIndex(indexBytes, meta); + shards.add(shard); + } + } + } + + private IndexShard deserializeIndex(byte[] indexBytes, GlobalIndexIOMeta meta) + throws IOException { + ByteArrayInputStream byteIn = new ByteArrayInputStream(indexBytes); + try (DataInputStream dataIn = new DataInputStream(byteIn)) { + // Read version + int version = dataIn.readInt(); + if (version != 1) { + throw new IOException("Unsupported index version: " + version); + } + + // Read row IDs + int size = dataIn.readInt(); + long[] rowIds = new long[size]; + for (int i = 0; i < size; i++) { + rowIds[i] = dataIn.readLong(); + } + + // Read vectors + List vectors = new ArrayList<>(); + for (int i = 0; i < size; i++) { + int dim = dataIn.readInt(); + float[] vector = new float[dim]; + for (int j = 0; j < dim; j++) { + vector[j] = dataIn.readFloat(); + } + vectors.add(vector); + } + + // Parse metadata + byte[] metaBytes = meta.metadata(); + VectorIndexMetadata metadata = deserializeMetadata(metaBytes); + + return new IndexShard(rowIds, vectors, metadata, meta.rowIdRange().from); + } + } + + private VectorIndexMetadata deserializeMetadata(byte[] metaBytes) throws IOException { + if (metaBytes == null || metaBytes.length == 0) { + return new VectorIndexMetadata(128, "COSINE", 16, 100); + } + + ByteArrayInputStream byteIn = new ByteArrayInputStream(metaBytes); + try (DataInputStream dataIn = new DataInputStream(byteIn)) { + int dimension = dataIn.readInt(); + String metricName = dataIn.readUTF(); + int m = dataIn.readInt(); + int efConstruction = dataIn.readInt(); + return new VectorIndexMetadata(dimension, metricName, m, efConstruction); + } + } + + /** + * Search for similar vectors. + * + * @param query query vector + * @param k number of results + * @return global index result containing row IDs + */ + public GlobalIndexResult search(float[] query, int k) { + Set resultIds = new HashSet<>(); + + // TODO: Implement vector similarity search with JVector or other engine + // For now, return empty results as placeholder + for (IndexShard shard : shards) { + // Brute force search as placeholder + List scored = new ArrayList<>(); + for (int i = 0; i < shard.vectors.size() && i < shard.rowIds.length; i++) { + float[] vec = shard.vectors.get(i); + float score = cosineSimilarity(query, vec); + scored.add(new ScoredResult(shard.rowIds[i], score)); + } + + // Sort and take top k + scored.sort((a, b) -> Float.compare(b.score, a.score)); + for (int i = 0; i < Math.min(k, scored.size()); i++) { + resultIds.add(scored.get(i).rowId); + } + } + + return GlobalIndexResult.wrap(resultIds); + } + + private float cosineSimilarity(float[] a, float[] b) { + if (a.length != b.length) { + return 0.0f; + } + float dotProduct = 0.0f; + float normA = 0.0f; + float normB = 0.0f; + for (int i = 0; i < a.length; i++) { + dotProduct += a[i] * b[i]; + normA += a[i] * a[i]; + normB += b[i] * b[i]; + } + return dotProduct / (float) (Math.sqrt(normA) * Math.sqrt(normB)); + } + + private static class ScoredResult { + final long rowId; + final float score; + + ScoredResult(long rowId, float score) { + this.rowId = rowId; + this.score = score; + } + } + + @Override + public void close() throws IOException { + shards.clear(); + } + + // Implementation of FunctionVisitor methods + @Override + public Optional visitIsNotNull(FieldRef fieldRef) { + throw new UnsupportedOperationException( + "Vector index does not support isNotNull predicate"); + } + + @Override + public Optional visitIsNull(FieldRef fieldRef) { + throw new UnsupportedOperationException("Vector index does not support isNull predicate"); + } + + @Override + public Optional visitStartsWith(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException( + "Vector index does not support startsWith predicate"); + } + + @Override + public Optional visitEndsWith(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support endsWith predicate"); + } + + @Override + public Optional visitContains(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support contains predicate"); + } + + @Override + public Optional visitLessThan(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support lessThan predicate"); + } + + @Override + public Optional visitGreaterOrEqual(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException( + "Vector index does not support greaterOrEqual predicate"); + } + + @Override + public Optional visitNotEqual(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support notEqual predicate"); + } + + @Override + public Optional visitLessOrEqual(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException( + "Vector index does not support lessOrEqual predicate"); + } + + @Override + public Optional visitEqual(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support equal predicate"); + } + + @Override + public Optional visitGreaterThan(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException( + "Vector index does not support greaterThan predicate"); + } + + @Override + public Optional visitIn(FieldRef fieldRef, List literals) { + throw new UnsupportedOperationException("Vector index does not support in predicate"); + } + + @Override + public Optional visitNotIn(FieldRef fieldRef, List literals) { + throw new UnsupportedOperationException("Vector index does not support notIn predicate"); + } + + private static class IndexShard { + final long[] rowIds; + final List vectors; + final VectorIndexMetadata metadata; + final long rowRangeStart; + + IndexShard( + long[] rowIds, + List vectors, + VectorIndexMetadata metadata, + long rowRangeStart) { + this.rowIds = rowIds; + this.vectors = vectors; + this.metadata = metadata; + this.rowRangeStart = rowRangeStart; + } + } + + private static class VectorIndexMetadata { + final int dimension; + final String similarityFunction; + final int m; + final int efConstruction; + + VectorIndexMetadata(int dimension, String similarityFunction, int m, int efConstruction) { + this.dimension = dimension; + this.similarityFunction = similarityFunction; + this.m = m; + this.efConstruction = efConstruction; + } + } +} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java new file mode 100644 index 000000000000..772afc8bb291 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -0,0 +1,202 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.paimon.globalindex.GlobalIndexWriter; +import org.apache.paimon.globalindex.io.GlobalIndexFileWriter; +import org.apache.paimon.options.Options; +import org.apache.paimon.types.ArrayType; +import org.apache.paimon.types.DataType; +import org.apache.paimon.utils.Preconditions; +import org.apache.paimon.utils.Range; + +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.List; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** + * Vector global index writer. + * + *

This is a framework implementation that provides the structure for vector indexing. In + * production, integrate with JVector or other vector search engines. + */ +public class VectorGlobalIndexWriter implements GlobalIndexWriter { + + private static final String OPTION_VECTOR_DIM = "vector.dim"; + private static final String OPTION_VECTOR_METRIC = "vector.metric"; + private static final String OPTION_VECTOR_M = "vector.M"; + private static final String OPTION_VECTOR_EF_CONSTRUCTION = "vector.ef-construction"; + + private final GlobalIndexFileWriter fileWriter; + private final DataType fieldType; + private final int dimension; + private final String similarityFunction; + private final int m; + private final int efConstruction; + + private final List vectors; + private long minRowId = Long.MAX_VALUE; + private long maxRowId = Long.MIN_VALUE; + + public VectorGlobalIndexWriter( + GlobalIndexFileWriter fileWriter, DataType fieldType, Options options) { + this.fileWriter = fileWriter; + this.fieldType = fieldType; + this.vectors = new ArrayList<>(); + + // Parse options + this.dimension = options.getInteger(OPTION_VECTOR_DIM, 128); + String metric = options.getString(OPTION_VECTOR_METRIC, "COSINE"); + this.similarityFunction = parseMetric(metric); + this.m = options.getInteger(OPTION_VECTOR_M, 16); + this.efConstruction = options.getInteger(OPTION_VECTOR_EF_CONSTRUCTION, 100); + + // Validate field type + validateFieldType(fieldType); + } + + private void validateFieldType(DataType type) { + checkArgument( + type instanceof ArrayType, "Vector field type must be ARRAY, but was: " + type); + } + + private String parseMetric(String metric) { + switch (metric.toUpperCase()) { + case "COSINE": + case "DOT_PRODUCT": + case "EUCLIDEAN": + return metric.toUpperCase(); + default: + throw new IllegalArgumentException("Unsupported metric: " + metric); + } + } + + @Override + public void write(Object key) { + Preconditions.checkArgument( + key instanceof VectorKey, "Key must be VectorKey, but was: " + key.getClass()); + + VectorKey vectorKey = (VectorKey) key; + long rowId = vectorKey.getRowId(); + float[] vector = vectorKey.getVector(); + + checkArgument( + vector.length == dimension, + "Vector dimension mismatch: expected " + dimension + ", but got " + vector.length); + + vectors.add(new VectorWithRowId(rowId, vector)); + minRowId = Math.min(minRowId, rowId); + maxRowId = Math.max(maxRowId, rowId); + } + + @Override + public List finish() { + try { + if (vectors.isEmpty()) { + return new ArrayList<>(); + } + + // Build vector index + // TODO: Integrate with JVector or other vector search engine + // For now, serialize the raw vectors with metadata + byte[] indexBytes = serializeVectors(); + byte[] metaBytes = serializeMetadata(); + + // Write to file + String fileName = fileWriter.newFileName(VectorGlobalIndexerFactory.IDENTIFIER); + try (OutputStream out = fileWriter.newOutputStream(fileName)) { + out.write(indexBytes); + } + + List results = new ArrayList<>(); + results.add(ResultEntry.of(fileName, metaBytes, new Range(minRowId, maxRowId + 1))); + return results; + } catch (IOException e) { + throw new RuntimeException("Failed to write vector global index", e); + } + } + + private byte[] serializeVectors() throws IOException { + ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); + try (DataOutputStream dataOut = new DataOutputStream(byteOut)) { + // Write version + dataOut.writeInt(1); + + // Write row IDs + dataOut.writeInt(vectors.size()); + for (VectorWithRowId vector : vectors) { + dataOut.writeLong(vector.rowId); + } + + // Write vectors + for (VectorWithRowId vector : vectors) { + dataOut.writeInt(vector.vector.length); + for (float v : vector.vector) { + dataOut.writeFloat(v); + } + } + } + return byteOut.toByteArray(); + } + + private byte[] serializeMetadata() throws IOException { + ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); + try (DataOutputStream dataOut = new DataOutputStream(byteOut)) { + dataOut.writeInt(dimension); + dataOut.writeUTF(similarityFunction); + dataOut.writeInt(m); + dataOut.writeInt(efConstruction); + } + return byteOut.toByteArray(); + } + + /** Key type for vector indexing. */ + public static class VectorKey { + private final long rowId; + private final float[] vector; + + public VectorKey(long rowId, float[] vector) { + this.rowId = rowId; + this.vector = vector; + } + + public long getRowId() { + return rowId; + } + + public float[] getVector() { + return vector; + } + } + + private static class VectorWithRowId { + final long rowId; + final float[] vector; + + VectorWithRowId(long rowId, float[] vector) { + this.rowId = rowId; + this.vector = vector; + } + } +} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexer.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexer.java new file mode 100644 index 000000000000..f6fc88050d14 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexer.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.paimon.globalindex.GlobalIndexIOMeta; +import org.apache.paimon.globalindex.GlobalIndexReader; +import org.apache.paimon.globalindex.GlobalIndexWriter; +import org.apache.paimon.globalindex.GlobalIndexer; +import org.apache.paimon.globalindex.io.GlobalIndexFileReader; +import org.apache.paimon.globalindex.io.GlobalIndexFileWriter; +import org.apache.paimon.options.Options; +import org.apache.paimon.types.DataType; + +import java.io.IOException; +import java.util.List; + +/** Vector global indexer using JVector. */ +public class VectorGlobalIndexer implements GlobalIndexer { + + private final DataType fieldType; + private final Options options; + + public VectorGlobalIndexer(DataType fieldType, Options options) { + this.fieldType = fieldType; + this.options = options; + } + + @Override + public GlobalIndexWriter createWriter(GlobalIndexFileWriter fileWriter) throws IOException { + return new VectorGlobalIndexWriter(fileWriter, fieldType, options); + } + + @Override + public GlobalIndexReader createReader( + GlobalIndexFileReader fileReader, List files) throws IOException { + return new VectorGlobalIndexReader(fileReader, files, fieldType, options); + } +} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexerFactory.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexerFactory.java new file mode 100644 index 000000000000..bd2215ae1a6a --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexerFactory.java @@ -0,0 +1,40 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.paimon.globalindex.GlobalIndexer; +import org.apache.paimon.globalindex.GlobalIndexerFactory; +import org.apache.paimon.options.Options; +import org.apache.paimon.types.DataType; + +/** Factory for creating vector global indexers using JVector. */ +public class VectorGlobalIndexerFactory implements GlobalIndexerFactory { + + public static final String IDENTIFIER = "vector"; + + @Override + public String identifier() { + return IDENTIFIER; + } + + @Override + public GlobalIndexer create(DataType type, Options options) { + return new VectorGlobalIndexer(type, options); + } +} diff --git a/paimon-vector/src/main/resources/META-INF/services/org.apache.paimon.globalindex.GlobalIndexerFactory b/paimon-vector/src/main/resources/META-INF/services/org.apache.paimon.globalindex.GlobalIndexerFactory new file mode 100644 index 000000000000..393c2249b908 --- /dev/null +++ b/paimon-vector/src/main/resources/META-INF/services/org.apache.paimon.globalindex.GlobalIndexerFactory @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +org.apache.paimon.vector.VectorGlobalIndexerFactory diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java new file mode 100644 index 000000000000..248d24b08ae0 --- /dev/null +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.paimon.globalindex.GlobalIndexer; +import org.apache.paimon.globalindex.GlobalIndexerFactory; +import org.apache.paimon.globalindex.GlobalIndexerFactoryUtils; +import org.apache.paimon.options.Options; +import org.apache.paimon.types.ArrayType; +import org.apache.paimon.types.DataType; +import org.apache.paimon.types.FloatType; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Test for {@link VectorGlobalIndexerFactory}. */ +public class VectorGlobalIndexTest { + + @Test + public void testFactoryRegistration() { + GlobalIndexerFactory factory = + GlobalIndexerFactoryUtils.load(VectorGlobalIndexerFactory.IDENTIFIER); + assertThat(factory).isNotNull(); + assertThat(factory).isInstanceOf(VectorGlobalIndexerFactory.class); + assertThat(factory.identifier()).isEqualTo("vector"); + } + + @Test + public void testCreateIndexer() { + GlobalIndexerFactory factory = new VectorGlobalIndexerFactory(); + DataType vectorType = new ArrayType(new FloatType()); + Options options = new Options(); + options.setInteger("vector.dim", 128); + options.setString("vector.metric", "COSINE"); + + GlobalIndexer indexer = factory.create(vectorType, options); + assertThat(indexer).isNotNull(); + assertThat(indexer).isInstanceOf(VectorGlobalIndexer.class); + } +} diff --git a/pom.xml b/pom.xml index 2ef51944f914..0dc9608c45f3 100644 --- a/pom.xml +++ b/pom.xml @@ -63,6 +63,7 @@ under the License. paimon-filesystems paimon-format paimon-lance + paimon-vector paimon-vfs paimon-bundle paimon-hive From 04f1c6e87f98fe1c63e6d4c4c23dfa9a315db393 Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 1 Dec 2025 11:22:52 +0800 Subject: [PATCH 02/11] add demo --- .../vector/VectorGlobalIndexReader.java | 201 ++++-- .../vector/VectorGlobalIndexWriter.java | 145 +++- .../paimon/vector/VectorIndexOptions.java | 82 +++ .../VectorGlobalIndexWriterReaderTest.java | 657 ++++++++++++++++++ 4 files changed, 1004 insertions(+), 81 deletions(-) create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java create mode 100644 paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexWriterReaderTest.java diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java index 7d252f073901..ba8c70591974 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -27,6 +27,16 @@ import org.apache.paimon.predicate.FieldRef; import org.apache.paimon.types.DataType; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.GraphSearcher; +import io.github.jbellis.jvector.graph.OnHeapGraphIndex; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.SearchResult; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; + import java.io.ByteArrayInputStream; import java.io.DataInputStream; import java.io.IOException; @@ -37,14 +47,18 @@ import java.util.Set; /** - * Vector global index reader. + * Vector global index reader using JVector. * - *

This is a framework implementation that provides the structure for vector search. In - * production, integrate with JVector or other vector search engines. + *

This implementation uses JVector's HNSW graph for efficient approximate nearest neighbor + * search. */ public class VectorGlobalIndexReader implements GlobalIndexReader { private final List shards; + private final String similarityFunction; + private final int dimension; + private final int m; + private final int efConstruction; public VectorGlobalIndexReader( GlobalIndexFileReader fileReader, @@ -53,6 +67,11 @@ public VectorGlobalIndexReader( Options options) throws IOException { this.shards = new ArrayList<>(); + VectorIndexOptions vectorOptions = new VectorIndexOptions(options); + this.similarityFunction = vectorOptions.metric(); + this.dimension = vectorOptions.dimension(); + this.m = vectorOptions.m(); + this.efConstruction = vectorOptions.efConstruction(); loadShards(fileReader, files); } @@ -82,9 +101,6 @@ private IndexShard deserializeIndex(byte[] indexBytes, GlobalIndexIOMeta meta) try (DataInputStream dataIn = new DataInputStream(byteIn)) { // Read version int version = dataIn.readInt(); - if (version != 1) { - throw new IOException("Unsupported index version: " + version); - } // Read row IDs int size = dataIn.readInt(); @@ -108,13 +124,71 @@ private IndexShard deserializeIndex(byte[] indexBytes, GlobalIndexIOMeta meta) byte[] metaBytes = meta.metadata(); VectorIndexMetadata metadata = deserializeMetadata(metaBytes); - return new IndexShard(rowIds, vectors, metadata, meta.rowIdRange().from); + // Rebuild graph index if version 2 (with JVector) + OnHeapGraphIndex graphIndex = null; + if (version == 2 && dataIn.available() > 0) { + // Read graph structure + int graphSize = dataIn.readInt(); + int maxDegree = dataIn.readInt(); + + // Rebuild the graph index from vectors + var vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); + List> vectorFloats = new ArrayList<>(); + for (float[] vec : vectors) { + vectorFloats.add(vectorTypeSupport.createFloatVector(vec)); + } + + // Create RandomAccessVectorValues + RandomAccessVectorValues vectorValues = + new RandomAccessVectorValues() { + @Override + public int size() { + return vectorFloats.size(); + } + + @Override + public int dimension() { + return metadata.dimension; + } + + @Override + public VectorFloat getVector(int i) { + return vectorFloats.get(i); + } + + @Override + public boolean isValueShared() { + return false; + } + + @Override + public RandomAccessVectorValues copy() { + return this; + } + }; + + // Rebuild graph using stored structure or rebuild from scratch + VectorSimilarityFunction similarityFunction = + VectorSimilarityFunction.valueOf(metadata.similarityFunction); + var builder = + new GraphIndexBuilder( + vectorValues, + similarityFunction, + metadata.m, + metadata.efConstruction, + 1.0f, // todo: need conf + 1.0f); // todo: need conf + graphIndex = builder.build(vectorValues); + } + + return new IndexShard(rowIds, vectors, metadata, meta.rowIdRange().from, graphIndex); } } private VectorIndexMetadata deserializeMetadata(byte[] metaBytes) throws IOException { if (metaBytes == null || metaBytes.length == 0) { - return new VectorIndexMetadata(128, "COSINE", 16, 100); + return new VectorIndexMetadata( + this.dimension, this.similarityFunction, this.m, this.efConstruction); } ByteArrayInputStream byteIn = new ByteArrayInputStream(metaBytes); @@ -128,7 +202,7 @@ private VectorIndexMetadata deserializeMetadata(byte[] metaBytes) throws IOExcep } /** - * Search for similar vectors. + * Search for similar vectors using JVector HNSW. * * @param query query vector * @param k number of results @@ -136,51 +210,79 @@ private VectorIndexMetadata deserializeMetadata(byte[] metaBytes) throws IOExcep */ public GlobalIndexResult search(float[] query, int k) { Set resultIds = new HashSet<>(); + var vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); + VectorFloat queryVector = vectorTypeSupport.createFloatVector(query); - // TODO: Implement vector similarity search with JVector or other engine - // For now, return empty results as placeholder for (IndexShard shard : shards) { - // Brute force search as placeholder - List scored = new ArrayList<>(); - for (int i = 0; i < shard.vectors.size() && i < shard.rowIds.length; i++) { - float[] vec = shard.vectors.get(i); - float score = cosineSimilarity(query, vec); - scored.add(new ScoredResult(shard.rowIds[i], score)); - } + try { + if (shard.graphIndex == null) { + // Fall back to brute force if no graph index + continue; + } - // Sort and take top k - scored.sort((a, b) -> Float.compare(b.score, a.score)); - for (int i = 0; i < Math.min(k, scored.size()); i++) { - resultIds.add(scored.get(i).rowId); - } - } + // Use JVector's GraphSearcher for efficient ANN search + VectorSimilarityFunction similarityFunction = + VectorSimilarityFunction.valueOf(shard.metadata.similarityFunction); - return GlobalIndexResult.wrap(resultIds); - } + // Create vector values for search context + var vectorTypeSupport2 = VectorizationProvider.getInstance().getVectorTypeSupport(); + List> vectorFloats = new ArrayList<>(); + for (float[] vec : shard.vectors) { + vectorFloats.add(vectorTypeSupport2.createFloatVector(vec)); + } - private float cosineSimilarity(float[] a, float[] b) { - if (a.length != b.length) { - return 0.0f; - } - float dotProduct = 0.0f; - float normA = 0.0f; - float normB = 0.0f; - for (int i = 0; i < a.length; i++) { - dotProduct += a[i] * b[i]; - normA += a[i] * a[i]; - normB += b[i] * b[i]; + RandomAccessVectorValues vectorValues = + new RandomAccessVectorValues() { + @Override + public int size() { + return vectorFloats.size(); + } + + @Override + public int dimension() { + return shard.metadata.dimension; + } + + @Override + public VectorFloat getVector(int i) { + return vectorFloats.get(i); + } + + @Override + public boolean isValueShared() { + return false; + } + + @Override + public RandomAccessVectorValues copy() { + return this; + } + }; + + // Search using static method + SearchResult result = + GraphSearcher.search( + queryVector, + k, + vectorValues, + similarityFunction, + shard.graphIndex, + Bits.ALL); + + // Collect row IDs from results + var nodes = result.getNodes(); + for (int i = 0; i < nodes.length && i < k; i++) { + int nodeId = nodes[i].node; + if (nodeId >= 0 && nodeId < shard.rowIds.length) { + resultIds.add(shard.rowIds[nodeId]); + } + } + } catch (Exception e) { + throw new RuntimeException("Failed to search vector index", e); + } } - return dotProduct / (float) (Math.sqrt(normA) * Math.sqrt(normB)); - } - - private static class ScoredResult { - final long rowId; - final float score; - ScoredResult(long rowId, float score) { - this.rowId = rowId; - this.score = score; - } + return GlobalIndexResult.wrap(resultIds); } @Override @@ -264,16 +366,19 @@ private static class IndexShard { final List vectors; final VectorIndexMetadata metadata; final long rowRangeStart; + final OnHeapGraphIndex graphIndex; IndexShard( long[] rowIds, List vectors, VectorIndexMetadata metadata, - long rowRangeStart) { + long rowRangeStart, + OnHeapGraphIndex graphIndex) { this.rowIds = rowIds; this.vectors = vectors; this.metadata = metadata; this.rowRangeStart = rowRangeStart; + this.graphIndex = graphIndex; } } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java index 772afc8bb291..f4b6f47a07a1 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -26,6 +26,13 @@ import org.apache.paimon.utils.Preconditions; import org.apache.paimon.utils.Range; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.OnHeapGraphIndex; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; + import java.io.ByteArrayOutputStream; import java.io.DataOutputStream; import java.io.IOException; @@ -36,24 +43,17 @@ import static org.apache.paimon.utils.Preconditions.checkArgument; /** - * Vector global index writer. + * Vector global index writer using JVector. * - *

This is a framework implementation that provides the structure for vector indexing. In - * production, integrate with JVector or other vector search engines. + *

This implementation uses JVector's HNSW (Hierarchical Navigable Small World) graph algorithm + * for efficient approximate nearest neighbor search. */ public class VectorGlobalIndexWriter implements GlobalIndexWriter { - private static final String OPTION_VECTOR_DIM = "vector.dim"; - private static final String OPTION_VECTOR_METRIC = "vector.metric"; - private static final String OPTION_VECTOR_M = "vector.M"; - private static final String OPTION_VECTOR_EF_CONSTRUCTION = "vector.ef-construction"; - private final GlobalIndexFileWriter fileWriter; private final DataType fieldType; - private final int dimension; - private final String similarityFunction; - private final int m; - private final int efConstruction; + private final VectorIndexOptions vectorOptions; + private final VectorSimilarityFunction similarityFunction; private final List vectors; private long minRowId = Long.MAX_VALUE; @@ -66,11 +66,8 @@ public VectorGlobalIndexWriter( this.vectors = new ArrayList<>(); // Parse options - this.dimension = options.getInteger(OPTION_VECTOR_DIM, 128); - String metric = options.getString(OPTION_VECTOR_METRIC, "COSINE"); - this.similarityFunction = parseMetric(metric); - this.m = options.getInteger(OPTION_VECTOR_M, 16); - this.efConstruction = options.getInteger(OPTION_VECTOR_EF_CONSTRUCTION, 100); + this.vectorOptions = new VectorIndexOptions(options); + this.similarityFunction = parseMetricToJVector(vectorOptions.metric()); // Validate field type validateFieldType(fieldType); @@ -81,12 +78,14 @@ private void validateFieldType(DataType type) { type instanceof ArrayType, "Vector field type must be ARRAY, but was: " + type); } - private String parseMetric(String metric) { + private VectorSimilarityFunction parseMetricToJVector(String metric) { switch (metric.toUpperCase()) { case "COSINE": + return VectorSimilarityFunction.COSINE; case "DOT_PRODUCT": + return VectorSimilarityFunction.DOT_PRODUCT; case "EUCLIDEAN": - return metric.toUpperCase(); + return VectorSimilarityFunction.EUCLIDEAN; default: throw new IllegalArgumentException("Unsupported metric: " + metric); } @@ -102,8 +101,11 @@ public void write(Object key) { float[] vector = vectorKey.getVector(); checkArgument( - vector.length == dimension, - "Vector dimension mismatch: expected " + dimension + ", but got " + vector.length); + vector.length == vectorOptions.dimension(), + "Vector dimension mismatch: expected " + + vectorOptions.dimension() + + ", but got " + + vector.length); vectors.add(new VectorWithRowId(rowId, vector)); minRowId = Math.min(minRowId, rowId); @@ -117,13 +119,16 @@ public List finish() { return new ArrayList<>(); } - // Build vector index - // TODO: Integrate with JVector or other vector search engine - // For now, serialize the raw vectors with metadata - byte[] indexBytes = serializeVectors(); + // Build JVector HNSW index + OnHeapGraphIndex graphIndex = buildJVectorIndex(); + + // Serialize index and metadata + byte[] indexBytes = serializeIndexWithJVector(graphIndex); byte[] metaBytes = serializeMetadata(); - // Write to file + // Write to file. todo: here could use OnDiskGraphIndex to write directly to file. + // RandomAccessVectorValues ravv + // OnDiskGraphIndex.write(index, ravv, indexPath); String fileName = fileWriter.newFileName(VectorGlobalIndexerFactory.IDENTIFIER); try (OutputStream out = fileWriter.newOutputStream(fileName)) { out.write(indexBytes); @@ -137,11 +142,64 @@ public List finish() { } } - private byte[] serializeVectors() throws IOException { + private OnHeapGraphIndex buildJVectorIndex() { + // Get the vectorization provider + var vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); + + // Create list of VectorFloat for JVector + List> vectorFloats = new ArrayList<>(); + for (VectorWithRowId vec : vectors) { + vectorFloats.add(vectorTypeSupport.createFloatVector(vec.vector)); + } + + // Create vector values adapter for JVector + RandomAccessVectorValues vectorValues = + new RandomAccessVectorValues() { + @Override + public int size() { + return vectorFloats.size(); + } + + @Override + public int dimension() { + return vectorOptions.dimension(); + } + + @Override + public VectorFloat getVector(int i) { + return vectorFloats.get(i); + } + + @Override + public boolean isValueShared() { + return false; + } + + @Override + public RandomAccessVectorValues copy() { + return this; + } + }; + + // Build HNSW graph index using JVector + GraphIndexBuilder builder = + new GraphIndexBuilder( + vectorValues, + similarityFunction, + vectorOptions.m(), + vectorOptions.efConstruction(), + 1.0f, // alpha (diversityWeight) + 1.0f // neighborOverflow + ); + + return builder.build(vectorValues); + } + + private byte[] serializeIndexWithJVector(OnHeapGraphIndex graphIndex) throws IOException { ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); try (DataOutputStream dataOut = new DataOutputStream(byteOut)) { // Write version - dataOut.writeInt(1); + dataOut.writeInt(2); // Version 2 with JVector index // Write row IDs dataOut.writeInt(vectors.size()); @@ -149,13 +207,34 @@ private byte[] serializeVectors() throws IOException { dataOut.writeLong(vector.rowId); } - // Write vectors + // Write vectors (still needed for graph index reconstruction) for (VectorWithRowId vector : vectors) { dataOut.writeInt(vector.vector.length); for (float v : vector.vector) { dataOut.writeFloat(v); } } + + // Write JVector graph structure + dataOut.writeInt(graphIndex.size()); + dataOut.writeInt(graphIndex.maxDegree()); + + // Serialize graph neighbors for each node + for (int i = 0; i < graphIndex.size(); i++) { + var view = graphIndex.getView(); + var neighbors = view.getNeighborsIterator(i); + + // Count neighbors + List neighborList = new ArrayList<>(); + while (neighbors.hasNext()) { + neighborList.add(neighbors.nextInt()); + } + + dataOut.writeInt(neighborList.size()); + for (int neighbor : neighborList) { + dataOut.writeInt(neighbor); + } + } } return byteOut.toByteArray(); } @@ -163,10 +242,10 @@ private byte[] serializeVectors() throws IOException { private byte[] serializeMetadata() throws IOException { ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); try (DataOutputStream dataOut = new DataOutputStream(byteOut)) { - dataOut.writeInt(dimension); - dataOut.writeUTF(similarityFunction); - dataOut.writeInt(m); - dataOut.writeInt(efConstruction); + dataOut.writeInt(vectorOptions.dimension()); + dataOut.writeUTF(similarityFunction.name()); + dataOut.writeInt(vectorOptions.m()); + dataOut.writeInt(vectorOptions.efConstruction()); } return byteOut.toByteArray(); } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java new file mode 100644 index 000000000000..79d3c14649f4 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.paimon.options.ConfigOption; +import org.apache.paimon.options.ConfigOptions; +import org.apache.paimon.options.Options; + +/** Options for vector index. */ +public class VectorIndexOptions { + + public static final ConfigOption VECTOR_DIM = + ConfigOptions.key("vector.dim") + .intType() + .defaultValue(128) + .withDescription("The dimension of the vector"); + + public static final ConfigOption VECTOR_METRIC = + ConfigOptions.key("vector.metric") + .stringType() + .defaultValue("EUCLIDEAN") + .withDescription( + "The similarity metric for vector search (COSINE, DOT_PRODUCT, EUCLIDEAN)"); + + public static final ConfigOption VECTOR_M = + ConfigOptions.key("vector.M") + .intType() + .defaultValue(16) + .withDescription( + "The maximum number of connections for each element during the index construction"); + + public static final ConfigOption VECTOR_EF_CONSTRUCTION = + ConfigOptions.key("vector.ef-construction") + .intType() + .defaultValue(100) + .withDescription( + "The size of the dynamic candidate list during the index construction"); + + private final int dimension; + private final String metric; + private final int m; + private final int efConstruction; + + public VectorIndexOptions(Options options) { + this.dimension = options.get(VECTOR_DIM); + this.metric = options.get(VECTOR_METRIC); + this.m = options.get(VECTOR_M); + this.efConstruction = options.get(VECTOR_EF_CONSTRUCTION); + } + + public int dimension() { + return dimension; + } + + public String metric() { + return metric; + } + + public int m() { + return m; + } + + public int efConstruction() { + return efConstruction; + } +} diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexWriterReaderTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexWriterReaderTest.java new file mode 100644 index 000000000000..f2b000b3a610 --- /dev/null +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexWriterReaderTest.java @@ -0,0 +1,657 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.globalindex.GlobalIndexIOMeta; +import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.globalindex.GlobalIndexWriter; +import org.apache.paimon.globalindex.io.GlobalIndexFileReader; +import org.apache.paimon.globalindex.io.GlobalIndexFileWriter; +import org.apache.paimon.options.Options; +import org.apache.paimon.types.ArrayType; +import org.apache.paimon.types.DataType; +import org.apache.paimon.types.FloatType; +import org.apache.paimon.utils.Range; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Test for {@link VectorGlobalIndexWriter} and {@link VectorGlobalIndexReader} with JVector. */ +public class VectorGlobalIndexWriterReaderTest { + + @TempDir java.nio.file.Path tempDir; + + private FileIO fileIO; + private Path indexPath; + private DataType vectorType; + + @BeforeEach + public void setup() { + fileIO = new LocalFileIO(); + indexPath = new Path(tempDir.toString()); + vectorType = new ArrayType(new FloatType()); + } + + @AfterEach + public void cleanup() throws IOException { + if (fileIO != null) { + fileIO.delete(indexPath, true); + } + } + + private GlobalIndexFileWriter createFileWriter(Path path) { + return new GlobalIndexFileWriter() { + @Override + public String newFileName(String prefix) { + return prefix + "-" + UUID.randomUUID(); + } + + @Override + public OutputStream newOutputStream(String fileName) throws IOException { + return fileIO.newOutputStream(new Path(path, fileName), false); + } + }; + } + + private GlobalIndexFileReader createFileReader(Path path) { + return fileName -> fileIO.newInputStream(new Path(path, fileName)); + } + + @Test + public void testWriteAndReadVectorIndex() throws IOException { + Options options = createDefaultOptions(); + int dimension = 128; + int numVectors = 100; + + // Create writer and write vectors + GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + List testVectors = generateRandomVectors(numVectors, dimension); + for (int i = 0; i < numVectors; i++) { + writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + } + + List results = writer.finish(); + assertThat(results).hasSize(1); + + GlobalIndexWriter.ResultEntry result = results.get(0); + assertThat(result.fileName()).isNotNull(); + assertThat(result.rowRange()).isEqualTo(new Range(0, numVectors)); + + // Create reader and verify + GlobalIndexFileReader fileReader = createFileReader(indexPath); + List metas = new ArrayList<>(); + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(indexPath, result.fileName())), + result.rowRange(), + result.meta())); + + VectorGlobalIndexReader reader = + new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + + // Test search - query with first vector should return itself + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); + assertThat(searchResult).isNotNull(); + + reader.close(); + } + + @Test + public void testSearchSimilarVectors() throws IOException { + Options options = createDefaultOptions(); + int dimension = 64; + options.setInteger("vector.dim", dimension); // Set correct dimension + int numVectors = 50; + + // Create vectors with known similarities + float[] baseVector = new float[dimension]; + for (int i = 0; i < dimension; i++) { + baseVector[i] = 1.0f; + } + + GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + // Write base vector and similar variants + writer.write(new VectorGlobalIndexWriter.VectorKey(0, baseVector)); + + for (int i = 1; i < numVectors; i++) { + float[] variant = baseVector.clone(); + // Add small noise + variant[i % dimension] += (i % 2 == 0 ? 0.1f : -0.1f); + writer.write(new VectorGlobalIndexWriter.VectorKey(i, variant)); + } + + List results = writer.finish(); + assertThat(results).hasSize(1); + + GlobalIndexWriter.ResultEntry result = results.get(0); + GlobalIndexFileReader fileReader = createFileReader(indexPath); + List metas = new ArrayList<>(); + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(indexPath, result.fileName())), + result.rowRange(), + result.meta())); + + VectorGlobalIndexReader reader = + new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + + // Search with base vector should return similar vectors + GlobalIndexResult searchResult = reader.search(baseVector, 10); + assertThat(searchResult).isNotNull(); + + reader.close(); + } + + @Test + public void testDifferentSimilarityFunctions() throws IOException { + int dimension = 32; + int numVectors = 20; + + String[] metrics = {"COSINE", "DOT_PRODUCT", "EUCLIDEAN"}; + + for (String metric : metrics) { + Options options = createDefaultOptions(); + options.setString("vector.metric", metric); + options.setInteger("vector.dim", dimension); // Set correct dimension + + Path metricIndexPath = new Path(indexPath, metric.toLowerCase()); + GlobalIndexFileWriter fileWriter = createFileWriter(metricIndexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + List testVectors = generateRandomVectors(numVectors, dimension); + for (int i = 0; i < numVectors; i++) { + writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + } + + List results = writer.finish(); + assertThat(results).hasSize(1); + + GlobalIndexWriter.ResultEntry result = results.get(0); + GlobalIndexFileReader fileReader = createFileReader(metricIndexPath); + List metas = new ArrayList<>(); + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(metricIndexPath, result.fileName())), + result.rowRange(), + result.meta())); + + VectorGlobalIndexReader reader = + new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + + // Verify search works with this metric + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 3); + assertThat(searchResult).isNotNull(); + + reader.close(); + } + } + + @Test + public void testDifferentDimensions() throws IOException { + int[] dimensions = {8, 32, 128, 256}; + + for (int dimension : dimensions) { + Options options = createDefaultOptions(); + options.setInteger("vector.dim", dimension); + + Path dimIndexPath = new Path(indexPath, "dim_" + dimension); + GlobalIndexFileWriter fileWriter = createFileWriter(dimIndexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + int numVectors = 10; + List testVectors = generateRandomVectors(numVectors, dimension); + for (int i = 0; i < numVectors; i++) { + writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + } + + List results = writer.finish(); + assertThat(results).hasSize(1); + + GlobalIndexWriter.ResultEntry result = results.get(0); + GlobalIndexFileReader fileReader = createFileReader(dimIndexPath); + List metas = new ArrayList<>(); + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(dimIndexPath, result.fileName())), + result.rowRange(), + result.meta())); + + VectorGlobalIndexReader reader = + new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + + // Verify search works with this dimension + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); + assertThat(searchResult).isNotNull(); + + reader.close(); + } + } + + @Test + public void testDimensionMismatch() throws IOException { + Options options = createDefaultOptions(); + options.setInteger("vector.dim", 64); + + GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + // Try to write vector with wrong dimension + float[] wrongDimVector = new float[32]; // Wrong dimension + assertThatThrownBy( + () -> + writer.write( + new VectorGlobalIndexWriter.VectorKey(0, wrongDimVector))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("dimension mismatch"); + } + + @Test + public void testEmptyIndex() throws IOException { + Options options = createDefaultOptions(); + + GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + List results = writer.finish(); + assertThat(results).isEmpty(); + } + + @Test + public void testHNSWParameters() throws IOException { + Options options = createDefaultOptions(); + options.setInteger("vector.M", 32); // Larger M for better recall + options.setInteger("vector.ef-construction", 200); // Larger ef for better quality + + int dimension = 64; + options.setInteger("vector.dim", dimension); // Set correct dimension + int numVectors = 50; + + GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + List testVectors = generateRandomVectors(numVectors, dimension); + for (int i = 0; i < numVectors; i++) { + writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + } + + List results = writer.finish(); + assertThat(results).hasSize(1); + + GlobalIndexWriter.ResultEntry result = results.get(0); + GlobalIndexFileReader fileReader = createFileReader(indexPath); + List metas = new ArrayList<>(); + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(indexPath, result.fileName())), + result.rowRange(), + result.meta())); + + VectorGlobalIndexReader reader = + new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + + // Verify search works with custom HNSW parameters + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); + assertThat(searchResult).isNotNull(); + + reader.close(); + } + + @Test + public void testLargeK() throws IOException { + Options options = createDefaultOptions(); + int dimension = 32; + options.setInteger("vector.dim", dimension); // Set correct dimension + int numVectors = 100; + + GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + List testVectors = generateRandomVectors(numVectors, dimension); + for (int i = 0; i < numVectors; i++) { + writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + } + + List results = writer.finish(); + assertThat(results).hasSize(1); + + GlobalIndexWriter.ResultEntry result = results.get(0); + GlobalIndexFileReader fileReader = createFileReader(indexPath); + List metas = new ArrayList<>(); + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(indexPath, result.fileName())), + result.rowRange(), + result.meta())); + + VectorGlobalIndexReader reader = + new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + + // Search with k larger than number of vectors + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 200); + assertThat(searchResult).isNotNull(); + assertThat(searchResult).isNotNull(); + + reader.close(); + } + + @Test + public void testEndToEndMultiShardScenario() throws IOException { + Options options = createDefaultOptions(); + int dimension = 128; + options.setInteger("vector.dim", dimension); + options.setString("vector.metric", "COSINE"); + options.setInteger("vector.M", 24); + options.setInteger("vector.ef-construction", 150); + + // Simulate a multi-shard scenario with incremental writes + int totalVectors = 300; + int shardsCount = 3; + int vectorsPerShard = totalVectors / shardsCount; + + List allMetas = new ArrayList<>(); + List allTestVectors = generateRandomVectors(totalVectors, dimension); + + // Write vectors in multiple shards (simulating incremental indexing) + for (int shardIdx = 0; shardIdx < shardsCount; shardIdx++) { + Path shardPath = new Path(indexPath, "shard_" + shardIdx); + GlobalIndexFileWriter fileWriter = createFileWriter(shardPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + int startIdx = shardIdx * vectorsPerShard; + int endIdx = Math.min(startIdx + vectorsPerShard, totalVectors); + + // Write vectors for this shard + for (int i = startIdx; i < endIdx; i++) { + writer.write(new VectorGlobalIndexWriter.VectorKey(i, allTestVectors.get(i))); + } + + List results = writer.finish(); + assertThat(results).hasSize(1); + + GlobalIndexWriter.ResultEntry result = results.get(0); + allMetas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(shardPath, result.fileName())), + result.rowRange(), + result.meta())); + } + + // Create a single reader for all shards + GlobalIndexFileReader fileReader = + new GlobalIndexFileReader() { + @Override + public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) + throws IOException { + // Try to find the file in any shard directory + for (int shardIdx = 0; shardIdx < shardsCount; shardIdx++) { + Path shardPath = new Path(indexPath, "shard_" + shardIdx); + Path filePath = new Path(shardPath, fileName); + if (fileIO.exists(filePath)) { + return fileIO.newInputStream(filePath); + } + } + throw new IOException("File not found: " + fileName); + } + }; + + VectorGlobalIndexReader reader = + new VectorGlobalIndexReader(fileReader, allMetas, vectorType, options); + + // Test 1: Search with various k values + int[] kValues = {1, 5, 10, 20, 50}; + for (int k : kValues) { + GlobalIndexResult searchResult = reader.search(allTestVectors.get(0), k); + assertThat(searchResult).isNotNull(); + assertThat(containsRowId(searchResult, 0)).isTrue(); + } + + // Test 2: Search with query vectors from different shards + for (int shardIdx = 0; shardIdx < shardsCount; shardIdx++) { + int queryIdx = shardIdx * vectorsPerShard + 5; + GlobalIndexResult searchResult = reader.search(allTestVectors.get(queryIdx), 10); + assertThat(searchResult).isNotNull(); + assertThat(containsRowId(searchResult, queryIdx)) + .as("Search result should contain the query vector's own ID") + .isTrue(); + } + + // Test 3: Create a known similar vector and verify top result + float[] originalVector = allTestVectors.get(50); + float[] slightlyModifiedVector = originalVector.clone(); + // Add minimal noise to create a very similar vector + for (int i = 0; i < dimension; i++) { + slightlyModifiedVector[i] += (i % 2 == 0 ? 0.001f : -0.001f); + } + normalize(slightlyModifiedVector); + + GlobalIndexResult similarSearchResult = reader.search(slightlyModifiedVector, 5); + assertThat(similarSearchResult).isNotNull(); + assertThat(containsRowId(similarSearchResult, 50)) + .as("Search with very similar vector should find the original") + .isTrue(); + + // Test 4: Verify search quality with a synthetic cluster + // Create a cluster center + float[] clusterCenter = new float[dimension]; + Random random = new Random(123); + for (int i = 0; i < dimension; i++) { + clusterCenter[i] = random.nextFloat() * 2 - 1; + } + normalize(clusterCenter); + + GlobalIndexResult clusterSearchResult = reader.search(clusterCenter, 20); + assertThat(clusterSearchResult).isNotNull(); + assertThat(clusterSearchResult.iterator().hasNext()).isTrue(); + + // Test 5: Edge case - search with k larger than total vectors + GlobalIndexResult largeKResult = reader.search(allTestVectors.get(10), totalVectors * 2); + assertThat(largeKResult).isNotNull(); + + // Test 6: Multiple consecutive searches (stress test) + for (int i = 0; i < 20; i++) { + int randomIdx = random.nextInt(totalVectors); + GlobalIndexResult result = reader.search(allTestVectors.get(randomIdx), 15); + assertThat(result).isNotNull(); + } + + reader.close(); + } + + @Test + public void testVectorIndexEndToEnd() throws IOException { + Options options = createDefaultOptions(); + int dimension = 2; + options.setInteger("vector.dim", dimension); + options.setString("vector.metric", "EUCLIDEAN"); + options.setInteger("vector.M", 16); + options.setInteger("vector.ef-construction", 100); + + // Create semantic vectors representing "documents" + // Apple [0.9, 0.1] and Banana [0.8, 0.2] are similar (fruits) + // Car [0.1, 0.9] is different (vehicle) + float[] appleVector = new float[] {0.9f, 0.1f}; + float[] bananaVector = new float[] {0.8f, 0.2f}; + float[] carVector = new float[] {0.1f, 0.9f}; + + GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, vectorType, options); + + // Write vectors with row IDs representing documents + long appleId = 0L; + long bananaId = 1L; + long carId = 2L; + + writer.write(new VectorGlobalIndexWriter.VectorKey(appleId, appleVector)); + writer.write(new VectorGlobalIndexWriter.VectorKey(bananaId, bananaVector)); + writer.write(new VectorGlobalIndexWriter.VectorKey(carId, carVector)); + + List results = writer.finish(); + assertThat(results).hasSize(1); + + GlobalIndexWriter.ResultEntry result = results.get(0); + GlobalIndexFileReader fileReader = createFileReader(indexPath); + List metas = new ArrayList<>(); + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(indexPath, result.fileName())), + result.rowRange(), + result.meta())); + + VectorGlobalIndexReader reader = + new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + + // Test 1: Query with vector similar to "Apple" - should find Apple and Banana + // Query: [0.85, 0.15] is between Apple and Banana + float[] queryVector = new float[] {0.85f, 0.15f}; + GlobalIndexResult searchResult = reader.search(queryVector, 2); + + assertThat(searchResult).isNotNull(); + assertThat(searchResult.iterator().hasNext()).as("Search should return results").isTrue(); + + // Collect all result IDs + List resultIds = new ArrayList<>(); + for (Long id : searchResult) { + resultIds.add(id); + } + + // Should find Apple and Banana (similar fruits), not Car + assertThat(resultIds).as("Should find Apple (most similar)").contains(appleId); + assertThat(resultIds).as("Should find Banana (second most similar)").contains(bananaId); + assertThat(resultIds.size()).as("Should return top 2 results").isLessThanOrEqualTo(2); + + // Test 2: Query with exact match - should find exact document + GlobalIndexResult exactMatchResult = reader.search(appleVector, 1); + assertThat(containsRowId(exactMatchResult, appleId)) + .as("Exact match query should find Apple") + .isTrue(); + + // Test 3: Query with vector similar to "Car" - should find Car first + float[] carQueryVector = new float[] {0.15f, 0.85f}; + GlobalIndexResult carSearchResult = reader.search(carQueryVector, 1); + + List carResultIds = new ArrayList<>(); + for (Long id : carSearchResult) { + carResultIds.add(id); + } + + assertThat(carResultIds).as("Query similar to Car should find Car").contains(carId); + + // Test 4: Search with larger k than available documents + GlobalIndexResult largeKResult = reader.search(queryVector, 10); + assertThat(largeKResult).isNotNull(); + assertThat(largeKResult.iterator().hasNext()).isTrue(); + + // Test 5: Verify all three documents are indexed + GlobalIndexResult allDocsResult = reader.search(new float[] {0.5f, 0.5f}, 3); + List allResultIds = new ArrayList<>(); + for (Long id : allDocsResult) { + allResultIds.add(id); + } + + assertThat(allResultIds.size()) + .as("Should be able to retrieve all 3 documents") + .isGreaterThanOrEqualTo(1) + .isLessThanOrEqualTo(3); + + reader.close(); + } + + private Options createDefaultOptions() { + Options options = new Options(); + options.setInteger("vector.dim", 128); + options.setString("vector.metric", "COSINE"); + options.setInteger("vector.M", 16); + options.setInteger("vector.ef-construction", 100); + return options; + } + + private List generateRandomVectors(int count, int dimension) { + Random random = new Random(42); // Fixed seed for reproducibility + List vectors = new ArrayList<>(); + for (int i = 0; i < count; i++) { + float[] vector = new float[dimension]; + for (int j = 0; j < dimension; j++) { + vector[j] = random.nextFloat() * 2 - 1; // Range [-1, 1] + } + // Normalize for cosine similarity + normalize(vector); + vectors.add(vector); + } + return vectors; + } + + private void normalize(float[] vector) { + float norm = 0; + for (float v : vector) { + norm += v * v; + } + norm = (float) Math.sqrt(norm); + if (norm > 0) { + for (int i = 0; i < vector.length; i++) { + vector[i] /= norm; + } + } + } + + private boolean containsRowId(GlobalIndexResult result, long rowId) { + for (Long id : result) { + if (id == rowId) { + return true; + } + } + return false; + } +} From 4401379e32213d9b6053b6466820f6cec3b1b620 Mon Sep 17 00:00:00 2001 From: yantian Date: Mon, 1 Dec 2025 17:29:54 +0800 Subject: [PATCH 03/11] fix --- pom.xml | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 0dc9608c45f3..5695a1049a1a 100644 --- a/pom.xml +++ b/pom.xml @@ -63,7 +63,6 @@ under the License. paimon-filesystems paimon-format paimon-lance - paimon-vector paimon-vfs paimon-bundle paimon-hive @@ -506,6 +505,16 @@ under the License. 11 + + + paimon-vector + + paimon-vector + + + 11 + + From 7ace2abb2181687e2bb0406fb311bb1bc135cf84 Mon Sep 17 00:00:00 2001 From: yantian Date: Tue, 2 Dec 2025 11:28:35 +0800 Subject: [PATCH 04/11] fix --- .../org/apache/paimon/vector/IndexShard.java | 45 +++++ .../PaimonRandomAccessVectorValues.java | 65 ++++++++ .../vector/VectorGlobalIndexReader.java | 156 ++++-------------- .../vector/VectorGlobalIndexWriter.java | 44 ++--- .../paimon/vector/VectorIndexMetadata.java | 84 ++++++++++ .../paimon/vector/VectorGlobalIndexTest.java | 57 ------- 6 files changed, 235 insertions(+), 216 deletions(-) create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/IndexShard.java create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/PaimonRandomAccessVectorValues.java create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java delete mode 100644 paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/IndexShard.java b/paimon-vector/src/main/java/org/apache/paimon/vector/IndexShard.java new file mode 100644 index 000000000000..78080c35aef9 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/IndexShard.java @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import io.github.jbellis.jvector.graph.OnHeapGraphIndex; + +import java.util.List; + +/** A shard of the vector index containing row IDs, vectors, metadata, and the graph index. */ +public class IndexShard { + final long[] rowIds; + final List vectors; + final VectorIndexMetadata metadata; + final long rowRangeStart; + final OnHeapGraphIndex graphIndex; + + IndexShard( + long[] rowIds, + List vectors, + VectorIndexMetadata metadata, + long rowRangeStart, + OnHeapGraphIndex graphIndex) { + this.rowIds = rowIds; + this.vectors = vectors; + this.metadata = metadata; + this.rowRangeStart = rowRangeStart; + this.graphIndex = graphIndex; + } +} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/PaimonRandomAccessVectorValues.java b/paimon-vector/src/main/java/org/apache/paimon/vector/PaimonRandomAccessVectorValues.java new file mode 100644 index 000000000000..bdf0700aca5b --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/PaimonRandomAccessVectorValues.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.vector.types.VectorFloat; + +import java.util.List; + +/** Implementation of RandomAccessVectorValues for Paimon vectors. */ +public class PaimonRandomAccessVectorValues implements RandomAccessVectorValues { + private final int size; + private final int dimension; + private final List> vectors; + private final boolean isValueShared; + + public PaimonRandomAccessVectorValues( + int size, int dimension, List> vectors, boolean isValueShared) { + this.size = size; + this.dimension = dimension; + this.vectors = vectors; + this.isValueShared = isValueShared; + } + + @Override + public int size() { + return this.size; + } + + @Override + public int dimension() { + return this.dimension; + } + + @Override + public VectorFloat getVector(int i) { + return this.vectors.get(i); + } + + @Override + public boolean isValueShared() { + return this.isValueShared; + } + + @Override + public RandomAccessVectorValues copy() { + return this; + } +} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java index ba8c70591974..3c40d2e59997 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -29,7 +29,6 @@ import io.github.jbellis.jvector.graph.GraphIndexBuilder; import io.github.jbellis.jvector.graph.GraphSearcher; -import io.github.jbellis.jvector.graph.OnHeapGraphIndex; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.SearchResult; import io.github.jbellis.jvector.util.Bits; @@ -122,16 +121,9 @@ private IndexShard deserializeIndex(byte[] indexBytes, GlobalIndexIOMeta meta) // Parse metadata byte[] metaBytes = meta.metadata(); - VectorIndexMetadata metadata = deserializeMetadata(metaBytes); + VectorIndexMetadata metadata = VectorIndexMetadata.deserializeMetadata(metaBytes); - // Rebuild graph index if version 2 (with JVector) - OnHeapGraphIndex graphIndex = null; - if (version == 2 && dataIn.available() > 0) { - // Read graph structure - int graphSize = dataIn.readInt(); - int maxDegree = dataIn.readInt(); - - // Rebuild the graph index from vectors + if (dataIn.available() > 0) { var vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); List> vectorFloats = new ArrayList<>(); for (float[] vec : vectors) { @@ -140,64 +132,29 @@ private IndexShard deserializeIndex(byte[] indexBytes, GlobalIndexIOMeta meta) // Create RandomAccessVectorValues RandomAccessVectorValues vectorValues = - new RandomAccessVectorValues() { - @Override - public int size() { - return vectorFloats.size(); - } - - @Override - public int dimension() { - return metadata.dimension; - } - - @Override - public VectorFloat getVector(int i) { - return vectorFloats.get(i); - } - - @Override - public boolean isValueShared() { - return false; - } - - @Override - public RandomAccessVectorValues copy() { - return this; - } - }; + new PaimonRandomAccessVectorValues( + vectorFloats.size(), metadata.dimension(), vectorFloats, false); // Rebuild graph using stored structure or rebuild from scratch VectorSimilarityFunction similarityFunction = - VectorSimilarityFunction.valueOf(metadata.similarityFunction); - var builder = + VectorSimilarityFunction.valueOf(metadata.similarityFunction()); + try (GraphIndexBuilder builder = new GraphIndexBuilder( vectorValues, similarityFunction, - metadata.m, - metadata.efConstruction, + metadata.m(), + metadata.efConstruction(), 1.0f, // todo: need conf - 1.0f); // todo: need conf - graphIndex = builder.build(vectorValues); + 1.0f)) { // todo: need conf + return new IndexShard( + rowIds, + vectors, + metadata, + meta.rowIdRange().from, + builder.build(vectorValues)); + } } - - return new IndexShard(rowIds, vectors, metadata, meta.rowIdRange().from, graphIndex); - } - } - - private VectorIndexMetadata deserializeMetadata(byte[] metaBytes) throws IOException { - if (metaBytes == null || metaBytes.length == 0) { - return new VectorIndexMetadata( - this.dimension, this.similarityFunction, this.m, this.efConstruction); - } - - ByteArrayInputStream byteIn = new ByteArrayInputStream(metaBytes); - try (DataInputStream dataIn = new DataInputStream(byteIn)) { - int dimension = dataIn.readInt(); - String metricName = dataIn.readUTF(); - int m = dataIn.readInt(); - int efConstruction = dataIn.readInt(); - return new VectorIndexMetadata(dimension, metricName, m, efConstruction); + return new IndexShard(rowIds, vectors, metadata, meta.rowIdRange().from, null); } } @@ -222,43 +179,21 @@ public GlobalIndexResult search(float[] query, int k) { // Use JVector's GraphSearcher for efficient ANN search VectorSimilarityFunction similarityFunction = - VectorSimilarityFunction.valueOf(shard.metadata.similarityFunction); + VectorSimilarityFunction.valueOf(shard.metadata.similarityFunction()); // Create vector values for search context - var vectorTypeSupport2 = VectorizationProvider.getInstance().getVectorTypeSupport(); + var vectorTypeSupportForIndex = + VectorizationProvider.getInstance().getVectorTypeSupport(); List> vectorFloats = new ArrayList<>(); for (float[] vec : shard.vectors) { - vectorFloats.add(vectorTypeSupport2.createFloatVector(vec)); + vectorFloats.add(vectorTypeSupportForIndex.createFloatVector(vec)); } - RandomAccessVectorValues vectorValues = - new RandomAccessVectorValues() { - @Override - public int size() { - return vectorFloats.size(); - } - - @Override - public int dimension() { - return shard.metadata.dimension; - } - - @Override - public VectorFloat getVector(int i) { - return vectorFloats.get(i); - } - - @Override - public boolean isValueShared() { - return false; - } - - @Override - public RandomAccessVectorValues copy() { - return this; - } - }; - + new PaimonRandomAccessVectorValues( + vectorFloats.size(), + shard.metadata.dimension(), + vectorFloats, + false); // Search using static method SearchResult result = GraphSearcher.search( @@ -272,7 +207,9 @@ public RandomAccessVectorValues copy() { // Collect row IDs from results var nodes = result.getNodes(); for (int i = 0; i < nodes.length && i < k; i++) { - int nodeId = nodes[i].node; + // todo: could get score here: nodeScore.score + SearchResult.NodeScore nodeScore = nodes[i]; + int nodeId = nodeScore.node; if (nodeId >= 0 && nodeId < shard.rowIds.length) { resultIds.add(shard.rowIds[nodeId]); } @@ -360,39 +297,4 @@ public Optional visitIn(FieldRef fieldRef, List liter public Optional visitNotIn(FieldRef fieldRef, List literals) { throw new UnsupportedOperationException("Vector index does not support notIn predicate"); } - - private static class IndexShard { - final long[] rowIds; - final List vectors; - final VectorIndexMetadata metadata; - final long rowRangeStart; - final OnHeapGraphIndex graphIndex; - - IndexShard( - long[] rowIds, - List vectors, - VectorIndexMetadata metadata, - long rowRangeStart, - OnHeapGraphIndex graphIndex) { - this.rowIds = rowIds; - this.vectors = vectors; - this.metadata = metadata; - this.rowRangeStart = rowRangeStart; - this.graphIndex = graphIndex; - } - } - - private static class VectorIndexMetadata { - final int dimension; - final String similarityFunction; - final int m; - final int efConstruction; - - VectorIndexMetadata(int dimension, String similarityFunction, int m, int efConstruction) { - this.dimension = dimension; - this.similarityFunction = similarityFunction; - this.m = m; - this.efConstruction = efConstruction; - } - } } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java index f4b6f47a07a1..0aa27f360d0d 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -124,11 +124,15 @@ public List finish() { // Serialize index and metadata byte[] indexBytes = serializeIndexWithJVector(graphIndex); - byte[] metaBytes = serializeMetadata(); - - // Write to file. todo: here could use OnDiskGraphIndex to write directly to file. - // RandomAccessVectorValues ravv - // OnDiskGraphIndex.write(index, ravv, indexPath); + VectorIndexMetadata metadata = + new VectorIndexMetadata( + vectorOptions.dimension(), + vectorOptions.metric(), + vectorOptions.m(), + vectorOptions.efConstruction()); + byte[] metaBytes = VectorIndexMetadata.serializeMetadata(metadata); + + // Write to file. String fileName = fileWriter.newFileName(VectorGlobalIndexerFactory.IDENTIFIER); try (OutputStream out = fileWriter.newOutputStream(fileName)) { out.write(indexBytes); @@ -154,32 +158,8 @@ private OnHeapGraphIndex buildJVectorIndex() { // Create vector values adapter for JVector RandomAccessVectorValues vectorValues = - new RandomAccessVectorValues() { - @Override - public int size() { - return vectorFloats.size(); - } - - @Override - public int dimension() { - return vectorOptions.dimension(); - } - - @Override - public VectorFloat getVector(int i) { - return vectorFloats.get(i); - } - - @Override - public boolean isValueShared() { - return false; - } - - @Override - public RandomAccessVectorValues copy() { - return this; - } - }; + new PaimonRandomAccessVectorValues( + vectorFloats.size(), vectorOptions.dimension(), vectorFloats, false); // Build HNSW graph index using JVector GraphIndexBuilder builder = @@ -199,7 +179,7 @@ private byte[] serializeIndexWithJVector(OnHeapGraphIndex graphIndex) throws IOE ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); try (DataOutputStream dataOut = new DataOutputStream(byteOut)) { // Write version - dataOut.writeInt(2); // Version 2 with JVector index + dataOut.writeInt(1); // Version 1 with JVector index // Write row IDs dataOut.writeInt(vectors.size()); diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java new file mode 100644 index 000000000000..7e990117d53c --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java @@ -0,0 +1,84 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.io.Serializable; + +/** Metadata for the vector index, including dimension, similarity function, and HNSW parameters. */ +public class VectorIndexMetadata implements Serializable { + + private static final long serialVersionUID = 1L; + + private final int dimension; + private final String similarityFunction; + private final int m; + private final int efConstruction; + + public VectorIndexMetadata( + int dimension, String similarityFunction, int m, int efConstruction) { + this.dimension = dimension; + this.similarityFunction = similarityFunction; + this.m = m; + this.efConstruction = efConstruction; + } + + public static VectorIndexMetadata deserializeMetadata(byte[] metaBytes) throws IOException { + ByteArrayInputStream byteIn = new ByteArrayInputStream(metaBytes); + try (DataInputStream dataIn = new DataInputStream(byteIn)) { + int dimension = dataIn.readInt(); + String metricName = dataIn.readUTF(); + int m = dataIn.readInt(); + int efConstruction = dataIn.readInt(); + return new VectorIndexMetadata(dimension, metricName, m, efConstruction); + } + } + + public static byte[] serializeMetadata(VectorIndexMetadata vectorIndexMetadata) + throws IOException { + ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); + try (DataOutputStream dataOut = new DataOutputStream(byteOut)) { + dataOut.writeInt(vectorIndexMetadata.dimension()); + dataOut.writeUTF(vectorIndexMetadata.similarityFunction()); + dataOut.writeInt(vectorIndexMetadata.m()); + dataOut.writeInt(vectorIndexMetadata.efConstruction()); + } + return byteOut.toByteArray(); + } + + public int dimension() { + return dimension; + } + + public String similarityFunction() { + return similarityFunction; + } + + public int m() { + return m; + } + + public int efConstruction() { + return efConstruction; + } +} diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java deleted file mode 100644 index 248d24b08ae0..000000000000 --- a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java +++ /dev/null @@ -1,57 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.paimon.vector; - -import org.apache.paimon.globalindex.GlobalIndexer; -import org.apache.paimon.globalindex.GlobalIndexerFactory; -import org.apache.paimon.globalindex.GlobalIndexerFactoryUtils; -import org.apache.paimon.options.Options; -import org.apache.paimon.types.ArrayType; -import org.apache.paimon.types.DataType; -import org.apache.paimon.types.FloatType; - -import org.junit.jupiter.api.Test; - -import static org.assertj.core.api.Assertions.assertThat; - -/** Test for {@link VectorGlobalIndexerFactory}. */ -public class VectorGlobalIndexTest { - - @Test - public void testFactoryRegistration() { - GlobalIndexerFactory factory = - GlobalIndexerFactoryUtils.load(VectorGlobalIndexerFactory.IDENTIFIER); - assertThat(factory).isNotNull(); - assertThat(factory).isInstanceOf(VectorGlobalIndexerFactory.class); - assertThat(factory.identifier()).isEqualTo("vector"); - } - - @Test - public void testCreateIndexer() { - GlobalIndexerFactory factory = new VectorGlobalIndexerFactory(); - DataType vectorType = new ArrayType(new FloatType()); - Options options = new Options(); - options.setInteger("vector.dim", 128); - options.setString("vector.metric", "COSINE"); - - GlobalIndexer indexer = factory.create(vectorType, options); - assertThat(indexer).isNotNull(); - assertThat(indexer).isInstanceOf(VectorGlobalIndexer.class); - } -} From 785754aab1517e2ad74b33ac19786204660bcf2f Mon Sep 17 00:00:00 2001 From: yantian Date: Tue, 2 Dec 2025 15:23:53 +0800 Subject: [PATCH 05/11] change to lucene --- paimon-vector/pom.xml | 14 +- .../org/apache/paimon/vector/IndexShard.java | 45 --- .../PaimonRandomAccessVectorValues.java | 65 ---- .../vector/VectorGlobalIndexReader.java | 288 ++++++++---------- .../vector/VectorGlobalIndexWriter.java | 243 +++++++++------ .../paimon/vector/VectorGlobalIndexer.java | 4 +- .../vector/VectorGlobalIndexerFactory.java | 2 +- ...erTest.java => VectorGlobalIndexTest.java} | 28 +- 8 files changed, 295 insertions(+), 394 deletions(-) delete mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/IndexShard.java delete mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/PaimonRandomAccessVectorValues.java rename paimon-vector/src/test/java/org/apache/paimon/vector/{VectorGlobalIndexWriterReaderTest.java => VectorGlobalIndexTest.java} (95%) diff --git a/paimon-vector/pom.xml b/paimon-vector/pom.xml index 499aa0d04c3e..274d876ffa4a 100644 --- a/paimon-vector/pom.xml +++ b/paimon-vector/pom.xml @@ -32,7 +32,7 @@ under the License. Paimon : Vector - 3.0.6 + 9.11.1 11 @@ -44,9 +44,15 @@ under the License. - io.github.jbellis - jvector - ${jvector.version} + org.apache.lucene + lucene-core + ${lucene.version} + + + + org.apache.lucene + lucene-codecs + ${lucene.version} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/IndexShard.java b/paimon-vector/src/main/java/org/apache/paimon/vector/IndexShard.java deleted file mode 100644 index 78080c35aef9..000000000000 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/IndexShard.java +++ /dev/null @@ -1,45 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.paimon.vector; - -import io.github.jbellis.jvector.graph.OnHeapGraphIndex; - -import java.util.List; - -/** A shard of the vector index containing row IDs, vectors, metadata, and the graph index. */ -public class IndexShard { - final long[] rowIds; - final List vectors; - final VectorIndexMetadata metadata; - final long rowRangeStart; - final OnHeapGraphIndex graphIndex; - - IndexShard( - long[] rowIds, - List vectors, - VectorIndexMetadata metadata, - long rowRangeStart, - OnHeapGraphIndex graphIndex) { - this.rowIds = rowIds; - this.vectors = vectors; - this.metadata = metadata; - this.rowRangeStart = rowRangeStart; - this.graphIndex = graphIndex; - } -} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/PaimonRandomAccessVectorValues.java b/paimon-vector/src/main/java/org/apache/paimon/vector/PaimonRandomAccessVectorValues.java deleted file mode 100644 index bdf0700aca5b..000000000000 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/PaimonRandomAccessVectorValues.java +++ /dev/null @@ -1,65 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.paimon.vector; - -import io.github.jbellis.jvector.graph.RandomAccessVectorValues; -import io.github.jbellis.jvector.vector.types.VectorFloat; - -import java.util.List; - -/** Implementation of RandomAccessVectorValues for Paimon vectors. */ -public class PaimonRandomAccessVectorValues implements RandomAccessVectorValues { - private final int size; - private final int dimension; - private final List> vectors; - private final boolean isValueShared; - - public PaimonRandomAccessVectorValues( - int size, int dimension, List> vectors, boolean isValueShared) { - this.size = size; - this.dimension = dimension; - this.vectors = vectors; - this.isValueShared = isValueShared; - } - - @Override - public int size() { - return this.size; - } - - @Override - public int dimension() { - return this.dimension; - } - - @Override - public VectorFloat getVector(int i) { - return this.vectors.get(i); - } - - @Override - public boolean isValueShared() { - return this.isValueShared; - } - - @Override - public RandomAccessVectorValues copy() { - return this; - } -} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java index 3c40d2e59997..f13c110e52f3 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -23,58 +23,110 @@ import org.apache.paimon.globalindex.GlobalIndexReader; import org.apache.paimon.globalindex.GlobalIndexResult; import org.apache.paimon.globalindex.io.GlobalIndexFileReader; -import org.apache.paimon.options.Options; import org.apache.paimon.predicate.FieldRef; -import org.apache.paimon.types.DataType; - -import io.github.jbellis.jvector.graph.GraphIndexBuilder; -import io.github.jbellis.jvector.graph.GraphSearcher; -import io.github.jbellis.jvector.graph.RandomAccessVectorValues; -import io.github.jbellis.jvector.graph.SearchResult; -import io.github.jbellis.jvector.util.Bits; -import io.github.jbellis.jvector.vector.VectorSimilarityFunction; -import io.github.jbellis.jvector.vector.VectorizationProvider; -import io.github.jbellis.jvector.vector.types.VectorFloat; - -import java.io.ByteArrayInputStream; -import java.io.DataInputStream; + +import org.apache.lucene.document.Document; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.StoredFields; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.store.MMapDirectory; + import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.UUID; /** - * Vector global index reader using JVector. + * Vector global index reader using Apache Lucene 9.x. * - *

This implementation uses JVector's HNSW graph for efficient approximate nearest neighbor - * search. + *

This implementation uses Lucene's native KnnFloatVectorQuery with HNSW graph for efficient + * approximate nearest neighbor search. */ public class VectorGlobalIndexReader implements GlobalIndexReader { - private final List shards; - private final String similarityFunction; - private final int dimension; - private final int m; - private final int efConstruction; - - public VectorGlobalIndexReader( - GlobalIndexFileReader fileReader, - List files, - DataType fieldType, - Options options) + private static final String VECTOR_FIELD = "vector"; + private static final String ROW_ID_FIELD = "rowId"; + + private final List searchers; + private final List directories; + private final List tempDirs; + + public VectorGlobalIndexReader(GlobalIndexFileReader fileReader, List files) throws IOException { - this.shards = new ArrayList<>(); - VectorIndexOptions vectorOptions = new VectorIndexOptions(options); - this.similarityFunction = vectorOptions.metric(); - this.dimension = vectorOptions.dimension(); - this.m = vectorOptions.m(); - this.efConstruction = vectorOptions.efConstruction(); - loadShards(fileReader, files); + this.searchers = new ArrayList<>(); + this.directories = new ArrayList<>(); + this.tempDirs = new ArrayList<>(); + loadIndices(fileReader, files); + } + + /** + * Search for similar vectors using Lucene KNN search. + * + * @param query query vector + * @param k number of results + * @return global index result containing row IDs + */ + public GlobalIndexResult search(float[] query, int k) { + Set resultIds = new HashSet<>(); + + for (IndexSearcher searcher : searchers) { + try { + // Create KNN query + KnnFloatVectorQuery knnQuery = new KnnFloatVectorQuery(VECTOR_FIELD, query, k); + + // Execute search + TopDocs topDocs = searcher.search(knnQuery, k); + StoredFields storedFields = searcher.storedFields(); + Set fieldsToLoad = Set.of(ROW_ID_FIELD); + // Collect row IDs from results + for (org.apache.lucene.search.ScoreDoc scoreDoc : topDocs.scoreDocs) { + float rawScore = scoreDoc.score; + Document doc = storedFields.document(scoreDoc.doc, fieldsToLoad); + long rowId = doc.getField(ROW_ID_FIELD).numericValue().longValue(); + resultIds.add(rowId); + } + } catch (IOException e) { + throw new RuntimeException("Failed to search vector index", e); + } + } + + return GlobalIndexResult.wrap(resultIds); + } + + @Override + public void close() throws IOException { + // Close readers + for (IndexSearcher searcher : searchers) { + searcher.getIndexReader().close(); + } + searchers.clear(); + + // Close directories + for (Directory directory : directories) { + directory.close(); + } + directories.clear(); + + // Clean up temp directories + for (java.nio.file.Path tempDir : tempDirs) { + deleteDirectory(tempDir); + } + tempDirs.clear(); } - private void loadShards(GlobalIndexFileReader fileReader, List files) + private void loadIndices(GlobalIndexFileReader fileReader, List files) throws IOException { for (GlobalIndexIOMeta meta : files) { try (SeekableInputStream in = fileReader.getInputStream(meta.fileName())) { @@ -88,143 +140,61 @@ private void loadShards(GlobalIndexFileReader fileReader, List vectors = new ArrayList<>(); - for (int i = 0; i < size; i++) { - int dim = dataIn.readInt(); - float[] vector = new float[dim]; - for (int j = 0; j < dim; j++) { - vector[j] = dataIn.readFloat(); - } - vectors.add(vector); - } - - // Parse metadata - byte[] metaBytes = meta.metadata(); - VectorIndexMetadata metadata = VectorIndexMetadata.deserializeMetadata(metaBytes); - - if (dataIn.available() > 0) { - var vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); - List> vectorFloats = new ArrayList<>(); - for (float[] vec : vectors) { - vectorFloats.add(vectorTypeSupport.createFloatVector(vec)); - } - - // Create RandomAccessVectorValues - RandomAccessVectorValues vectorValues = - new PaimonRandomAccessVectorValues( - vectorFloats.size(), metadata.dimension(), vectorFloats, false); - - // Rebuild graph using stored structure or rebuild from scratch - VectorSimilarityFunction similarityFunction = - VectorSimilarityFunction.valueOf(metadata.similarityFunction()); - try (GraphIndexBuilder builder = - new GraphIndexBuilder( - vectorValues, - similarityFunction, - metadata.m(), - metadata.efConstruction(), - 1.0f, // todo: need conf - 1.0f)) { // todo: need conf - return new IndexShard( - rowIds, - vectors, - metadata, - meta.rowIdRange().from, - builder.build(vectorValues)); - } - } - return new IndexShard(rowIds, vectors, metadata, meta.rowIdRange().from, null); } } - /** - * Search for similar vectors using JVector HNSW. - * - * @param query query vector - * @param k number of results - * @return global index result containing row IDs - */ - public GlobalIndexResult search(float[] query, int k) { - Set resultIds = new HashSet<>(); - var vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); - VectorFloat queryVector = vectorTypeSupport.createFloatVector(query); + private Directory deserializeDirectory(byte[] data) throws IOException { + // Create temporary directory for MMap + Path tempDir = Files.createTempDirectory("paimon-vector-read" + UUID.randomUUID()); + tempDirs.add(tempDir); + Directory directory = new MMapDirectory(tempDir); - for (IndexShard shard : shards) { - try { - if (shard.graphIndex == null) { - // Fall back to brute force if no graph index - continue; - } + ByteBuffer buffer = ByteBuffer.wrap(data); - // Use JVector's GraphSearcher for efficient ANN search - VectorSimilarityFunction similarityFunction = - VectorSimilarityFunction.valueOf(shard.metadata.similarityFunction()); + // Read number of files + int numFiles = buffer.getInt(); - // Create vector values for search context - var vectorTypeSupportForIndex = - VectorizationProvider.getInstance().getVectorTypeSupport(); - List> vectorFloats = new ArrayList<>(); - for (float[] vec : shard.vectors) { - vectorFloats.add(vectorTypeSupportForIndex.createFloatVector(vec)); - } - RandomAccessVectorValues vectorValues = - new PaimonRandomAccessVectorValues( - vectorFloats.size(), - shard.metadata.dimension(), - vectorFloats, - false); - // Search using static method - SearchResult result = - GraphSearcher.search( - queryVector, - k, - vectorValues, - similarityFunction, - shard.graphIndex, - Bits.ALL); + for (int i = 0; i < numFiles; i++) { + // Read file name + int nameLength = buffer.getInt(); + byte[] nameBytes = new byte[nameLength]; + buffer.get(nameBytes); + String fileName = new String(nameBytes, StandardCharsets.UTF_8); - // Collect row IDs from results - var nodes = result.getNodes(); - for (int i = 0; i < nodes.length && i < k; i++) { - // todo: could get score here: nodeScore.score - SearchResult.NodeScore nodeScore = nodes[i]; - int nodeId = nodeScore.node; - if (nodeId >= 0 && nodeId < shard.rowIds.length) { - resultIds.add(shard.rowIds[nodeId]); - } - } - } catch (Exception e) { - throw new RuntimeException("Failed to search vector index", e); + // Read file content + long fileLength = buffer.getLong(); + byte[] fileContent = new byte[(int) fileLength]; + buffer.get(fileContent); + + // Write to directory + try (IndexOutput output = directory.createOutput(fileName, null)) { + output.writeBytes(fileContent, 0, fileContent.length); } } - return GlobalIndexResult.wrap(resultIds); + return directory; } - @Override - public void close() throws IOException { - shards.clear(); + private void deleteDirectory(java.nio.file.Path path) throws IOException { + if (java.nio.file.Files.exists(path)) { + java.nio.file.Files.walk(path) + .sorted(java.util.Comparator.reverseOrder()) + .forEach( + p -> { + try { + java.nio.file.Files.delete(p); + } catch (IOException e) { + // Ignore cleanup errors + } + }); + } } // Implementation of FunctionVisitor methods diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java index 0aa27f360d0d..a782d434a456 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -26,30 +26,38 @@ import org.apache.paimon.utils.Preconditions; import org.apache.paimon.utils.Range; -import io.github.jbellis.jvector.graph.GraphIndexBuilder; -import io.github.jbellis.jvector.graph.OnHeapGraphIndex; -import io.github.jbellis.jvector.graph.RandomAccessVectorValues; -import io.github.jbellis.jvector.vector.VectorSimilarityFunction; -import io.github.jbellis.jvector.vector.VectorizationProvider; -import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StoredField; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.MMapDirectory; import java.io.ByteArrayOutputStream; -import java.io.DataOutputStream; import java.io.IOException; import java.io.OutputStream; +import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; import static org.apache.paimon.utils.Preconditions.checkArgument; /** - * Vector global index writer using JVector. + * Vector global index writer using Apache Lucene 9.x. * - *

This implementation uses JVector's HNSW (Hierarchical Navigable Small World) graph algorithm - * for efficient approximate nearest neighbor search. + *

This implementation uses Lucene's native KnnFloatVectorField with HNSW algorithm for efficient + * approximate nearest neighbor search. */ public class VectorGlobalIndexWriter implements GlobalIndexWriter { + private static final String VECTOR_FIELD = "vector"; + private static final String ROW_ID_FIELD = "rowId"; + private final GlobalIndexFileWriter fileWriter; private final DataType fieldType; private final VectorIndexOptions vectorOptions; @@ -67,30 +75,12 @@ public VectorGlobalIndexWriter( // Parse options this.vectorOptions = new VectorIndexOptions(options); - this.similarityFunction = parseMetricToJVector(vectorOptions.metric()); + this.similarityFunction = parseMetricToLucene(vectorOptions.metric()); // Validate field type validateFieldType(fieldType); } - private void validateFieldType(DataType type) { - checkArgument( - type instanceof ArrayType, "Vector field type must be ARRAY, but was: " + type); - } - - private VectorSimilarityFunction parseMetricToJVector(String metric) { - switch (metric.toUpperCase()) { - case "COSINE": - return VectorSimilarityFunction.COSINE; - case "DOT_PRODUCT": - return VectorSimilarityFunction.DOT_PRODUCT; - case "EUCLIDEAN": - return VectorSimilarityFunction.EUCLIDEAN; - default: - throw new IllegalArgumentException("Unsupported metric: " + metric); - } - } - @Override public void write(Object key) { Preconditions.checkArgument( @@ -119,11 +109,11 @@ public List finish() { return new ArrayList<>(); } - // Build JVector HNSW index - OnHeapGraphIndex graphIndex = buildJVectorIndex(); + // Build Lucene index in memory + byte[] indexBytes = + buildLuceneIndex(this.vectorOptions.m(), this.vectorOptions.efConstruction()); - // Serialize index and metadata - byte[] indexBytes = serializeIndexWithJVector(graphIndex); + // Create metadata VectorIndexMetadata metadata = new VectorIndexMetadata( vectorOptions.dimension(), @@ -132,7 +122,7 @@ public List finish() { vectorOptions.efConstruction()); byte[] metaBytes = VectorIndexMetadata.serializeMetadata(metadata); - // Write to file. + // Write to file String fileName = fileWriter.newFileName(VectorGlobalIndexerFactory.IDENTIFIER); try (OutputStream out = fileWriter.newOutputStream(fileName)) { out.write(indexBytes); @@ -146,88 +136,141 @@ public List finish() { } } - private OnHeapGraphIndex buildJVectorIndex() { - // Get the vectorization provider - var vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); + private void validateFieldType(DataType type) { + checkArgument( + type instanceof ArrayType, "Vector field type must be ARRAY, but was: " + type); + } + + private VectorSimilarityFunction parseMetricToLucene(String metric) { + switch (metric.toUpperCase()) { + case "COSINE": + return VectorSimilarityFunction.COSINE; + case "DOT_PRODUCT": + return VectorSimilarityFunction.DOT_PRODUCT; + case "EUCLIDEAN": + return VectorSimilarityFunction.EUCLIDEAN; + case "MAX_INNER_PRODUCT": + return VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; + default: + throw new IllegalArgumentException("Unsupported metric: " + metric); + } + } + + private byte[] buildLuceneIndex(int m, int efConstruction) throws IOException { + // Create temporary directory for MMap + java.nio.file.Path tempDir = + java.nio.file.Files.createTempDirectory( + fileWriter.newFileName("paimon-vector-index")); + Directory directory = new MMapDirectory(tempDir); + + try { + // Configure index writer + IndexWriterConfig config = new IndexWriterConfig(); + config.setRAMBufferSizeMB(256); // Increase buffer for better performance + config.setCodec( + new Lucene99Codec() { + @Override + public KnnVectorsFormat getKnnVectorsFormatForField(String field) { + return new Lucene99HnswVectorsFormat(m, efConstruction); + } + }); + + try (IndexWriter writer = new IndexWriter(directory, config)) { + // Add each vector as a document + for (VectorWithRowId vectorWithRowId : vectors) { + Document doc = new Document(); + + // Add KNN vector field + doc.add( + new KnnFloatVectorField( + VECTOR_FIELD, vectorWithRowId.vector, similarityFunction)); + + // Store row ID + doc.add(new StoredField(ROW_ID_FIELD, vectorWithRowId.rowId)); + + writer.addDocument(doc); + } - // Create list of VectorFloat for JVector - List> vectorFloats = new ArrayList<>(); - for (VectorWithRowId vec : vectors) { - vectorFloats.add(vectorTypeSupport.createFloatVector(vec.vector)); + // Commit changes + writer.commit(); + } + + // Serialize directory to byte array + return serializeDirectory(directory); + } finally { + // Clean up + directory.close(); + deleteDirectory(tempDir); } + } - // Create vector values adapter for JVector - RandomAccessVectorValues vectorValues = - new PaimonRandomAccessVectorValues( - vectorFloats.size(), vectorOptions.dimension(), vectorFloats, false); - - // Build HNSW graph index using JVector - GraphIndexBuilder builder = - new GraphIndexBuilder( - vectorValues, - similarityFunction, - vectorOptions.m(), - vectorOptions.efConstruction(), - 1.0f, // alpha (diversityWeight) - 1.0f // neighborOverflow - ); - - return builder.build(vectorValues); + private void deleteDirectory(java.nio.file.Path path) throws IOException { + if (java.nio.file.Files.exists(path)) { + java.nio.file.Files.walk(path) + .sorted(java.util.Comparator.reverseOrder()) + .forEach( + p -> { + try { + java.nio.file.Files.delete(p); + } catch (IOException e) { + // Ignore cleanup errors + } + }); + } } - private byte[] serializeIndexWithJVector(OnHeapGraphIndex graphIndex) throws IOException { - ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); - try (DataOutputStream dataOut = new DataOutputStream(byteOut)) { - // Write version - dataOut.writeInt(1); // Version 1 with JVector index + private byte[] serializeDirectory(Directory directory) throws IOException { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); - // Write row IDs - dataOut.writeInt(vectors.size()); - for (VectorWithRowId vector : vectors) { - dataOut.writeLong(vector.rowId); - } + // Write all files from the directory + String[] files = directory.listAll(); - // Write vectors (still needed for graph index reconstruction) - for (VectorWithRowId vector : vectors) { - dataOut.writeInt(vector.vector.length); - for (float v : vector.vector) { - dataOut.writeFloat(v); - } - } + // Write number of files + baos.write(intToBytes(files.length)); - // Write JVector graph structure - dataOut.writeInt(graphIndex.size()); - dataOut.writeInt(graphIndex.maxDegree()); + for (String fileName : files) { + // Write file name length and name + byte[] nameBytes = fileName.getBytes(StandardCharsets.UTF_8); + baos.write(intToBytes(nameBytes.length)); + baos.write(nameBytes); - // Serialize graph neighbors for each node - for (int i = 0; i < graphIndex.size(); i++) { - var view = graphIndex.getView(); - var neighbors = view.getNeighborsIterator(i); + // Write file content length and content + long fileLength = directory.fileLength(fileName); + baos.write(longToBytes(fileLength)); - // Count neighbors - List neighborList = new ArrayList<>(); - while (neighbors.hasNext()) { - neighborList.add(neighbors.nextInt()); - } + try (org.apache.lucene.store.IndexInput input = directory.openInput(fileName, null)) { + byte[] buffer = new byte[8192]; + long remaining = fileLength; - dataOut.writeInt(neighborList.size()); - for (int neighbor : neighborList) { - dataOut.writeInt(neighbor); + while (remaining > 0) { + int toRead = (int) Math.min(buffer.length, remaining); + input.readBytes(buffer, 0, toRead); + baos.write(buffer, 0, toRead); + remaining -= toRead; } } } - return byteOut.toByteArray(); + + return baos.toByteArray(); } - private byte[] serializeMetadata() throws IOException { - ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); - try (DataOutputStream dataOut = new DataOutputStream(byteOut)) { - dataOut.writeInt(vectorOptions.dimension()); - dataOut.writeUTF(similarityFunction.name()); - dataOut.writeInt(vectorOptions.m()); - dataOut.writeInt(vectorOptions.efConstruction()); - } - return byteOut.toByteArray(); + private byte[] intToBytes(int value) { + return new byte[] { + (byte) (value >>> 24), (byte) (value >>> 16), (byte) (value >>> 8), (byte) value + }; + } + + private byte[] longToBytes(long value) { + return new byte[] { + (byte) (value >>> 56), + (byte) (value >>> 48), + (byte) (value >>> 40), + (byte) (value >>> 32), + (byte) (value >>> 24), + (byte) (value >>> 16), + (byte) (value >>> 8), + (byte) value + }; } /** Key type for vector indexing. */ diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexer.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexer.java index f6fc88050d14..2b10306a9cd8 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexer.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexer.java @@ -30,7 +30,7 @@ import java.io.IOException; import java.util.List; -/** Vector global indexer using JVector. */ +/** Vector global indexer. */ public class VectorGlobalIndexer implements GlobalIndexer { private final DataType fieldType; @@ -49,6 +49,6 @@ public GlobalIndexWriter createWriter(GlobalIndexFileWriter fileWriter) throws I @Override public GlobalIndexReader createReader( GlobalIndexFileReader fileReader, List files) throws IOException { - return new VectorGlobalIndexReader(fileReader, files, fieldType, options); + return new VectorGlobalIndexReader(fileReader, files); } } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexerFactory.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexerFactory.java index bd2215ae1a6a..b76864f30f6c 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexerFactory.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexerFactory.java @@ -23,7 +23,7 @@ import org.apache.paimon.options.Options; import org.apache.paimon.types.DataType; -/** Factory for creating vector global indexers using JVector. */ +/** Factory for creating vector global indexers. */ public class VectorGlobalIndexerFactory implements GlobalIndexerFactory { public static final String IDENTIFIER = "vector"; diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexWriterReaderTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java similarity index 95% rename from paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexWriterReaderTest.java rename to paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java index f2b000b3a610..f44fc028b327 100644 --- a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexWriterReaderTest.java +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java @@ -47,8 +47,8 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; -/** Test for {@link VectorGlobalIndexWriter} and {@link VectorGlobalIndexReader} with JVector. */ -public class VectorGlobalIndexWriterReaderTest { +/** Test for {@link VectorGlobalIndexWriter} and {@link VectorGlobalIndexReader}. */ +public class VectorGlobalIndexTest { @TempDir java.nio.file.Path tempDir; @@ -121,8 +121,7 @@ public void testWriteAndReadVectorIndex() throws IOException { result.rowRange(), result.meta())); - VectorGlobalIndexReader reader = - new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Test search - query with first vector should return itself GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); @@ -171,8 +170,7 @@ public void testSearchSimilarVectors() throws IOException { result.rowRange(), result.meta())); - VectorGlobalIndexReader reader = - new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Search with base vector should return similar vectors GlobalIndexResult searchResult = reader.search(baseVector, 10); @@ -216,8 +214,7 @@ public void testDifferentSimilarityFunctions() throws IOException { result.rowRange(), result.meta())); - VectorGlobalIndexReader reader = - new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Verify search works with this metric GlobalIndexResult searchResult = reader.search(testVectors.get(0), 3); @@ -259,8 +256,7 @@ public void testDifferentDimensions() throws IOException { result.rowRange(), result.meta())); - VectorGlobalIndexReader reader = - new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Verify search works with this dimension GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); @@ -333,8 +329,7 @@ public void testHNSWParameters() throws IOException { result.rowRange(), result.meta())); - VectorGlobalIndexReader reader = - new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Verify search works with custom HNSW parameters GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); @@ -372,8 +367,7 @@ public void testLargeK() throws IOException { result.rowRange(), result.meta())); - VectorGlobalIndexReader reader = - new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Search with k larger than number of vectors GlobalIndexResult searchResult = reader.search(testVectors.get(0), 200); @@ -445,8 +439,7 @@ public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) } }; - VectorGlobalIndexReader reader = - new VectorGlobalIndexReader(fileReader, allMetas, vectorType, options); + VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, allMetas); // Test 1: Search with various k values int[] kValues = {1, 5, 10, 20, 50}; @@ -550,8 +543,7 @@ public void testVectorIndexEndToEnd() throws IOException { result.rowRange(), result.meta())); - VectorGlobalIndexReader reader = - new VectorGlobalIndexReader(fileReader, metas, vectorType, options); + VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Test 1: Query with vector similar to "Apple" - should find Apple and Banana // Query: [0.85, 0.15] is between Apple and Banana From 62f7db52b08eba3ae7d21f2acfcce7a69c37ba35 Mon Sep 17 00:00:00 2001 From: yantian Date: Tue, 2 Dec 2025 15:58:03 +0800 Subject: [PATCH 06/11] fix --- .../vector/VectorGlobalIndexWriter.java | 25 +++++++++++-------- .../paimon/vector/VectorIndexMetadata.java | 13 ---------- .../paimon/vector/VectorIndexOptions.java | 4 +-- .../paimon/vector/VectorGlobalIndexTest.java | 8 +++--- 4 files changed, 21 insertions(+), 29 deletions(-) diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java index a782d434a456..61762228bd32 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -28,7 +28,7 @@ import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99Codec; -import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; +import org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat; import org.apache.lucene.document.Document; import org.apache.lucene.document.KnnFloatVectorField; import org.apache.lucene.document.StoredField; @@ -165,15 +165,7 @@ private byte[] buildLuceneIndex(int m, int efConstruction) throws IOException { try { // Configure index writer - IndexWriterConfig config = new IndexWriterConfig(); - config.setRAMBufferSizeMB(256); // Increase buffer for better performance - config.setCodec( - new Lucene99Codec() { - @Override - public KnnVectorsFormat getKnnVectorsFormatForField(String field) { - return new Lucene99HnswVectorsFormat(m, efConstruction); - } - }); + IndexWriterConfig config = getIndexWriterConfig(m, efConstruction); try (IndexWriter writer = new IndexWriter(directory, config)) { // Add each vector as a document @@ -204,6 +196,19 @@ public KnnVectorsFormat getKnnVectorsFormatForField(String field) { } } + private static IndexWriterConfig getIndexWriterConfig(int m, int efConstruction) { + IndexWriterConfig config = new IndexWriterConfig(); + config.setRAMBufferSizeMB(256); // Increase buffer for better performance + config.setCodec( + new Lucene99Codec(Lucene99Codec.Mode.BEST_SPEED) { + @Override + public KnnVectorsFormat getKnnVectorsFormatForField(String field) { + return new Lucene99HnswScalarQuantizedVectorsFormat(m, efConstruction); + } + }); + return config; + } + private void deleteDirectory(java.nio.file.Path path) throws IOException { if (java.nio.file.Files.exists(path)) { java.nio.file.Files.walk(path) diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java index 7e990117d53c..08321d278942 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java @@ -18,9 +18,7 @@ package org.apache.paimon.vector; -import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; -import java.io.DataInputStream; import java.io.DataOutputStream; import java.io.IOException; import java.io.Serializable; @@ -43,17 +41,6 @@ public VectorIndexMetadata( this.efConstruction = efConstruction; } - public static VectorIndexMetadata deserializeMetadata(byte[] metaBytes) throws IOException { - ByteArrayInputStream byteIn = new ByteArrayInputStream(metaBytes); - try (DataInputStream dataIn = new DataInputStream(byteIn)) { - int dimension = dataIn.readInt(); - String metricName = dataIn.readUTF(); - int m = dataIn.readInt(); - int efConstruction = dataIn.readInt(); - return new VectorIndexMetadata(dimension, metricName, m, efConstruction); - } - } - public static byte[] serializeMetadata(VectorIndexMetadata vectorIndexMetadata) throws IOException { ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java index 79d3c14649f4..571c1afb0344 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java @@ -36,10 +36,10 @@ public class VectorIndexOptions { .stringType() .defaultValue("EUCLIDEAN") .withDescription( - "The similarity metric for vector search (COSINE, DOT_PRODUCT, EUCLIDEAN)"); + "The similarity metric for vector search (COSINE, DOT_PRODUCT, EUCLIDEAN, MAX_INNER_PRODUCT)"); public static final ConfigOption VECTOR_M = - ConfigOptions.key("vector.M") + ConfigOptions.key("vector.m") .intType() .defaultValue(16) .withDescription( diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java index f44fc028b327..862b63809519 100644 --- a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java @@ -300,7 +300,7 @@ public void testEmptyIndex() throws IOException { @Test public void testHNSWParameters() throws IOException { Options options = createDefaultOptions(); - options.setInteger("vector.M", 32); // Larger M for better recall + options.setInteger("vector.m", 32); // Larger M for better recall options.setInteger("vector.ef-construction", 200); // Larger ef for better quality int dimension = 64; @@ -383,7 +383,7 @@ public void testEndToEndMultiShardScenario() throws IOException { int dimension = 128; options.setInteger("vector.dim", dimension); options.setString("vector.metric", "COSINE"); - options.setInteger("vector.M", 24); + options.setInteger("vector.m", 24); options.setInteger("vector.ef-construction", 150); // Simulate a multi-shard scenario with incremental writes @@ -507,7 +507,7 @@ public void testVectorIndexEndToEnd() throws IOException { int dimension = 2; options.setInteger("vector.dim", dimension); options.setString("vector.metric", "EUCLIDEAN"); - options.setInteger("vector.M", 16); + options.setInteger("vector.m", 16); options.setInteger("vector.ef-construction", 100); // Create semantic vectors representing "documents" @@ -605,7 +605,7 @@ private Options createDefaultOptions() { Options options = new Options(); options.setInteger("vector.dim", 128); options.setString("vector.metric", "COSINE"); - options.setInteger("vector.M", 16); + options.setInteger("vector.m", 16); options.setInteger("vector.ef-construction", 100); return options; } From e9e6ed5634b919875d016cf53392719e5e79d367 Mon Sep 17 00:00:00 2001 From: yantian Date: Tue, 2 Dec 2025 16:30:59 +0800 Subject: [PATCH 07/11] fix --- .../paimon/globalindex/GlobalIndexReader.java | 4 +++ .../vector/VectorGlobalIndexReader.java | 30 ++++++++-------- .../paimon/vector/VectorGlobalIndexTest.java | 34 +++++++++---------- 3 files changed, 36 insertions(+), 32 deletions(-) diff --git a/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java b/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java index 2ab40805f84f..09ecbb2821f3 100644 --- a/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java +++ b/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java @@ -41,4 +41,8 @@ default GlobalIndexResult visitOr(List children) { default GlobalIndexResult visit(TransformPredicate predicate) { throw new UnsupportedOperationException("Should not invoke this"); } + + default GlobalIndexResult visit(float[] query, int k) { + throw new UnsupportedOperationException("Should not invoke this"); + } } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java index f13c110e52f3..91a0216b4660 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -44,7 +44,6 @@ import java.util.ArrayList; import java.util.HashSet; import java.util.List; -import java.util.Optional; import java.util.Set; import java.util.UUID; @@ -78,7 +77,8 @@ public VectorGlobalIndexReader(GlobalIndexFileReader fileReader, List resultIds = new HashSet<>(); for (IndexSearcher searcher : searchers) { @@ -199,72 +199,72 @@ private void deleteDirectory(java.nio.file.Path path) throws IOException { // Implementation of FunctionVisitor methods @Override - public Optional visitIsNotNull(FieldRef fieldRef) { + public GlobalIndexResult visitIsNotNull(FieldRef fieldRef) { throw new UnsupportedOperationException( "Vector index does not support isNotNull predicate"); } @Override - public Optional visitIsNull(FieldRef fieldRef) { + public GlobalIndexResult visitIsNull(FieldRef fieldRef) { throw new UnsupportedOperationException("Vector index does not support isNull predicate"); } @Override - public Optional visitStartsWith(FieldRef fieldRef, Object literal) { + public GlobalIndexResult visitStartsWith(FieldRef fieldRef, Object literal) { throw new UnsupportedOperationException( "Vector index does not support startsWith predicate"); } @Override - public Optional visitEndsWith(FieldRef fieldRef, Object literal) { + public GlobalIndexResult visitEndsWith(FieldRef fieldRef, Object literal) { throw new UnsupportedOperationException("Vector index does not support endsWith predicate"); } @Override - public Optional visitContains(FieldRef fieldRef, Object literal) { + public GlobalIndexResult visitContains(FieldRef fieldRef, Object literal) { throw new UnsupportedOperationException("Vector index does not support contains predicate"); } @Override - public Optional visitLessThan(FieldRef fieldRef, Object literal) { + public GlobalIndexResult visitLessThan(FieldRef fieldRef, Object literal) { throw new UnsupportedOperationException("Vector index does not support lessThan predicate"); } @Override - public Optional visitGreaterOrEqual(FieldRef fieldRef, Object literal) { + public GlobalIndexResult visitGreaterOrEqual(FieldRef fieldRef, Object literal) { throw new UnsupportedOperationException( "Vector index does not support greaterOrEqual predicate"); } @Override - public Optional visitNotEqual(FieldRef fieldRef, Object literal) { + public GlobalIndexResult visitNotEqual(FieldRef fieldRef, Object literal) { throw new UnsupportedOperationException("Vector index does not support notEqual predicate"); } @Override - public Optional visitLessOrEqual(FieldRef fieldRef, Object literal) { + public GlobalIndexResult visitLessOrEqual(FieldRef fieldRef, Object literal) { throw new UnsupportedOperationException( "Vector index does not support lessOrEqual predicate"); } @Override - public Optional visitEqual(FieldRef fieldRef, Object literal) { + public GlobalIndexResult visitEqual(FieldRef fieldRef, Object literal) { throw new UnsupportedOperationException("Vector index does not support equal predicate"); } @Override - public Optional visitGreaterThan(FieldRef fieldRef, Object literal) { + public GlobalIndexResult visitGreaterThan(FieldRef fieldRef, Object literal) { throw new UnsupportedOperationException( "Vector index does not support greaterThan predicate"); } @Override - public Optional visitIn(FieldRef fieldRef, List literals) { + public GlobalIndexResult visitIn(FieldRef fieldRef, List literals) { throw new UnsupportedOperationException("Vector index does not support in predicate"); } @Override - public Optional visitNotIn(FieldRef fieldRef, List literals) { + public GlobalIndexResult visitNotIn(FieldRef fieldRef, List literals) { throw new UnsupportedOperationException("Vector index does not support notIn predicate"); } } diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java index 862b63809519..ba8b2668be5b 100644 --- a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java @@ -124,7 +124,7 @@ public void testWriteAndReadVectorIndex() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Test search - query with first vector should return itself - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); + GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 5); assertThat(searchResult).isNotNull(); reader.close(); @@ -173,7 +173,7 @@ public void testSearchSimilarVectors() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Search with base vector should return similar vectors - GlobalIndexResult searchResult = reader.search(baseVector, 10); + GlobalIndexResult searchResult = reader.visit(baseVector, 10); assertThat(searchResult).isNotNull(); reader.close(); @@ -217,7 +217,7 @@ public void testDifferentSimilarityFunctions() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Verify search works with this metric - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 3); + GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 3); assertThat(searchResult).isNotNull(); reader.close(); @@ -259,7 +259,7 @@ public void testDifferentDimensions() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Verify search works with this dimension - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); + GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 5); assertThat(searchResult).isNotNull(); reader.close(); @@ -332,7 +332,7 @@ public void testHNSWParameters() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Verify search works with custom HNSW parameters - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); + GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 5); assertThat(searchResult).isNotNull(); reader.close(); @@ -370,7 +370,7 @@ public void testLargeK() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Search with k larger than number of vectors - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 200); + GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 200); assertThat(searchResult).isNotNull(); assertThat(searchResult).isNotNull(); @@ -444,7 +444,7 @@ public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) // Test 1: Search with various k values int[] kValues = {1, 5, 10, 20, 50}; for (int k : kValues) { - GlobalIndexResult searchResult = reader.search(allTestVectors.get(0), k); + GlobalIndexResult searchResult = reader.visit(allTestVectors.get(0), k); assertThat(searchResult).isNotNull(); assertThat(containsRowId(searchResult, 0)).isTrue(); } @@ -452,7 +452,7 @@ public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) // Test 2: Search with query vectors from different shards for (int shardIdx = 0; shardIdx < shardsCount; shardIdx++) { int queryIdx = shardIdx * vectorsPerShard + 5; - GlobalIndexResult searchResult = reader.search(allTestVectors.get(queryIdx), 10); + GlobalIndexResult searchResult = reader.visit(allTestVectors.get(queryIdx), 10); assertThat(searchResult).isNotNull(); assertThat(containsRowId(searchResult, queryIdx)) .as("Search result should contain the query vector's own ID") @@ -468,7 +468,7 @@ public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) } normalize(slightlyModifiedVector); - GlobalIndexResult similarSearchResult = reader.search(slightlyModifiedVector, 5); + GlobalIndexResult similarSearchResult = reader.visit(slightlyModifiedVector, 5); assertThat(similarSearchResult).isNotNull(); assertThat(containsRowId(similarSearchResult, 50)) .as("Search with very similar vector should find the original") @@ -483,18 +483,18 @@ public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) } normalize(clusterCenter); - GlobalIndexResult clusterSearchResult = reader.search(clusterCenter, 20); + GlobalIndexResult clusterSearchResult = reader.visit(clusterCenter, 20); assertThat(clusterSearchResult).isNotNull(); assertThat(clusterSearchResult.iterator().hasNext()).isTrue(); // Test 5: Edge case - search with k larger than total vectors - GlobalIndexResult largeKResult = reader.search(allTestVectors.get(10), totalVectors * 2); + GlobalIndexResult largeKResult = reader.visit(allTestVectors.get(10), totalVectors * 2); assertThat(largeKResult).isNotNull(); // Test 6: Multiple consecutive searches (stress test) for (int i = 0; i < 20; i++) { int randomIdx = random.nextInt(totalVectors); - GlobalIndexResult result = reader.search(allTestVectors.get(randomIdx), 15); + GlobalIndexResult result = reader.visit(allTestVectors.get(randomIdx), 15); assertThat(result).isNotNull(); } @@ -548,7 +548,7 @@ public void testVectorIndexEndToEnd() throws IOException { // Test 1: Query with vector similar to "Apple" - should find Apple and Banana // Query: [0.85, 0.15] is between Apple and Banana float[] queryVector = new float[] {0.85f, 0.15f}; - GlobalIndexResult searchResult = reader.search(queryVector, 2); + GlobalIndexResult searchResult = reader.visit(queryVector, 2); assertThat(searchResult).isNotNull(); assertThat(searchResult.iterator().hasNext()).as("Search should return results").isTrue(); @@ -565,14 +565,14 @@ public void testVectorIndexEndToEnd() throws IOException { assertThat(resultIds.size()).as("Should return top 2 results").isLessThanOrEqualTo(2); // Test 2: Query with exact match - should find exact document - GlobalIndexResult exactMatchResult = reader.search(appleVector, 1); + GlobalIndexResult exactMatchResult = reader.visit(appleVector, 1); assertThat(containsRowId(exactMatchResult, appleId)) .as("Exact match query should find Apple") .isTrue(); // Test 3: Query with vector similar to "Car" - should find Car first float[] carQueryVector = new float[] {0.15f, 0.85f}; - GlobalIndexResult carSearchResult = reader.search(carQueryVector, 1); + GlobalIndexResult carSearchResult = reader.visit(carQueryVector, 1); List carResultIds = new ArrayList<>(); for (Long id : carSearchResult) { @@ -582,12 +582,12 @@ public void testVectorIndexEndToEnd() throws IOException { assertThat(carResultIds).as("Query similar to Car should find Car").contains(carId); // Test 4: Search with larger k than available documents - GlobalIndexResult largeKResult = reader.search(queryVector, 10); + GlobalIndexResult largeKResult = reader.visit(queryVector, 10); assertThat(largeKResult).isNotNull(); assertThat(largeKResult.iterator().hasNext()).isTrue(); // Test 5: Verify all three documents are indexed - GlobalIndexResult allDocsResult = reader.search(new float[] {0.5f, 0.5f}, 3); + GlobalIndexResult allDocsResult = reader.visit(new float[] {0.5f, 0.5f}, 3); List allResultIds = new ArrayList<>(); for (Long id : allDocsResult) { allResultIds.add(id); From 08477e1ceec21688d293490a0aadba291d8b308c Mon Sep 17 00:00:00 2001 From: yantian Date: Wed, 3 Dec 2025 09:18:58 +0800 Subject: [PATCH 08/11] fix --- .../paimon/globalindex/GlobalIndexReader.java | 4 - paimon-vector/pom.xml | 2 +- .../apache/paimon/vector/ByteVectorKey.java | 48 ++++++ .../apache/paimon/vector/FloatVectorKey.java | 49 ++++++ .../vector/VectorGlobalIndexReader.java | 57 +++--- .../vector/VectorGlobalIndexWriter.java | 90 ++++------ .../org/apache/paimon/vector/VectorKey.java | 29 ++++ .../paimon/vector/VectorGlobalIndexTest.java | 163 ++++++++++++++---- 8 files changed, 326 insertions(+), 116 deletions(-) create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorKey.java create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorKey.java create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/VectorKey.java diff --git a/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java b/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java index 09ecbb2821f3..2ab40805f84f 100644 --- a/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java +++ b/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java @@ -41,8 +41,4 @@ default GlobalIndexResult visitOr(List children) { default GlobalIndexResult visit(TransformPredicate predicate) { throw new UnsupportedOperationException("Should not invoke this"); } - - default GlobalIndexResult visit(float[] query, int k) { - throw new UnsupportedOperationException("Should not invoke this"); - } } diff --git a/paimon-vector/pom.xml b/paimon-vector/pom.xml index 274d876ffa4a..1268b63c6a4f 100644 --- a/paimon-vector/pom.xml +++ b/paimon-vector/pom.xml @@ -32,7 +32,7 @@ under the License. Paimon : Vector - 9.11.1 + 9.12.3 11 diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorKey.java b/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorKey.java new file mode 100644 index 000000000000..06bb5c9f784f --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorKey.java @@ -0,0 +1,48 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.lucene.document.KnnByteVectorField; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.VectorSimilarityFunction; + +/** Vector key for byte vector. */ +public class ByteVectorKey implements VectorKey { + private final long rowId; + private final byte[] vector; + + public ByteVectorKey(long rowId, byte[] vector) { + this.rowId = rowId; + this.vector = vector; + } + + public long rowId() { + return rowId; + } + + @Override + public IndexableField toIndexableField( + String fieldName, VectorSimilarityFunction similarityFunction) { + return new KnnByteVectorField(fieldName, this.vector, similarityFunction); + } + + public byte[] vector() { + return vector; + } +} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorKey.java b/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorKey.java new file mode 100644 index 000000000000..1ec4baf75242 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorKey.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.VectorSimilarityFunction; + +/** Vector key for float vector. */ +public class FloatVectorKey implements VectorKey { + private final long rowId; + private final float[] vector; + + public FloatVectorKey(long rowId, float[] vector) { + this.rowId = rowId; + this.vector = vector; + } + + public float[] vector() { + return vector; + } + + @Override + public long rowId() { + return rowId; + } + + @Override + public IndexableField toIndexableField( + String fieldName, VectorSimilarityFunction similarityFunction) { + return new KnnFloatVectorField(fieldName, this.vector(), similarityFunction); + } +} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java index 91a0216b4660..46bfd67e730a 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -30,7 +30,9 @@ import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.StoredFields; import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnByteVectorQuery; import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; import org.apache.lucene.store.IndexOutput; @@ -77,32 +79,14 @@ public VectorGlobalIndexReader(GlobalIndexFileReader fileReader, List resultIds = new HashSet<>(); - - for (IndexSearcher searcher : searchers) { - try { - // Create KNN query - KnnFloatVectorQuery knnQuery = new KnnFloatVectorQuery(VECTOR_FIELD, query, k); - - // Execute search - TopDocs topDocs = searcher.search(knnQuery, k); - StoredFields storedFields = searcher.storedFields(); - Set fieldsToLoad = Set.of(ROW_ID_FIELD); - // Collect row IDs from results - for (org.apache.lucene.search.ScoreDoc scoreDoc : topDocs.scoreDocs) { - float rawScore = scoreDoc.score; - Document doc = storedFields.document(scoreDoc.doc, fieldsToLoad); - long rowId = doc.getField(ROW_ID_FIELD).numericValue().longValue(); - resultIds.add(rowId); - } - } catch (IOException e) { - throw new RuntimeException("Failed to search vector index", e); - } - } + public GlobalIndexResult search(float[] query, int k) { + KnnFloatVectorQuery knnQuery = new KnnFloatVectorQuery(VECTOR_FIELD, query, k); + return search(knnQuery, k); + } - return GlobalIndexResult.wrap(resultIds); + public GlobalIndexResult search(byte[] query, int k) { + KnnByteVectorQuery knnQuery = new KnnByteVectorQuery(VECTOR_FIELD, query, k); + return search(knnQuery, k); } @Override @@ -126,6 +110,29 @@ public void close() throws IOException { tempDirs.clear(); } + private GlobalIndexResult search(Query query, int k) { + Set resultIds = new HashSet<>(); + for (IndexSearcher searcher : searchers) { + try { + // Execute search + TopDocs topDocs = searcher.search(query, k); + StoredFields storedFields = searcher.storedFields(); + Set fieldsToLoad = Set.of(ROW_ID_FIELD); + // Collect row IDs from results + for (org.apache.lucene.search.ScoreDoc scoreDoc : topDocs.scoreDocs) { + float rawScore = scoreDoc.score; + Document doc = storedFields.document(scoreDoc.doc, fieldsToLoad); + long rowId = doc.getField(ROW_ID_FIELD).numericValue().longValue(); + resultIds.add(rowId); + } + } catch (IOException e) { + throw new RuntimeException("Failed to search vector index", e); + } + } + + return GlobalIndexResult.wrap(resultIds); + } + private void loadIndices(GlobalIndexFileReader fileReader, List files) throws IOException { for (GlobalIndexIOMeta meta : files) { diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java index 61762228bd32..71a04ac6b577 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -23,14 +23,12 @@ import org.apache.paimon.options.Options; import org.apache.paimon.types.ArrayType; import org.apache.paimon.types.DataType; -import org.apache.paimon.utils.Preconditions; import org.apache.paimon.utils.Range; import org.apache.lucene.codecs.KnnVectorsFormat; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat; import org.apache.lucene.document.Document; -import org.apache.lucene.document.KnnFloatVectorField; import org.apache.lucene.document.StoredField; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; @@ -63,7 +61,7 @@ public class VectorGlobalIndexWriter implements GlobalIndexWriter { private final VectorIndexOptions vectorOptions; private final VectorSimilarityFunction similarityFunction; - private final List vectors; + private final List vectors; private long minRowId = Long.MAX_VALUE; private long maxRowId = Long.MIN_VALUE; @@ -83,21 +81,36 @@ public VectorGlobalIndexWriter( @Override public void write(Object key) { - Preconditions.checkArgument( - key instanceof VectorKey, "Key must be VectorKey, but was: " + key.getClass()); - - VectorKey vectorKey = (VectorKey) key; - long rowId = vectorKey.getRowId(); - float[] vector = vectorKey.getVector(); - - checkArgument( - vector.length == vectorOptions.dimension(), - "Vector dimension mismatch: expected " - + vectorOptions.dimension() - + ", but got " - + vector.length); - - vectors.add(new VectorWithRowId(rowId, vector)); + Long rowId; + if (key instanceof FloatVectorKey) { + FloatVectorKey vectorKey = (FloatVectorKey) key; + rowId = vectorKey.rowId(); + float[] vector = vectorKey.vector(); + + checkArgument( + vector.length == vectorOptions.dimension(), + "Vector dimension mismatch: expected " + + vectorOptions.dimension() + + ", but got " + + vector.length); + + vectors.add(vectorKey); + } else if (key instanceof ByteVectorKey) { + ByteVectorKey vectorKey = (ByteVectorKey) key; + rowId = vectorKey.rowId(); + byte[] byteVector = vectorKey.vector(); + + checkArgument( + byteVector.length == vectorOptions.dimension(), + "Vector dimension mismatch: expected " + + vectorOptions.dimension() + + ", but got " + + byteVector.length); + + vectors.add(vectorKey); + } else { + throw new IllegalArgumentException("Unsupported key type: " + key.getClass().getName()); + } minRowId = Math.min(minRowId, rowId); maxRowId = Math.max(maxRowId, rowId); } @@ -169,16 +182,14 @@ private byte[] buildLuceneIndex(int m, int efConstruction) throws IOException { try (IndexWriter writer = new IndexWriter(directory, config)) { // Add each vector as a document - for (VectorWithRowId vectorWithRowId : vectors) { + for (VectorKey vectorKey : vectors) { Document doc = new Document(); // Add KNN vector field - doc.add( - new KnnFloatVectorField( - VECTOR_FIELD, vectorWithRowId.vector, similarityFunction)); + doc.add(vectorKey.toIndexableField(VECTOR_FIELD, similarityFunction)); // Store row ID - doc.add(new StoredField(ROW_ID_FIELD, vectorWithRowId.rowId)); + doc.add(new StoredField(ROW_ID_FIELD, vectorKey.rowId())); writer.addDocument(doc); } @@ -200,7 +211,7 @@ private static IndexWriterConfig getIndexWriterConfig(int m, int efConstruction) IndexWriterConfig config = new IndexWriterConfig(); config.setRAMBufferSizeMB(256); // Increase buffer for better performance config.setCodec( - new Lucene99Codec(Lucene99Codec.Mode.BEST_SPEED) { + new Lucene912Codec(Lucene912Codec.Mode.BEST_SPEED) { @Override public KnnVectorsFormat getKnnVectorsFormatForField(String field) { return new Lucene99HnswScalarQuantizedVectorsFormat(m, efConstruction); @@ -277,33 +288,4 @@ private byte[] longToBytes(long value) { (byte) value }; } - - /** Key type for vector indexing. */ - public static class VectorKey { - private final long rowId; - private final float[] vector; - - public VectorKey(long rowId, float[] vector) { - this.rowId = rowId; - this.vector = vector; - } - - public long getRowId() { - return rowId; - } - - public float[] getVector() { - return vector; - } - } - - private static class VectorWithRowId { - final long rowId; - final float[] vector; - - VectorWithRowId(long rowId, float[] vector) { - this.rowId = rowId; - this.vector = vector; - } - } } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorKey.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorKey.java new file mode 100644 index 000000000000..7adff60845e4 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorKey.java @@ -0,0 +1,29 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.VectorSimilarityFunction; + +/** Vector key interface. */ +public interface VectorKey { + long rowId(); + + IndexableField toIndexableField(String fieldName, VectorSimilarityFunction similarityFunction); +} diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java index ba8b2668be5b..5d3e8ca1b2ac 100644 --- a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java @@ -30,6 +30,7 @@ import org.apache.paimon.types.ArrayType; import org.apache.paimon.types.DataType; import org.apache.paimon.types.FloatType; +import org.apache.paimon.types.TinyIntType; import org.apache.paimon.utils.Range; import org.junit.jupiter.api.AfterEach; @@ -101,7 +102,7 @@ public void testWriteAndReadVectorIndex() throws IOException { List testVectors = generateRandomVectors(numVectors, dimension); for (int i = 0; i < numVectors; i++) { - writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + writer.write(new FloatVectorKey(i, testVectors.get(i))); } List results = writer.finish(); @@ -124,7 +125,7 @@ public void testWriteAndReadVectorIndex() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Test search - query with first vector should return itself - GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 5); + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); assertThat(searchResult).isNotNull(); reader.close(); @@ -148,13 +149,13 @@ public void testSearchSimilarVectors() throws IOException { new VectorGlobalIndexWriter(fileWriter, vectorType, options); // Write base vector and similar variants - writer.write(new VectorGlobalIndexWriter.VectorKey(0, baseVector)); + writer.write(new FloatVectorKey(0, baseVector)); for (int i = 1; i < numVectors; i++) { float[] variant = baseVector.clone(); // Add small noise variant[i % dimension] += (i % 2 == 0 ? 0.1f : -0.1f); - writer.write(new VectorGlobalIndexWriter.VectorKey(i, variant)); + writer.write(new FloatVectorKey(i, variant)); } List results = writer.finish(); @@ -173,7 +174,7 @@ public void testSearchSimilarVectors() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Search with base vector should return similar vectors - GlobalIndexResult searchResult = reader.visit(baseVector, 10); + GlobalIndexResult searchResult = reader.search(baseVector, 10); assertThat(searchResult).isNotNull(); reader.close(); @@ -198,7 +199,7 @@ public void testDifferentSimilarityFunctions() throws IOException { List testVectors = generateRandomVectors(numVectors, dimension); for (int i = 0; i < numVectors; i++) { - writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + writer.write(new FloatVectorKey(i, testVectors.get(i))); } List results = writer.finish(); @@ -217,7 +218,7 @@ public void testDifferentSimilarityFunctions() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Verify search works with this metric - GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 3); + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 3); assertThat(searchResult).isNotNull(); reader.close(); @@ -240,7 +241,7 @@ public void testDifferentDimensions() throws IOException { int numVectors = 10; List testVectors = generateRandomVectors(numVectors, dimension); for (int i = 0; i < numVectors; i++) { - writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + writer.write(new FloatVectorKey(i, testVectors.get(i))); } List results = writer.finish(); @@ -259,7 +260,7 @@ public void testDifferentDimensions() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Verify search works with this dimension - GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 5); + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); assertThat(searchResult).isNotNull(); reader.close(); @@ -277,10 +278,7 @@ public void testDimensionMismatch() throws IOException { // Try to write vector with wrong dimension float[] wrongDimVector = new float[32]; // Wrong dimension - assertThatThrownBy( - () -> - writer.write( - new VectorGlobalIndexWriter.VectorKey(0, wrongDimVector))) + assertThatThrownBy(() -> writer.write(new FloatVectorKey(0, wrongDimVector))) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("dimension mismatch"); } @@ -313,7 +311,7 @@ public void testHNSWParameters() throws IOException { List testVectors = generateRandomVectors(numVectors, dimension); for (int i = 0; i < numVectors; i++) { - writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + writer.write(new FloatVectorKey(i, testVectors.get(i))); } List results = writer.finish(); @@ -332,7 +330,7 @@ public void testHNSWParameters() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Verify search works with custom HNSW parameters - GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 5); + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); assertThat(searchResult).isNotNull(); reader.close(); @@ -351,7 +349,7 @@ public void testLargeK() throws IOException { List testVectors = generateRandomVectors(numVectors, dimension); for (int i = 0; i < numVectors; i++) { - writer.write(new VectorGlobalIndexWriter.VectorKey(i, testVectors.get(i))); + writer.write(new FloatVectorKey(i, testVectors.get(i))); } List results = writer.finish(); @@ -370,7 +368,7 @@ public void testLargeK() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); // Search with k larger than number of vectors - GlobalIndexResult searchResult = reader.visit(testVectors.get(0), 200); + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 200); assertThat(searchResult).isNotNull(); assertThat(searchResult).isNotNull(); @@ -406,7 +404,7 @@ public void testEndToEndMultiShardScenario() throws IOException { // Write vectors for this shard for (int i = startIdx; i < endIdx; i++) { - writer.write(new VectorGlobalIndexWriter.VectorKey(i, allTestVectors.get(i))); + writer.write(new FloatVectorKey(i, allTestVectors.get(i))); } List results = writer.finish(); @@ -444,7 +442,7 @@ public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) // Test 1: Search with various k values int[] kValues = {1, 5, 10, 20, 50}; for (int k : kValues) { - GlobalIndexResult searchResult = reader.visit(allTestVectors.get(0), k); + GlobalIndexResult searchResult = reader.search(allTestVectors.get(0), k); assertThat(searchResult).isNotNull(); assertThat(containsRowId(searchResult, 0)).isTrue(); } @@ -452,7 +450,7 @@ public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) // Test 2: Search with query vectors from different shards for (int shardIdx = 0; shardIdx < shardsCount; shardIdx++) { int queryIdx = shardIdx * vectorsPerShard + 5; - GlobalIndexResult searchResult = reader.visit(allTestVectors.get(queryIdx), 10); + GlobalIndexResult searchResult = reader.search(allTestVectors.get(queryIdx), 10); assertThat(searchResult).isNotNull(); assertThat(containsRowId(searchResult, queryIdx)) .as("Search result should contain the query vector's own ID") @@ -468,7 +466,7 @@ public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) } normalize(slightlyModifiedVector); - GlobalIndexResult similarSearchResult = reader.visit(slightlyModifiedVector, 5); + GlobalIndexResult similarSearchResult = reader.search(slightlyModifiedVector, 5); assertThat(similarSearchResult).isNotNull(); assertThat(containsRowId(similarSearchResult, 50)) .as("Search with very similar vector should find the original") @@ -483,18 +481,18 @@ public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) } normalize(clusterCenter); - GlobalIndexResult clusterSearchResult = reader.visit(clusterCenter, 20); + GlobalIndexResult clusterSearchResult = reader.search(clusterCenter, 20); assertThat(clusterSearchResult).isNotNull(); assertThat(clusterSearchResult.iterator().hasNext()).isTrue(); // Test 5: Edge case - search with k larger than total vectors - GlobalIndexResult largeKResult = reader.visit(allTestVectors.get(10), totalVectors * 2); + GlobalIndexResult largeKResult = reader.search(allTestVectors.get(10), totalVectors * 2); assertThat(largeKResult).isNotNull(); // Test 6: Multiple consecutive searches (stress test) for (int i = 0; i < 20; i++) { int randomIdx = random.nextInt(totalVectors); - GlobalIndexResult result = reader.visit(allTestVectors.get(randomIdx), 15); + GlobalIndexResult result = reader.search(allTestVectors.get(randomIdx), 15); assertThat(result).isNotNull(); } @@ -526,9 +524,9 @@ public void testVectorIndexEndToEnd() throws IOException { long bananaId = 1L; long carId = 2L; - writer.write(new VectorGlobalIndexWriter.VectorKey(appleId, appleVector)); - writer.write(new VectorGlobalIndexWriter.VectorKey(bananaId, bananaVector)); - writer.write(new VectorGlobalIndexWriter.VectorKey(carId, carVector)); + writer.write(new FloatVectorKey(appleId, appleVector)); + writer.write(new FloatVectorKey(bananaId, bananaVector)); + writer.write(new FloatVectorKey(carId, carVector)); List results = writer.finish(); assertThat(results).hasSize(1); @@ -548,7 +546,7 @@ public void testVectorIndexEndToEnd() throws IOException { // Test 1: Query with vector similar to "Apple" - should find Apple and Banana // Query: [0.85, 0.15] is between Apple and Banana float[] queryVector = new float[] {0.85f, 0.15f}; - GlobalIndexResult searchResult = reader.visit(queryVector, 2); + GlobalIndexResult searchResult = reader.search(queryVector, 2); assertThat(searchResult).isNotNull(); assertThat(searchResult.iterator().hasNext()).as("Search should return results").isTrue(); @@ -565,14 +563,14 @@ public void testVectorIndexEndToEnd() throws IOException { assertThat(resultIds.size()).as("Should return top 2 results").isLessThanOrEqualTo(2); // Test 2: Query with exact match - should find exact document - GlobalIndexResult exactMatchResult = reader.visit(appleVector, 1); + GlobalIndexResult exactMatchResult = reader.search(appleVector, 1); assertThat(containsRowId(exactMatchResult, appleId)) .as("Exact match query should find Apple") .isTrue(); // Test 3: Query with vector similar to "Car" - should find Car first float[] carQueryVector = new float[] {0.15f, 0.85f}; - GlobalIndexResult carSearchResult = reader.visit(carQueryVector, 1); + GlobalIndexResult carSearchResult = reader.search(carQueryVector, 1); List carResultIds = new ArrayList<>(); for (Long id : carSearchResult) { @@ -582,12 +580,12 @@ public void testVectorIndexEndToEnd() throws IOException { assertThat(carResultIds).as("Query similar to Car should find Car").contains(carId); // Test 4: Search with larger k than available documents - GlobalIndexResult largeKResult = reader.visit(queryVector, 10); + GlobalIndexResult largeKResult = reader.search(queryVector, 10); assertThat(largeKResult).isNotNull(); assertThat(largeKResult.iterator().hasNext()).isTrue(); // Test 5: Verify all three documents are indexed - GlobalIndexResult allDocsResult = reader.visit(new float[] {0.5f, 0.5f}, 3); + GlobalIndexResult allDocsResult = reader.search(new float[] {0.5f, 0.5f}, 3); List allResultIds = new ArrayList<>(); for (Long id : allDocsResult) { allResultIds.add(id); @@ -646,4 +644,105 @@ private boolean containsRowId(GlobalIndexResult result, long rowId) { } return false; } + + @Test + public void testByteVectorIndexEndToEnd() throws IOException { + Options options = createDefaultOptions(); + int dimension = 2; + options.setInteger("vector.dim", dimension); + options.setString("vector.metric", "EUCLIDEAN"); + options.setInteger("vector.m", 16); + options.setInteger("vector.ef-construction", 100); + + // Create semantic vectors representing "documents" using byte vectors + // Apple [90, 10] and Banana [80, 20] are similar (fruits) + // Car [10, 90] is different (vehicle) + byte[] appleVector = new byte[] {90, 10}; + byte[] bananaVector = new byte[] {80, 20}; + byte[] carVector = new byte[] {10, 90}; + + DataType byteVectorType = new ArrayType(new TinyIntType()); + GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); + VectorGlobalIndexWriter writer = + new VectorGlobalIndexWriter(fileWriter, byteVectorType, options); + + // Write vectors with row IDs representing documents + long appleId = 0L; + long bananaId = 1L; + long carId = 2L; + + writer.write(new ByteVectorKey(appleId, appleVector)); + writer.write(new ByteVectorKey(bananaId, bananaVector)); + writer.write(new ByteVectorKey(carId, carVector)); + + List results = writer.finish(); + assertThat(results).hasSize(1); + + GlobalIndexWriter.ResultEntry result = results.get(0); + GlobalIndexFileReader fileReader = createFileReader(indexPath); + List metas = new ArrayList<>(); + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(indexPath, result.fileName())), + result.rowRange(), + result.meta())); + + VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); + + // Test 1: Query with vector similar to "Apple" - should find Apple and Banana + // Query: [85, 15] is between Apple and Banana + byte[] queryVector = new byte[] {85, 15}; + GlobalIndexResult searchResult = reader.search(queryVector, 2); + + assertThat(searchResult).isNotNull(); + assertThat(searchResult.iterator().hasNext()).as("Search should return results").isTrue(); + + // Collect all result IDs + List resultIds = new ArrayList<>(); + for (Long id : searchResult) { + resultIds.add(id); + } + + // Should find Apple and Banana (similar fruits), not Car + assertThat(resultIds).as("Should find Apple (most similar)").contains(appleId); + assertThat(resultIds).as("Should find Banana (second most similar)").contains(bananaId); + assertThat(resultIds.size()).as("Should return top 2 results").isLessThanOrEqualTo(2); + + // Test 2: Query with exact match - should find exact document + GlobalIndexResult exactMatchResult = reader.search(appleVector, 1); + assertThat(containsRowId(exactMatchResult, appleId)) + .as("Exact match query should find Apple") + .isTrue(); + + // Test 3: Query with vector similar to "Car" - should find Car first + byte[] carQueryVector = new byte[] {15, 85}; + GlobalIndexResult carSearchResult = reader.search(carQueryVector, 1); + + List carResultIds = new ArrayList<>(); + for (Long id : carSearchResult) { + carResultIds.add(id); + } + + assertThat(carResultIds).as("Query similar to Car should find Car").contains(carId); + + // Test 4: Search with larger k than available documents + GlobalIndexResult largeKResult = reader.search(queryVector, 10); + assertThat(largeKResult).isNotNull(); + assertThat(largeKResult.iterator().hasNext()).isTrue(); + + // Test 5: Verify all three documents are indexed + GlobalIndexResult allDocsResult = reader.search(new byte[] {50, 50}, 3); + List allResultIds = new ArrayList<>(); + for (Long id : allDocsResult) { + allResultIds.add(id); + } + + assertThat(allResultIds.size()) + .as("Should be able to retrieve all 3 documents") + .isGreaterThanOrEqualTo(1) + .isLessThanOrEqualTo(3); + + reader.close(); + } } From fec2404b9a1c547ea8a6c8c9acb329d4d0868085 Mon Sep 17 00:00:00 2001 From: yantian Date: Wed, 3 Dec 2025 18:00:15 +0800 Subject: [PATCH 09/11] add test script, split, notice --- .github/workflows/utitcase-paimon-vector.yml | 65 ++ paimon-vector/pom.xml | 49 ++ ...yteVectorKey.java => ByteVectorIndex.java} | 12 +- ...atVectorKey.java => FloatVectorIndex.java} | 11 +- .../paimon/vector/IndexMMapDirectory.java | 57 ++ .../vector/VectorGlobalIndexReader.java | 125 ++-- .../vector/VectorGlobalIndexWriter.java | 147 +++-- .../{VectorKey.java => VectorIndex.java} | 17 +- .../paimon/vector/VectorIndexOptions.java | 29 +- .../src/main/resources/META-INF/NOTICE | 9 + .../paimon/vector/VectorGlobalIndexTest.java | 561 +++--------------- 11 files changed, 431 insertions(+), 651 deletions(-) create mode 100644 .github/workflows/utitcase-paimon-vector.yml rename paimon-vector/src/main/java/org/apache/paimon/vector/{ByteVectorKey.java => ByteVectorIndex.java} (78%) rename paimon-vector/src/main/java/org/apache/paimon/vector/{FloatVectorKey.java => FloatVectorIndex.java} (78%) create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/IndexMMapDirectory.java rename paimon-vector/src/main/java/org/apache/paimon/vector/{VectorKey.java => VectorIndex.java} (67%) create mode 100644 paimon-vector/src/main/resources/META-INF/NOTICE diff --git a/.github/workflows/utitcase-paimon-vector.yml b/.github/workflows/utitcase-paimon-vector.yml new file mode 100644 index 000000000000..1699efb3e12c --- /dev/null +++ b/.github/workflows/utitcase-paimon-vector.yml @@ -0,0 +1,65 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################ + +name: UTCase and ITCase Paimon Vector + +on: + push: + pull_request: + paths-ignore: + - 'docs/**' + - '**/*.md' + - 'paimon-python/**' + - '.github/workflows/paimon-python-checks.yml' + +env: + JDK_VERSION: 11 + MAVEN_OPTS: -Dmaven.wagon.httpconnectionManager.ttlSeconds=30 -Dmaven.wagon.http.retryHandler.requestSentEnabled=true + +concurrency: + group: ${{ github.workflow }}-${{ github.event_name }}-${{ github.event.number || github.run_id }} + cancel-in-progress: true + +jobs: + build_test: + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up JDK ${{ env.JDK_VERSION }} + uses: actions/setup-java@v4 + with: + java-version: ${{ env.JDK_VERSION }} + distribution: 'temurin' + + - name: Build Paimon Vector + run: mvn -T 2C -B clean install -DskipTests -Ppaimon-vector + + - name: Test Paimon Vector + run: | + . .github/workflows/utils.sh + jvm_timezone=$(random_timezone) + echo "JVM timezone is set to $jvm_timezone" + TEST_MODULE="org.apache.paimon:paimon-vector" + echo "Start testing module: $TEST_MODULE" + mvn -T 2C -B test verify -Ppaimon-vector -pl "${TEST_MODULE}" -Duser.timezone=$jvm_timezone + echo "All modules tested" + env: + MAVEN_OPTS: -Xmx4096m -XX:+UseG1GC -XX:CICompilerCount=2 \ No newline at end of file diff --git a/paimon-vector/pom.xml b/paimon-vector/pom.xml index 1268b63c6a4f..11c4d96afc76 100644 --- a/paimon-vector/pom.xml +++ b/paimon-vector/pom.xml @@ -72,4 +72,53 @@ under the License. test + + + + + maven-jar-plugin + + + + test-jar + + + + + + maven-shade-plugin + + + shade-paimon + package + + shade + + + + + org.apache.lucene:lucene-core + org.apache.lucene:lucene-codecs + + + + + org.apache.lucene + org.apache.paimon.shade.lucene + + + + + * + + META-INF/LICENSE.txt + + + + + + + + + diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorKey.java b/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorIndex.java similarity index 78% rename from paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorKey.java rename to paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorIndex.java index 06bb5c9f784f..33239b92389e 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorKey.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorIndex.java @@ -22,24 +22,24 @@ import org.apache.lucene.index.IndexableField; import org.apache.lucene.index.VectorSimilarityFunction; -/** Vector key for byte vector. */ -public class ByteVectorKey implements VectorKey { +/** Vector index for byte vector. */ +public class ByteVectorIndex extends VectorIndex { private final long rowId; private final byte[] vector; - public ByteVectorKey(long rowId, byte[] vector) { + public ByteVectorIndex(long rowId, byte[] vector) { this.rowId = rowId; this.vector = vector; } + @Override public long rowId() { return rowId; } @Override - public IndexableField toIndexableField( - String fieldName, VectorSimilarityFunction similarityFunction) { - return new KnnByteVectorField(fieldName, this.vector, similarityFunction); + public IndexableField indexableField(VectorSimilarityFunction similarityFunction) { + return new KnnByteVectorField(VECTOR_FIELD, this.vector, similarityFunction); } public byte[] vector() { diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorKey.java b/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorIndex.java similarity index 78% rename from paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorKey.java rename to paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorIndex.java index 1ec4baf75242..f2229b49e411 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorKey.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorIndex.java @@ -22,12 +22,12 @@ import org.apache.lucene.index.IndexableField; import org.apache.lucene.index.VectorSimilarityFunction; -/** Vector key for float vector. */ -public class FloatVectorKey implements VectorKey { +/** Vector index for float vector. */ +public class FloatVectorIndex extends VectorIndex { private final long rowId; private final float[] vector; - public FloatVectorKey(long rowId, float[] vector) { + public FloatVectorIndex(long rowId, float[] vector) { this.rowId = rowId; this.vector = vector; } @@ -42,8 +42,7 @@ public long rowId() { } @Override - public IndexableField toIndexableField( - String fieldName, VectorSimilarityFunction similarityFunction) { - return new KnnFloatVectorField(fieldName, this.vector(), similarityFunction); + public IndexableField indexableField(VectorSimilarityFunction similarityFunction) { + return new KnnFloatVectorField(VECTOR_FIELD, this.vector(), similarityFunction); } } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/IndexMMapDirectory.java b/paimon-vector/src/main/java/org/apache/paimon/vector/IndexMMapDirectory.java new file mode 100644 index 000000000000..4bfb235dca41 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/IndexMMapDirectory.java @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.lucene.store.MMapDirectory; + +import java.io.IOException; +import java.nio.file.Path; +import java.util.UUID; + +/** A wrapper of MMapDirectory for vector index. */ +public class IndexMMapDirectory implements AutoCloseable { + private final Path path; + private final MMapDirectory mmapDirectory; + + public IndexMMapDirectory() throws IOException { + this.path = + java.nio.file.Files.createTempDirectory("paimon-vector-index" + UUID.randomUUID()); + this.mmapDirectory = new MMapDirectory(path); + } + + public MMapDirectory directory() { + return mmapDirectory; + } + + public void close() throws Exception { + mmapDirectory.close(); + if (java.nio.file.Files.exists(path)) { + java.nio.file.Files.walk(path) + .sorted(java.util.Comparator.reverseOrder()) + .forEach( + p -> { + try { + java.nio.file.Files.delete(p); + } catch (IOException e) { + // Ignore cleanup errors + } + }); + } + } +} diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java index 46bfd67e730a..cb38d68c6081 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -24,6 +24,7 @@ import org.apache.paimon.globalindex.GlobalIndexResult; import org.apache.paimon.globalindex.io.GlobalIndexFileReader; import org.apache.paimon.predicate.FieldRef; +import org.apache.paimon.utils.RoaringNavigableMap64; import org.apache.lucene.document.Document; import org.apache.lucene.index.DirectoryReader; @@ -34,20 +35,14 @@ import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.TopDocs; -import org.apache.lucene.store.Directory; import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.store.MMapDirectory; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; -import java.nio.file.Files; -import java.nio.file.Path; import java.util.ArrayList; -import java.util.HashSet; import java.util.List; import java.util.Set; -import java.util.UUID; /** * Vector global index reader using Apache Lucene 9.x. @@ -57,18 +52,15 @@ */ public class VectorGlobalIndexReader implements GlobalIndexReader { - private static final String VECTOR_FIELD = "vector"; - private static final String ROW_ID_FIELD = "rowId"; + private static final int BUFFER_SIZE = 8192; // 8KB buffer for streaming private final List searchers; - private final List directories; - private final List tempDirs; + private final List directories; public VectorGlobalIndexReader(GlobalIndexFileReader fileReader, List files) throws IOException { this.searchers = new ArrayList<>(); this.directories = new ArrayList<>(); - this.tempDirs = new ArrayList<>(); loadIndices(fileReader, files); } @@ -80,12 +72,12 @@ public VectorGlobalIndexReader(GlobalIndexFileReader fileReader, List resultIds = new HashSet<>(); + RoaringNavigableMap64 roaringBitmap64 = new RoaringNavigableMap64(); for (IndexSearcher searcher : searchers) { try { // Execute search TopDocs topDocs = searcher.search(query, k); StoredFields storedFields = searcher.storedFields(); - Set fieldsToLoad = Set.of(ROW_ID_FIELD); + Set fieldsToLoad = Set.of(VectorIndex.ROW_ID_FIELD); // Collect row IDs from results for (org.apache.lucene.search.ScoreDoc scoreDoc : topDocs.scoreDocs) { float rawScore = scoreDoc.score; Document doc = storedFields.document(scoreDoc.doc, fieldsToLoad); - long rowId = doc.getField(ROW_ID_FIELD).numericValue().longValue(); - resultIds.add(rowId); + long rowId = doc.getField(VectorIndex.ROW_ID_FIELD).numericValue().longValue(); + roaringBitmap64.add(rowId); } } catch (IOException e) { throw new RuntimeException("Failed to search vector index", e); } } - return GlobalIndexResult.wrap(resultIds); + return GlobalIndexResult.create(() -> roaringBitmap64); } private void loadIndices(GlobalIndexFileReader fileReader, List files) throws IOException { for (GlobalIndexIOMeta meta : files) { try (SeekableInputStream in = fileReader.getInputStream(meta.fileName())) { - byte[] indexBytes = new byte[(int) meta.fileSize()]; - int totalRead = 0; - while (totalRead < indexBytes.length) { - int read = in.read(indexBytes, totalRead, indexBytes.length - totalRead); - if (read == -1) { - throw new IOException("Unexpected end of stream"); - } - totalRead += read; - } - - Directory directory = deserializeDirectory(indexBytes); + IndexMMapDirectory directory = deserializeDirectory(in); directories.add(directory); - - IndexReader reader = DirectoryReader.open(directory); + IndexReader reader = DirectoryReader.open(directory.directory()); IndexSearcher searcher = new IndexSearcher(reader); searchers.add(searcher); } } } - private Directory deserializeDirectory(byte[] data) throws IOException { - // Create temporary directory for MMap - Path tempDir = Files.createTempDirectory("paimon-vector-read" + UUID.randomUUID()); - tempDirs.add(tempDir); - Directory directory = new MMapDirectory(tempDir); - - ByteBuffer buffer = ByteBuffer.wrap(data); + private IndexMMapDirectory deserializeDirectory(SeekableInputStream in) throws IOException { + IndexMMapDirectory indexMMapDirectory = new IndexMMapDirectory(); // Read number of files - int numFiles = buffer.getInt(); + int numFiles = readInt(in); + + // Reusable buffer for streaming + byte[] buffer = new byte[BUFFER_SIZE]; for (int i = 0; i < numFiles; i++) { // Read file name - int nameLength = buffer.getInt(); + int nameLength = readInt(in); byte[] nameBytes = new byte[nameLength]; - buffer.get(nameBytes); + readFully(in, nameBytes); String fileName = new String(nameBytes, StandardCharsets.UTF_8); - // Read file content - long fileLength = buffer.getLong(); - byte[] fileContent = new byte[(int) fileLength]; - buffer.get(fileContent); - - // Write to directory - try (IndexOutput output = directory.createOutput(fileName, null)) { - output.writeBytes(fileContent, 0, fileContent.length); + // Read file content length + long fileLength = readLong(in); + + // Stream file content directly to directory + try (IndexOutput output = indexMMapDirectory.directory().createOutput(fileName, null)) { + long remaining = fileLength; + while (remaining > 0) { + int toRead = (int) Math.min(buffer.length, remaining); + readFully(in, buffer, 0, toRead); + output.writeBytes(buffer, 0, toRead); + remaining -= toRead; + } } } - return directory; + return indexMMapDirectory; + } + + private int readInt(SeekableInputStream in) throws IOException { + byte[] bytes = new byte[4]; + readFully(in, bytes); + return ByteBuffer.wrap(bytes).getInt(); + } + + private long readLong(SeekableInputStream in) throws IOException { + byte[] bytes = new byte[8]; + readFully(in, bytes); + return ByteBuffer.wrap(bytes).getLong(); + } + + private void readFully(SeekableInputStream in, byte[] buffer) throws IOException { + readFully(in, buffer, 0, buffer.length); + } + + private void readFully(SeekableInputStream in, byte[] buffer, int offset, int length) + throws IOException { + int totalRead = 0; + while (totalRead < length) { + int read = in.read(buffer, offset + totalRead, length - totalRead); + if (read == -1) { + throw new IOException("Unexpected end of stream"); + } + totalRead += read; + } } private void deleteDirectory(java.nio.file.Path path) throws IOException { diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java index 71a04ac6b577..3ed77b1bdb69 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -29,12 +29,10 @@ import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat; import org.apache.lucene.document.Document; -import org.apache.lucene.document.StoredField; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.store.Directory; -import org.apache.lucene.store.MMapDirectory; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -53,38 +51,31 @@ */ public class VectorGlobalIndexWriter implements GlobalIndexWriter { - private static final String VECTOR_FIELD = "vector"; - private static final String ROW_ID_FIELD = "rowId"; - private final GlobalIndexFileWriter fileWriter; private final DataType fieldType; private final VectorIndexOptions vectorOptions; private final VectorSimilarityFunction similarityFunction; + private final int sizePerIndex; - private final List vectors; - private long minRowId = Long.MAX_VALUE; - private long maxRowId = Long.MIN_VALUE; + private final List vectors; public VectorGlobalIndexWriter( GlobalIndexFileWriter fileWriter, DataType fieldType, Options options) { + checkArgument( + fieldType instanceof ArrayType, + "Vector field type must be ARRAY, but was: " + fieldType); this.fileWriter = fileWriter; this.fieldType = fieldType; this.vectors = new ArrayList<>(); - - // Parse options this.vectorOptions = new VectorIndexOptions(options); this.similarityFunction = parseMetricToLucene(vectorOptions.metric()); - - // Validate field type - validateFieldType(fieldType); + this.sizePerIndex = vectorOptions.sizePerIndex(); } @Override public void write(Object key) { - Long rowId; - if (key instanceof FloatVectorKey) { - FloatVectorKey vectorKey = (FloatVectorKey) key; - rowId = vectorKey.rowId(); + if (key instanceof FloatVectorIndex) { + FloatVectorIndex vectorKey = (FloatVectorIndex) key; float[] vector = vectorKey.vector(); checkArgument( @@ -95,9 +86,8 @@ public void write(Object key) { + vector.length); vectors.add(vectorKey); - } else if (key instanceof ByteVectorKey) { - ByteVectorKey vectorKey = (ByteVectorKey) key; - rowId = vectorKey.rowId(); + } else if (key instanceof ByteVectorIndex) { + ByteVectorIndex vectorKey = (ByteVectorIndex) key; byte[] byteVector = vectorKey.vector(); checkArgument( @@ -109,10 +99,9 @@ public void write(Object key) { vectors.add(vectorKey); } else { - throw new IllegalArgumentException("Unsupported key type: " + key.getClass().getName()); + throw new IllegalArgumentException( + "Unsupported index type: " + key.getClass().getName()); } - minRowId = Math.min(minRowId, rowId); - maxRowId = Math.max(maxRowId, rowId); } @Override @@ -122,38 +111,52 @@ public List finish() { return new ArrayList<>(); } - // Build Lucene index in memory - byte[] indexBytes = - buildLuceneIndex(this.vectorOptions.m(), this.vectorOptions.efConstruction()); - - // Create metadata - VectorIndexMetadata metadata = - new VectorIndexMetadata( - vectorOptions.dimension(), - vectorOptions.metric(), - vectorOptions.m(), - vectorOptions.efConstruction()); - byte[] metaBytes = VectorIndexMetadata.serializeMetadata(metadata); - - // Write to file - String fileName = fileWriter.newFileName(VectorGlobalIndexerFactory.IDENTIFIER); - try (OutputStream out = fileWriter.newOutputStream(fileName)) { - out.write(indexBytes); + List results = new ArrayList<>(); + + // Split vectors into batches if size exceeds sizePerIndex + int totalVectors = vectors.size(); + int numBatches = (int) Math.ceil((double) totalVectors / sizePerIndex); + + for (int batchIndex = 0; batchIndex < numBatches; batchIndex++) { + int startIdx = batchIndex * sizePerIndex; + int endIdx = Math.min(startIdx + sizePerIndex, totalVectors); + List batchVectors = vectors.subList(startIdx, endIdx); + + // Build index + byte[] indexBytes = + buildIndex( + batchVectors, + this.vectorOptions.m(), + this.vectorOptions.efConstruction(), + this.vectorOptions.writeBufferSize()); + + // Create metadata + VectorIndexMetadata metadata = + new VectorIndexMetadata( + vectorOptions.dimension(), + vectorOptions.metric(), + vectorOptions.m(), + vectorOptions.efConstruction()); + byte[] metaBytes = VectorIndexMetadata.serializeMetadata(metadata); + + // Write to file + String fileName = fileWriter.newFileName(VectorGlobalIndexerFactory.IDENTIFIER); + try (OutputStream out = fileWriter.newOutputStream(fileName)) { + out.write(indexBytes); + } + long minRowIdInBatch = batchVectors.get(0).rowId(); + long maxRowIdInBatch = batchVectors.get(batchVectors.size() - 1).rowId(); + results.add( + ResultEntry.of( + fileName, metaBytes, new Range(minRowIdInBatch, maxRowIdInBatch))); } - List results = new ArrayList<>(); - results.add(ResultEntry.of(fileName, metaBytes, new Range(minRowId, maxRowId + 1))); return results; } catch (IOException e) { throw new RuntimeException("Failed to write vector global index", e); } } - private void validateFieldType(DataType type) { - checkArgument( - type instanceof ArrayType, "Vector field type must be ARRAY, but was: " + type); - } - private VectorSimilarityFunction parseMetricToLucene(String metric) { switch (metric.toUpperCase()) { case "COSINE": @@ -169,27 +172,24 @@ private VectorSimilarityFunction parseMetricToLucene(String metric) { } } - private byte[] buildLuceneIndex(int m, int efConstruction) throws IOException { - // Create temporary directory for MMap - java.nio.file.Path tempDir = - java.nio.file.Files.createTempDirectory( - fileWriter.newFileName("paimon-vector-index")); - Directory directory = new MMapDirectory(tempDir); + private byte[] buildIndex( + List batchVectors, int m, int efConstruction, int writeBufferSize) + throws IOException { - try { + try (IndexMMapDirectory indexMMapDirectory = new IndexMMapDirectory()) { // Configure index writer - IndexWriterConfig config = getIndexWriterConfig(m, efConstruction); + IndexWriterConfig config = getIndexWriterConfig(m, efConstruction, writeBufferSize); - try (IndexWriter writer = new IndexWriter(directory, config)) { + try (IndexWriter writer = new IndexWriter(indexMMapDirectory.directory(), config)) { // Add each vector as a document - for (VectorKey vectorKey : vectors) { + for (VectorIndex vectorIndex : batchVectors) { Document doc = new Document(); // Add KNN vector field - doc.add(vectorKey.toIndexableField(VECTOR_FIELD, similarityFunction)); + doc.add(vectorIndex.indexableField(similarityFunction)); // Store row ID - doc.add(new StoredField(ROW_ID_FIELD, vectorKey.rowId())); + doc.add(vectorIndex.rowIdStoredField()); writer.addDocument(doc); } @@ -199,17 +199,17 @@ private byte[] buildLuceneIndex(int m, int efConstruction) throws IOException { } // Serialize directory to byte array - return serializeDirectory(directory); - } finally { - // Clean up - directory.close(); - deleteDirectory(tempDir); + return serializeDirectory(indexMMapDirectory.directory()); + } catch (Exception e) { + throw new RuntimeException(e); } } - private static IndexWriterConfig getIndexWriterConfig(int m, int efConstruction) { + private static IndexWriterConfig getIndexWriterConfig( + int m, int efConstruction, int writeBufferSize) { IndexWriterConfig config = new IndexWriterConfig(); - config.setRAMBufferSizeMB(256); // Increase buffer for better performance + config.setRAMBufferSizeMB( + writeBufferSize); // Configure RAM buffer size based on user settings config.setCodec( new Lucene912Codec(Lucene912Codec.Mode.BEST_SPEED) { @Override @@ -220,21 +220,6 @@ public KnnVectorsFormat getKnnVectorsFormatForField(String field) { return config; } - private void deleteDirectory(java.nio.file.Path path) throws IOException { - if (java.nio.file.Files.exists(path)) { - java.nio.file.Files.walk(path) - .sorted(java.util.Comparator.reverseOrder()) - .forEach( - p -> { - try { - java.nio.file.Files.delete(p); - } catch (IOException e) { - // Ignore cleanup errors - } - }); - } - } - private byte[] serializeDirectory(Directory directory) throws IOException { ByteArrayOutputStream baos = new ByteArrayOutputStream(); diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorKey.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndex.java similarity index 67% rename from paimon-vector/src/main/java/org/apache/paimon/vector/VectorKey.java rename to paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndex.java index 7adff60845e4..744b65a9c546 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorKey.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndex.java @@ -18,12 +18,21 @@ package org.apache.paimon.vector; +import org.apache.lucene.document.StoredField; import org.apache.lucene.index.IndexableField; import org.apache.lucene.index.VectorSimilarityFunction; -/** Vector key interface. */ -public interface VectorKey { - long rowId(); +/** Vector index interface. */ +public abstract class VectorIndex { - IndexableField toIndexableField(String fieldName, VectorSimilarityFunction similarityFunction); + public static final String VECTOR_FIELD = "vector"; + public static final String ROW_ID_FIELD = "rowId"; + + public abstract long rowId(); + + public abstract IndexableField indexableField(VectorSimilarityFunction similarityFunction); + + public StoredField rowIdStoredField() { + return new StoredField(ROW_ID_FIELD, rowId()); + } } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java index 571c1afb0344..31883e928aca 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java @@ -36,7 +36,7 @@ public class VectorIndexOptions { .stringType() .defaultValue("EUCLIDEAN") .withDescription( - "The similarity metric for vector search (COSINE, DOT_PRODUCT, EUCLIDEAN, MAX_INNER_PRODUCT)"); + "The similarity metric for vector search (COSINE, DOT_PRODUCT, EUCLIDEAN, MAX_INNER_PRODUCT), and EUCLIDEAN is the default"); public static final ConfigOption VECTOR_M = ConfigOptions.key("vector.m") @@ -52,16 +52,35 @@ public class VectorIndexOptions { .withDescription( "The size of the dynamic candidate list during the index construction"); + public static final ConfigOption VECTOR_SIZE_PER_INDEX = + ConfigOptions.key("vector.size-per-index") + .intType() + .defaultValue(10000) + .withDescription("The size of vectors stored in each vector index file"); + + public static final ConfigOption VECTOR_WRITE_BUFFER_SIZE = + ConfigOptions.key("vector.write-buffer-size") + .intType() + .defaultValue(256) + .withDescription("write buffer size in MB for vector index"); + private final int dimension; private final String metric; private final int m; private final int efConstruction; + private final int sizePerIndex; + private final int writeBufferSize; public VectorIndexOptions(Options options) { this.dimension = options.get(VECTOR_DIM); this.metric = options.get(VECTOR_METRIC); this.m = options.get(VECTOR_M); this.efConstruction = options.get(VECTOR_EF_CONSTRUCTION); + this.sizePerIndex = + options.get(VECTOR_SIZE_PER_INDEX) > 0 + ? options.get(VECTOR_SIZE_PER_INDEX) + : VECTOR_SIZE_PER_INDEX.defaultValue(); + this.writeBufferSize = options.get(VECTOR_WRITE_BUFFER_SIZE); } public int dimension() { @@ -79,4 +98,12 @@ public int m() { public int efConstruction() { return efConstruction; } + + public int sizePerIndex() { + return sizePerIndex; + } + + public int writeBufferSize() { + return writeBufferSize; + } } diff --git a/paimon-vector/src/main/resources/META-INF/NOTICE b/paimon-vector/src/main/resources/META-INF/NOTICE new file mode 100644 index 000000000000..0eb327d920a5 --- /dev/null +++ b/paimon-vector/src/main/resources/META-INF/NOTICE @@ -0,0 +1,9 @@ +paimon-vector +Copyright 2023-2025 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +This project bundles the following dependencies under the Apache Software License 2.0 (http://www.apache.org/licenses/LICENSE-2.0.txt) +- org.apache.lucene:lucene-core:9.12.3 +- org.apache.lucene:lucene-codecs:9.12.3 \ No newline at end of file diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java index 5d3e8ca1b2ac..47b3dd5d703d 100644 --- a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java @@ -30,8 +30,6 @@ import org.apache.paimon.types.ArrayType; import org.apache.paimon.types.DataType; import org.apache.paimon.types.FloatType; -import org.apache.paimon.types.TinyIntType; -import org.apache.paimon.utils.Range; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -89,97 +87,6 @@ private GlobalIndexFileReader createFileReader(Path path) { return fileName -> fileIO.newInputStream(new Path(path, fileName)); } - @Test - public void testWriteAndReadVectorIndex() throws IOException { - Options options = createDefaultOptions(); - int dimension = 128; - int numVectors = 100; - - // Create writer and write vectors - GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); - VectorGlobalIndexWriter writer = - new VectorGlobalIndexWriter(fileWriter, vectorType, options); - - List testVectors = generateRandomVectors(numVectors, dimension); - for (int i = 0; i < numVectors; i++) { - writer.write(new FloatVectorKey(i, testVectors.get(i))); - } - - List results = writer.finish(); - assertThat(results).hasSize(1); - - GlobalIndexWriter.ResultEntry result = results.get(0); - assertThat(result.fileName()).isNotNull(); - assertThat(result.rowRange()).isEqualTo(new Range(0, numVectors)); - - // Create reader and verify - GlobalIndexFileReader fileReader = createFileReader(indexPath); - List metas = new ArrayList<>(); - metas.add( - new GlobalIndexIOMeta( - result.fileName(), - fileIO.getFileSize(new Path(indexPath, result.fileName())), - result.rowRange(), - result.meta())); - - VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); - - // Test search - query with first vector should return itself - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); - assertThat(searchResult).isNotNull(); - - reader.close(); - } - - @Test - public void testSearchSimilarVectors() throws IOException { - Options options = createDefaultOptions(); - int dimension = 64; - options.setInteger("vector.dim", dimension); // Set correct dimension - int numVectors = 50; - - // Create vectors with known similarities - float[] baseVector = new float[dimension]; - for (int i = 0; i < dimension; i++) { - baseVector[i] = 1.0f; - } - - GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); - VectorGlobalIndexWriter writer = - new VectorGlobalIndexWriter(fileWriter, vectorType, options); - - // Write base vector and similar variants - writer.write(new FloatVectorKey(0, baseVector)); - - for (int i = 1; i < numVectors; i++) { - float[] variant = baseVector.clone(); - // Add small noise - variant[i % dimension] += (i % 2 == 0 ? 0.1f : -0.1f); - writer.write(new FloatVectorKey(i, variant)); - } - - List results = writer.finish(); - assertThat(results).hasSize(1); - - GlobalIndexWriter.ResultEntry result = results.get(0); - GlobalIndexFileReader fileReader = createFileReader(indexPath); - List metas = new ArrayList<>(); - metas.add( - new GlobalIndexIOMeta( - result.fileName(), - fileIO.getFileSize(new Path(indexPath, result.fileName())), - result.rowRange(), - result.meta())); - - VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); - - // Search with base vector should return similar vectors - GlobalIndexResult searchResult = reader.search(baseVector, 10); - assertThat(searchResult).isNotNull(); - - reader.close(); - } - @Test public void testDifferentSimilarityFunctions() throws IOException { int dimension = 32; @@ -188,9 +95,8 @@ public void testDifferentSimilarityFunctions() throws IOException { String[] metrics = {"COSINE", "DOT_PRODUCT", "EUCLIDEAN"}; for (String metric : metrics) { - Options options = createDefaultOptions(); + Options options = createDefaultOptions(dimension); options.setString("vector.metric", metric); - options.setInteger("vector.dim", dimension); // Set correct dimension Path metricIndexPath = new Path(indexPath, metric.toLowerCase()); GlobalIndexFileWriter fileWriter = createFileWriter(metricIndexPath); @@ -199,7 +105,7 @@ public void testDifferentSimilarityFunctions() throws IOException { List testVectors = generateRandomVectors(numVectors, dimension); for (int i = 0; i < numVectors; i++) { - writer.write(new FloatVectorKey(i, testVectors.get(i))); + writer.write(new FloatVectorIndex(i, testVectors.get(i))); } List results = writer.finish(); @@ -230,8 +136,7 @@ public void testDifferentDimensions() throws IOException { int[] dimensions = {8, 32, 128, 256}; for (int dimension : dimensions) { - Options options = createDefaultOptions(); - options.setInteger("vector.dim", dimension); + Options options = createDefaultOptions(dimension); Path dimIndexPath = new Path(indexPath, "dim_" + dimension); GlobalIndexFileWriter fileWriter = createFileWriter(dimIndexPath); @@ -241,7 +146,7 @@ public void testDifferentDimensions() throws IOException { int numVectors = 10; List testVectors = generateRandomVectors(numVectors, dimension); for (int i = 0; i < numVectors; i++) { - writer.write(new FloatVectorKey(i, testVectors.get(i))); + writer.write(new FloatVectorIndex(i, testVectors.get(i))); } List results = writer.finish(); @@ -269,8 +174,7 @@ public void testDifferentDimensions() throws IOException { @Test public void testDimensionMismatch() throws IOException { - Options options = createDefaultOptions(); - options.setInteger("vector.dim", 64); + Options options = createDefaultOptions(64); GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); VectorGlobalIndexWriter writer = @@ -278,242 +182,27 @@ public void testDimensionMismatch() throws IOException { // Try to write vector with wrong dimension float[] wrongDimVector = new float[32]; // Wrong dimension - assertThatThrownBy(() -> writer.write(new FloatVectorKey(0, wrongDimVector))) + assertThatThrownBy(() -> writer.write(new FloatVectorIndex(0, wrongDimVector))) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("dimension mismatch"); } @Test - public void testEmptyIndex() throws IOException { - Options options = createDefaultOptions(); - - GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); - VectorGlobalIndexWriter writer = - new VectorGlobalIndexWriter(fileWriter, vectorType, options); - - List results = writer.finish(); - assertThat(results).isEmpty(); - } - - @Test - public void testHNSWParameters() throws IOException { - Options options = createDefaultOptions(); - options.setInteger("vector.m", 32); // Larger M for better recall - options.setInteger("vector.ef-construction", 200); // Larger ef for better quality - - int dimension = 64; - options.setInteger("vector.dim", dimension); // Set correct dimension - int numVectors = 50; - - GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); - VectorGlobalIndexWriter writer = - new VectorGlobalIndexWriter(fileWriter, vectorType, options); - - List testVectors = generateRandomVectors(numVectors, dimension); - for (int i = 0; i < numVectors; i++) { - writer.write(new FloatVectorKey(i, testVectors.get(i))); - } - - List results = writer.finish(); - assertThat(results).hasSize(1); - - GlobalIndexWriter.ResultEntry result = results.get(0); - GlobalIndexFileReader fileReader = createFileReader(indexPath); - List metas = new ArrayList<>(); - metas.add( - new GlobalIndexIOMeta( - result.fileName(), - fileIO.getFileSize(new Path(indexPath, result.fileName())), - result.rowRange(), - result.meta())); - - VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); - - // Verify search works with custom HNSW parameters - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); - assertThat(searchResult).isNotNull(); - - reader.close(); - } - - @Test - public void testLargeK() throws IOException { - Options options = createDefaultOptions(); - int dimension = 32; - options.setInteger("vector.dim", dimension); // Set correct dimension - int numVectors = 100; - - GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); - VectorGlobalIndexWriter writer = - new VectorGlobalIndexWriter(fileWriter, vectorType, options); - - List testVectors = generateRandomVectors(numVectors, dimension); - for (int i = 0; i < numVectors; i++) { - writer.write(new FloatVectorKey(i, testVectors.get(i))); - } - - List results = writer.finish(); - assertThat(results).hasSize(1); - - GlobalIndexWriter.ResultEntry result = results.get(0); - GlobalIndexFileReader fileReader = createFileReader(indexPath); - List metas = new ArrayList<>(); - metas.add( - new GlobalIndexIOMeta( - result.fileName(), - fileIO.getFileSize(new Path(indexPath, result.fileName())), - result.rowRange(), - result.meta())); - - VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); - - // Search with k larger than number of vectors - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 200); - assertThat(searchResult).isNotNull(); - assertThat(searchResult).isNotNull(); - - reader.close(); - } - - @Test - public void testEndToEndMultiShardScenario() throws IOException { - Options options = createDefaultOptions(); - int dimension = 128; - options.setInteger("vector.dim", dimension); - options.setString("vector.metric", "COSINE"); - options.setInteger("vector.m", 24); - options.setInteger("vector.ef-construction", 150); - - // Simulate a multi-shard scenario with incremental writes - int totalVectors = 300; - int shardsCount = 3; - int vectorsPerShard = totalVectors / shardsCount; - - List allMetas = new ArrayList<>(); - List allTestVectors = generateRandomVectors(totalVectors, dimension); - - // Write vectors in multiple shards (simulating incremental indexing) - for (int shardIdx = 0; shardIdx < shardsCount; shardIdx++) { - Path shardPath = new Path(indexPath, "shard_" + shardIdx); - GlobalIndexFileWriter fileWriter = createFileWriter(shardPath); - VectorGlobalIndexWriter writer = - new VectorGlobalIndexWriter(fileWriter, vectorType, options); - - int startIdx = shardIdx * vectorsPerShard; - int endIdx = Math.min(startIdx + vectorsPerShard, totalVectors); - - // Write vectors for this shard - for (int i = startIdx; i < endIdx; i++) { - writer.write(new FloatVectorKey(i, allTestVectors.get(i))); - } - - List results = writer.finish(); - assertThat(results).hasSize(1); - - GlobalIndexWriter.ResultEntry result = results.get(0); - allMetas.add( - new GlobalIndexIOMeta( - result.fileName(), - fileIO.getFileSize(new Path(shardPath, result.fileName())), - result.rowRange(), - result.meta())); - } - - // Create a single reader for all shards - GlobalIndexFileReader fileReader = - new GlobalIndexFileReader() { - @Override - public org.apache.paimon.fs.SeekableInputStream getInputStream(String fileName) - throws IOException { - // Try to find the file in any shard directory - for (int shardIdx = 0; shardIdx < shardsCount; shardIdx++) { - Path shardPath = new Path(indexPath, "shard_" + shardIdx); - Path filePath = new Path(shardPath, fileName); - if (fileIO.exists(filePath)) { - return fileIO.newInputStream(filePath); - } - } - throw new IOException("File not found: " + fileName); - } - }; - - VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, allMetas); - - // Test 1: Search with various k values - int[] kValues = {1, 5, 10, 20, 50}; - for (int k : kValues) { - GlobalIndexResult searchResult = reader.search(allTestVectors.get(0), k); - assertThat(searchResult).isNotNull(); - assertThat(containsRowId(searchResult, 0)).isTrue(); - } - - // Test 2: Search with query vectors from different shards - for (int shardIdx = 0; shardIdx < shardsCount; shardIdx++) { - int queryIdx = shardIdx * vectorsPerShard + 5; - GlobalIndexResult searchResult = reader.search(allTestVectors.get(queryIdx), 10); - assertThat(searchResult).isNotNull(); - assertThat(containsRowId(searchResult, queryIdx)) - .as("Search result should contain the query vector's own ID") - .isTrue(); - } - - // Test 3: Create a known similar vector and verify top result - float[] originalVector = allTestVectors.get(50); - float[] slightlyModifiedVector = originalVector.clone(); - // Add minimal noise to create a very similar vector - for (int i = 0; i < dimension; i++) { - slightlyModifiedVector[i] += (i % 2 == 0 ? 0.001f : -0.001f); - } - normalize(slightlyModifiedVector); - - GlobalIndexResult similarSearchResult = reader.search(slightlyModifiedVector, 5); - assertThat(similarSearchResult).isNotNull(); - assertThat(containsRowId(similarSearchResult, 50)) - .as("Search with very similar vector should find the original") - .isTrue(); - - // Test 4: Verify search quality with a synthetic cluster - // Create a cluster center - float[] clusterCenter = new float[dimension]; - Random random = new Random(123); - for (int i = 0; i < dimension; i++) { - clusterCenter[i] = random.nextFloat() * 2 - 1; - } - normalize(clusterCenter); - - GlobalIndexResult clusterSearchResult = reader.search(clusterCenter, 20); - assertThat(clusterSearchResult).isNotNull(); - assertThat(clusterSearchResult.iterator().hasNext()).isTrue(); - - // Test 5: Edge case - search with k larger than total vectors - GlobalIndexResult largeKResult = reader.search(allTestVectors.get(10), totalVectors * 2); - assertThat(largeKResult).isNotNull(); - - // Test 6: Multiple consecutive searches (stress test) - for (int i = 0; i < 20; i++) { - int randomIdx = random.nextInt(totalVectors); - GlobalIndexResult result = reader.search(allTestVectors.get(randomIdx), 15); - assertThat(result).isNotNull(); - } - - reader.close(); - } - - @Test - public void testVectorIndexEndToEnd() throws IOException { - Options options = createDefaultOptions(); + public void testFloatVectorIndexEndToEnd() throws IOException { int dimension = 2; - options.setInteger("vector.dim", dimension); - options.setString("vector.metric", "EUCLIDEAN"); - options.setInteger("vector.m", 16); - options.setInteger("vector.ef-construction", 100); + Options options = createDefaultOptions(dimension); + options.setInteger("vector.size-per-index", 3); - // Create semantic vectors representing "documents" - // Apple [0.9, 0.1] and Banana [0.8, 0.2] are similar (fruits) - // Car [0.1, 0.9] is different (vehicle) + // Apple [0.9, 0.1], Banana [0.8, 0.2], and Orange [0.85, 0.15] are close in Euclidean + // distance (fruits) + // Car [0.1, 0.9], Bike [0.15, 0.85], and Truck [0.2, 0.8] are close to each other in + // Euclidean distance (vehicles) float[] appleVector = new float[] {0.9f, 0.1f}; float[] bananaVector = new float[] {0.8f, 0.2f}; float[] carVector = new float[] {0.1f, 0.9f}; + float[] orangeVector = new float[] {0.85f, 0.15f}; + float[] bikeVector = new float[] {0.15f, 0.85f}; + float[] truckVector = new float[] {0.2f, 0.8f}; GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); VectorGlobalIndexWriter writer = @@ -523,93 +212,83 @@ public void testVectorIndexEndToEnd() throws IOException { long appleId = 0L; long bananaId = 1L; long carId = 2L; + long orangeId = 3L; + long bikeId = 4L; + long truckId = 5L; - writer.write(new FloatVectorKey(appleId, appleVector)); - writer.write(new FloatVectorKey(bananaId, bananaVector)); - writer.write(new FloatVectorKey(carId, carVector)); + writer.write(new FloatVectorIndex(appleId, appleVector)); + writer.write(new FloatVectorIndex(bananaId, bananaVector)); + writer.write(new FloatVectorIndex(carId, carVector)); + writer.write(new FloatVectorIndex(orangeId, orangeVector)); + writer.write(new FloatVectorIndex(bikeId, bikeVector)); + writer.write(new FloatVectorIndex(truckId, truckVector)); List results = writer.finish(); - assertThat(results).hasSize(1); + assertThat(results).hasSize(2); - GlobalIndexWriter.ResultEntry result = results.get(0); GlobalIndexFileReader fileReader = createFileReader(indexPath); List metas = new ArrayList<>(); - metas.add( - new GlobalIndexIOMeta( - result.fileName(), - fileIO.getFileSize(new Path(indexPath, result.fileName())), - result.rowRange(), - result.meta())); + for (GlobalIndexWriter.ResultEntry result : results) { + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(indexPath, result.fileName())), + result.rowRange(), + result.meta())); + } VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); - // Test 1: Query with vector similar to "Apple" - should find Apple and Banana + // Test 1: Query with vector similar to "Apple" - should find Apple, Banana, and Orange // Query: [0.85, 0.15] is between Apple and Banana float[] queryVector = new float[] {0.85f, 0.15f}; - GlobalIndexResult searchResult = reader.search(queryVector, 2); + GlobalIndexResult searchResult = reader.search(queryVector, 1); assertThat(searchResult).isNotNull(); - assertThat(searchResult.iterator().hasNext()).as("Search should return results").isTrue(); + assertThat(searchResult.results().iterator().hasNext()) + .as("Search should return results") + .isTrue(); // Collect all result IDs List resultIds = new ArrayList<>(); - for (Long id : searchResult) { - resultIds.add(id); - } + searchResult.results().iterator().forEachRemaining(resultIds::add); - // Should find Apple and Banana (similar fruits), not Car - assertThat(resultIds).as("Should find Apple (most similar)").contains(appleId); - assertThat(resultIds).as("Should find Banana (second most similar)").contains(bananaId); - assertThat(resultIds.size()).as("Should return top 2 results").isLessThanOrEqualTo(2); + // Should find Apple, Banana, and Orange (similar fruits), not vehicles + assertThat(resultIds).containsAnyOf(appleId, bananaId, orangeId); // Test 2: Query with exact match - should find exact document GlobalIndexResult exactMatchResult = reader.search(appleVector, 1); - assertThat(containsRowId(exactMatchResult, appleId)) - .as("Exact match query should find Apple") - .isTrue(); - - // Test 3: Query with vector similar to "Car" - should find Car first - float[] carQueryVector = new float[] {0.15f, 0.85f}; - GlobalIndexResult carSearchResult = reader.search(carQueryVector, 1); + assertThat(containsRowId(exactMatchResult, appleId)).isTrue(); - List carResultIds = new ArrayList<>(); - for (Long id : carSearchResult) { - carResultIds.add(id); - } - - assertThat(carResultIds).as("Query similar to Car should find Car").contains(carId); - - // Test 4: Search with larger k than available documents + // Test 3: Search with larger k than available documents GlobalIndexResult largeKResult = reader.search(queryVector, 10); assertThat(largeKResult).isNotNull(); - assertThat(largeKResult.iterator().hasNext()).isTrue(); + assertThat(largeKResult.results().iterator().hasNext()).isTrue(); - // Test 5: Verify all three documents are indexed - GlobalIndexResult allDocsResult = reader.search(new float[] {0.5f, 0.5f}, 3); + // Test 4: Verify all six documents are indexed + GlobalIndexResult allDocsResult = reader.search(new float[] {0.5f, 0.5f}, 6); List allResultIds = new ArrayList<>(); - for (Long id : allDocsResult) { - allResultIds.add(id); - } + allDocsResult.results().iterator().forEachRemaining(allResultIds::add); assertThat(allResultIds.size()) - .as("Should be able to retrieve all 3 documents") + .as("Should be able to retrieve all 6 documents") .isGreaterThanOrEqualTo(1) - .isLessThanOrEqualTo(3); + .isLessThanOrEqualTo(6); reader.close(); } - private Options createDefaultOptions() { + private Options createDefaultOptions(int dimension) { Options options = new Options(); - options.setInteger("vector.dim", 128); - options.setString("vector.metric", "COSINE"); + options.setInteger("vector.dim", dimension); + options.setString("vector.metric", "EUCLIDEAN"); options.setInteger("vector.m", 16); options.setInteger("vector.ef-construction", 100); return options; } private List generateRandomVectors(int count, int dimension) { - Random random = new Random(42); // Fixed seed for reproducibility + Random random = new Random(42); List vectors = new ArrayList<>(); for (int i = 0; i < count; i++) { float[] vector = new float[dimension]; @@ -617,132 +296,24 @@ private List generateRandomVectors(int count, int dimension) { vector[j] = random.nextFloat() * 2 - 1; // Range [-1, 1] } // Normalize for cosine similarity - normalize(vector); + float norm = 0; + for (float v : vector) { + norm += v * v; + } + norm = (float) Math.sqrt(norm); + if (norm > 0) { + for (int m = 0; m < vector.length; m++) { + vector[m] /= norm; + } + } vectors.add(vector); } return vectors; } - private void normalize(float[] vector) { - float norm = 0; - for (float v : vector) { - norm += v * v; - } - norm = (float) Math.sqrt(norm); - if (norm > 0) { - for (int i = 0; i < vector.length; i++) { - vector[i] /= norm; - } - } - } - private boolean containsRowId(GlobalIndexResult result, long rowId) { - for (Long id : result) { - if (id == rowId) { - return true; - } - } - return false; - } - - @Test - public void testByteVectorIndexEndToEnd() throws IOException { - Options options = createDefaultOptions(); - int dimension = 2; - options.setInteger("vector.dim", dimension); - options.setString("vector.metric", "EUCLIDEAN"); - options.setInteger("vector.m", 16); - options.setInteger("vector.ef-construction", 100); - - // Create semantic vectors representing "documents" using byte vectors - // Apple [90, 10] and Banana [80, 20] are similar (fruits) - // Car [10, 90] is different (vehicle) - byte[] appleVector = new byte[] {90, 10}; - byte[] bananaVector = new byte[] {80, 20}; - byte[] carVector = new byte[] {10, 90}; - - DataType byteVectorType = new ArrayType(new TinyIntType()); - GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); - VectorGlobalIndexWriter writer = - new VectorGlobalIndexWriter(fileWriter, byteVectorType, options); - - // Write vectors with row IDs representing documents - long appleId = 0L; - long bananaId = 1L; - long carId = 2L; - - writer.write(new ByteVectorKey(appleId, appleVector)); - writer.write(new ByteVectorKey(bananaId, bananaVector)); - writer.write(new ByteVectorKey(carId, carVector)); - - List results = writer.finish(); - assertThat(results).hasSize(1); - - GlobalIndexWriter.ResultEntry result = results.get(0); - GlobalIndexFileReader fileReader = createFileReader(indexPath); - List metas = new ArrayList<>(); - metas.add( - new GlobalIndexIOMeta( - result.fileName(), - fileIO.getFileSize(new Path(indexPath, result.fileName())), - result.rowRange(), - result.meta())); - - VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); - - // Test 1: Query with vector similar to "Apple" - should find Apple and Banana - // Query: [85, 15] is between Apple and Banana - byte[] queryVector = new byte[] {85, 15}; - GlobalIndexResult searchResult = reader.search(queryVector, 2); - - assertThat(searchResult).isNotNull(); - assertThat(searchResult.iterator().hasNext()).as("Search should return results").isTrue(); - - // Collect all result IDs List resultIds = new ArrayList<>(); - for (Long id : searchResult) { - resultIds.add(id); - } - - // Should find Apple and Banana (similar fruits), not Car - assertThat(resultIds).as("Should find Apple (most similar)").contains(appleId); - assertThat(resultIds).as("Should find Banana (second most similar)").contains(bananaId); - assertThat(resultIds.size()).as("Should return top 2 results").isLessThanOrEqualTo(2); - - // Test 2: Query with exact match - should find exact document - GlobalIndexResult exactMatchResult = reader.search(appleVector, 1); - assertThat(containsRowId(exactMatchResult, appleId)) - .as("Exact match query should find Apple") - .isTrue(); - - // Test 3: Query with vector similar to "Car" - should find Car first - byte[] carQueryVector = new byte[] {15, 85}; - GlobalIndexResult carSearchResult = reader.search(carQueryVector, 1); - - List carResultIds = new ArrayList<>(); - for (Long id : carSearchResult) { - carResultIds.add(id); - } - - assertThat(carResultIds).as("Query similar to Car should find Car").contains(carId); - - // Test 4: Search with larger k than available documents - GlobalIndexResult largeKResult = reader.search(queryVector, 10); - assertThat(largeKResult).isNotNull(); - assertThat(largeKResult.iterator().hasNext()).isTrue(); - - // Test 5: Verify all three documents are indexed - GlobalIndexResult allDocsResult = reader.search(new byte[] {50, 50}, 3); - List allResultIds = new ArrayList<>(); - for (Long id : allDocsResult) { - allResultIds.add(id); - } - - assertThat(allResultIds.size()) - .as("Should be able to retrieve all 3 documents") - .isGreaterThanOrEqualTo(1) - .isLessThanOrEqualTo(3); - - reader.close(); + result.results().iterator().forEachRemaining(resultIds::add); + return resultIds.contains(rowId); } } From 0219ab8f59ac3572ec2a806768718ddcac583cc1 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 5 Dec 2025 16:31:58 +0800 Subject: [PATCH 10/11] delete no need case --- .../vector/VectorGlobalIndexReader.java | 33 ++++++-- .../vector/VectorGlobalIndexWriter.java | 11 +-- .../paimon/vector/VectorIndexMetadata.java | 71 ----------------- .../paimon/vector/VectorGlobalIndexTest.java | 76 ++++--------------- 4 files changed, 44 insertions(+), 147 deletions(-) delete mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java index cb38d68c6081..495704526551 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -41,7 +41,9 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.ArrayList; +import java.util.Comparator; import java.util.List; +import java.util.PriorityQueue; import java.util.Set; /** @@ -101,28 +103,47 @@ public void close() throws IOException { } private GlobalIndexResult search(Query query, int k) { - RoaringNavigableMap64 roaringBitmap64 = new RoaringNavigableMap64(); + PriorityQueue topK = + new PriorityQueue<>(Comparator.comparingDouble(sr -> sr.score)); for (IndexSearcher searcher : searchers) { try { - // Execute search TopDocs topDocs = searcher.search(query, k); StoredFields storedFields = searcher.storedFields(); Set fieldsToLoad = Set.of(VectorIndex.ROW_ID_FIELD); - // Collect row IDs from results for (org.apache.lucene.search.ScoreDoc scoreDoc : topDocs.scoreDocs) { - float rawScore = scoreDoc.score; Document doc = storedFields.document(scoreDoc.doc, fieldsToLoad); long rowId = doc.getField(VectorIndex.ROW_ID_FIELD).numericValue().longValue(); - roaringBitmap64.add(rowId); + if (topK.size() < k) { + topK.offer(new ScoredRow(rowId, scoreDoc.score)); + } else { + if (topK.peek() != null && scoreDoc.score > topK.peek().score) { + topK.poll(); + topK.offer(new ScoredRow(rowId, scoreDoc.score)); + } + } } } catch (IOException e) { throw new RuntimeException("Failed to search vector index", e); } } - + RoaringNavigableMap64 roaringBitmap64 = new RoaringNavigableMap64(); + for (ScoredRow scoredRow : topK) { + roaringBitmap64.add(scoredRow.rowId); + } return GlobalIndexResult.create(() -> roaringBitmap64); } + /** Helper class to store row ID with its score. */ + private static class ScoredRow { + final long rowId; + final float score; + + ScoredRow(long rowId, float score) { + this.rowId = rowId; + this.score = score; + } + } + private void loadIndices(GlobalIndexFileReader fileReader, List files) throws IOException { for (GlobalIndexIOMeta meta : files) { diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java index 3ed77b1bdb69..31b1e5d62afa 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -130,15 +130,6 @@ public List finish() { this.vectorOptions.efConstruction(), this.vectorOptions.writeBufferSize()); - // Create metadata - VectorIndexMetadata metadata = - new VectorIndexMetadata( - vectorOptions.dimension(), - vectorOptions.metric(), - vectorOptions.m(), - vectorOptions.efConstruction()); - byte[] metaBytes = VectorIndexMetadata.serializeMetadata(metadata); - // Write to file String fileName = fileWriter.newFileName(VectorGlobalIndexerFactory.IDENTIFIER); try (OutputStream out = fileWriter.newOutputStream(fileName)) { @@ -148,7 +139,7 @@ public List finish() { long maxRowIdInBatch = batchVectors.get(batchVectors.size() - 1).rowId(); results.add( ResultEntry.of( - fileName, metaBytes, new Range(minRowIdInBatch, maxRowIdInBatch))); + fileName, null, new Range(minRowIdInBatch, maxRowIdInBatch))); } return results; diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java deleted file mode 100644 index 08321d278942..000000000000 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexMetadata.java +++ /dev/null @@ -1,71 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.paimon.vector; - -import java.io.ByteArrayOutputStream; -import java.io.DataOutputStream; -import java.io.IOException; -import java.io.Serializable; - -/** Metadata for the vector index, including dimension, similarity function, and HNSW parameters. */ -public class VectorIndexMetadata implements Serializable { - - private static final long serialVersionUID = 1L; - - private final int dimension; - private final String similarityFunction; - private final int m; - private final int efConstruction; - - public VectorIndexMetadata( - int dimension, String similarityFunction, int m, int efConstruction) { - this.dimension = dimension; - this.similarityFunction = similarityFunction; - this.m = m; - this.efConstruction = efConstruction; - } - - public static byte[] serializeMetadata(VectorIndexMetadata vectorIndexMetadata) - throws IOException { - ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); - try (DataOutputStream dataOut = new DataOutputStream(byteOut)) { - dataOut.writeInt(vectorIndexMetadata.dimension()); - dataOut.writeUTF(vectorIndexMetadata.similarityFunction()); - dataOut.writeInt(vectorIndexMetadata.m()); - dataOut.writeInt(vectorIndexMetadata.efConstruction()); - } - return byteOut.toByteArray(); - } - - public int dimension() { - return dimension; - } - - public String similarityFunction() { - return similarityFunction; - } - - public int m() { - return m; - } - - public int efConstruction() { - return efConstruction; - } -} diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java index 47b3dd5d703d..d4a528d155aa 100644 --- a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java @@ -193,35 +193,18 @@ public void testFloatVectorIndexEndToEnd() throws IOException { Options options = createDefaultOptions(dimension); options.setInteger("vector.size-per-index", 3); - // Apple [0.9, 0.1], Banana [0.8, 0.2], and Orange [0.85, 0.15] are close in Euclidean - // distance (fruits) - // Car [0.1, 0.9], Bike [0.15, 0.85], and Truck [0.2, 0.8] are close to each other in - // Euclidean distance (vehicles) - float[] appleVector = new float[] {0.9f, 0.1f}; - float[] bananaVector = new float[] {0.8f, 0.2f}; - float[] carVector = new float[] {0.1f, 0.9f}; - float[] orangeVector = new float[] {0.85f, 0.15f}; - float[] bikeVector = new float[] {0.15f, 0.85f}; - float[] truckVector = new float[] {0.2f, 0.8f}; + float[][] vectors = + new float[][] { + new float[] {1.0f, 0.0f}, new float[] {0.95f, 0.1f}, new float[] {0.1f, 0.95f}, + new float[] {0.98f, 0.05f}, new float[] {0.0f, 1.0f}, new float[] {0.05f, 0.98f} + }; GlobalIndexFileWriter fileWriter = createFileWriter(indexPath); VectorGlobalIndexWriter writer = new VectorGlobalIndexWriter(fileWriter, vectorType, options); - - // Write vectors with row IDs representing documents - long appleId = 0L; - long bananaId = 1L; - long carId = 2L; - long orangeId = 3L; - long bikeId = 4L; - long truckId = 5L; - - writer.write(new FloatVectorIndex(appleId, appleVector)); - writer.write(new FloatVectorIndex(bananaId, bananaVector)); - writer.write(new FloatVectorIndex(carId, carVector)); - writer.write(new FloatVectorIndex(orangeId, orangeVector)); - writer.write(new FloatVectorIndex(bikeId, bikeVector)); - writer.write(new FloatVectorIndex(truckId, truckVector)); + for (int i = 0; i < vectors.length; i++) { + writer.write(new FloatVectorIndex(i, vectors[i])); + } List results = writer.finish(); assertThat(results).hasSize(2); @@ -239,42 +222,15 @@ public void testFloatVectorIndexEndToEnd() throws IOException { VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); - // Test 1: Query with vector similar to "Apple" - should find Apple, Banana, and Orange - // Query: [0.85, 0.15] is between Apple and Banana - float[] queryVector = new float[] {0.85f, 0.15f}; - GlobalIndexResult searchResult = reader.search(queryVector, 1); - - assertThat(searchResult).isNotNull(); - assertThat(searchResult.results().iterator().hasNext()) - .as("Search should return results") - .isTrue(); - - // Collect all result IDs - List resultIds = new ArrayList<>(); - searchResult.results().iterator().forEachRemaining(resultIds::add); - - // Should find Apple, Banana, and Orange (similar fruits), not vehicles - assertThat(resultIds).containsAnyOf(appleId, bananaId, orangeId); - - // Test 2: Query with exact match - should find exact document - GlobalIndexResult exactMatchResult = reader.search(appleVector, 1); - assertThat(containsRowId(exactMatchResult, appleId)).isTrue(); - - // Test 3: Search with larger k than available documents - GlobalIndexResult largeKResult = reader.search(queryVector, 10); - assertThat(largeKResult).isNotNull(); - assertThat(largeKResult.results().iterator().hasNext()).isTrue(); - - // Test 4: Verify all six documents are indexed - GlobalIndexResult allDocsResult = reader.search(new float[] {0.5f, 0.5f}, 6); - List allResultIds = new ArrayList<>(); - allDocsResult.results().iterator().forEachRemaining(allResultIds::add); - - assertThat(allResultIds.size()) - .as("Should be able to retrieve all 6 documents") - .isGreaterThanOrEqualTo(1) - .isLessThanOrEqualTo(6); + GlobalIndexResult result = reader.search(vectors[0], 1); + assertThat(result.results().getLongCardinality()).isEqualTo(1); + assertThat(containsRowId(result, 0)).isTrue(); + float[] queryVector = new float[] {0.85f, 0.15f}; + result = reader.search(queryVector, 2); + assertThat(result.results().getLongCardinality()).isEqualTo(2); + assertThat(containsRowId(result, 1)).isTrue(); + assertThat(containsRowId(result, 3)).isTrue(); reader.close(); } From 72c038910345d3454e8d6100b574219f3c93921e Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 5 Dec 2025 17:58:55 +0800 Subject: [PATCH 11/11] update test workflow and delete no need case --- .github/workflows/check-licensing.yml | 1 + .github/workflows/docs-tests.yml | 3 + .github/workflows/e2e-tests-flink-1.x.yml | 5 + .../workflows/e2e-tests-flink-2.x-jdk11.yml | 5 + .github/workflows/utitcase-flink-1.x.yml | 5 + .../workflows/utitcase-flink-2.x-jdk11.yml | 5 + .github/workflows/utitcase-jdk11.yml | 5 + .github/workflows/utitcase-paimon-vector.yml | 10 +- .github/workflows/utitcase-spark-3.x.yml | 5 + .github/workflows/utitcase-spark-4.x.yml | 5 + .github/workflows/utitcase.yml | 2 + .../apache/paimon/vector/ByteVectorIndex.java | 8 +- .../paimon/vector/FloatVectorIndex.java | 8 +- .../paimon/vector/IndexMMapDirectory.java | 2 +- .../vector/VectorGlobalIndexReader.java | 134 ++++++++++------ .../vector/VectorGlobalIndexWriter.java | 148 +++++------------- .../org/apache/paimon/vector/VectorIndex.java | 15 +- .../paimon/vector/VectorIndexOptions.java | 12 +- .../apache/paimon/vector/VectorMetric.java | 51 ++++++ .../paimon/vector/VectorGlobalIndexTest.java | 44 +++--- 20 files changed, 274 insertions(+), 199 deletions(-) create mode 100644 paimon-vector/src/main/java/org/apache/paimon/vector/VectorMetric.java diff --git a/.github/workflows/check-licensing.yml b/.github/workflows/check-licensing.yml index 0cb422ac3323..f9c8aae6c1a6 100644 --- a/.github/workflows/check-licensing.yml +++ b/.github/workflows/check-licensing.yml @@ -22,6 +22,7 @@ on: - 'docs/**' - '**/*.md' - 'paimon-python/**' + - 'paimon-vector/**' env: MAVEN_OPTS: -Dmaven.wagon.httpconnectionManager.ttlSeconds=30 -Dmaven.wagon.http.retryHandler.requestSentEnabled=true diff --git a/.github/workflows/docs-tests.yml b/.github/workflows/docs-tests.yml index 65fef1c77ee1..c714e76af042 100644 --- a/.github/workflows/docs-tests.yml +++ b/.github/workflows/docs-tests.yml @@ -16,6 +16,9 @@ name: docs-tests on: push: + paths-ignore: + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' pull_request: paths: - 'docs/**' diff --git a/.github/workflows/e2e-tests-flink-1.x.yml b/.github/workflows/e2e-tests-flink-1.x.yml index 48a754b9b91d..beaee9124141 100644 --- a/.github/workflows/e2e-tests-flink-1.x.yml +++ b/.github/workflows/e2e-tests-flink-1.x.yml @@ -20,12 +20,17 @@ name: End to End Tests Flink 1.x on: push: + paths-ignore: + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' pull_request: paths-ignore: - 'docs/**' - '**/*.md' - 'paimon-python/**' - '.github/workflows/paimon-python-checks.yml' + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' env: JDK_VERSION: 8 diff --git a/.github/workflows/e2e-tests-flink-2.x-jdk11.yml b/.github/workflows/e2e-tests-flink-2.x-jdk11.yml index 5f977784678f..40607890da7d 100644 --- a/.github/workflows/e2e-tests-flink-2.x-jdk11.yml +++ b/.github/workflows/e2e-tests-flink-2.x-jdk11.yml @@ -20,12 +20,17 @@ name: End to End Tests Flink 2.x on JDK 11 on: push: + paths-ignore: + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' pull_request: paths-ignore: - 'docs/**' - '**/*.md' - 'paimon-python/**' - '.github/workflows/paimon-python-checks.yml' + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' env: JDK_VERSION: 11 diff --git a/.github/workflows/utitcase-flink-1.x.yml b/.github/workflows/utitcase-flink-1.x.yml index 7ed24e564481..488ed5b528c3 100644 --- a/.github/workflows/utitcase-flink-1.x.yml +++ b/.github/workflows/utitcase-flink-1.x.yml @@ -19,12 +19,17 @@ name: UTCase and ITCase Flink 1.x on: push: + paths-ignore: + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' pull_request: paths-ignore: - 'docs/**' - '**/*.md' - 'paimon-python/**' - '.github/workflows/paimon-python-checks.yml' + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' env: JDK_VERSION: 8 diff --git a/.github/workflows/utitcase-flink-2.x-jdk11.yml b/.github/workflows/utitcase-flink-2.x-jdk11.yml index 8e90e679cbb5..fedfe95eb656 100644 --- a/.github/workflows/utitcase-flink-2.x-jdk11.yml +++ b/.github/workflows/utitcase-flink-2.x-jdk11.yml @@ -20,12 +20,17 @@ name: UTCase and ITCase Flink 2.x on JDK 11 on: push: + paths-ignore: + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' pull_request: paths-ignore: - 'docs/**' - '**/*.md' - 'paimon-python/**' - '.github/workflows/paimon-python-checks.yml' + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' env: JDK_VERSION: 11 diff --git a/.github/workflows/utitcase-jdk11.yml b/.github/workflows/utitcase-jdk11.yml index 0386a02e14d5..870ed13dc53c 100644 --- a/.github/workflows/utitcase-jdk11.yml +++ b/.github/workflows/utitcase-jdk11.yml @@ -20,12 +20,17 @@ name: UTCase and ITCase Others on JDK 11 on: push: + paths-ignore: + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' pull_request: paths-ignore: - 'docs/**' - '**/*.md' - 'paimon-python/**' - '.github/workflows/paimon-python-checks.yml' + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' env: JDK_VERSION: 11 diff --git a/.github/workflows/utitcase-paimon-vector.yml b/.github/workflows/utitcase-paimon-vector.yml index 1699efb3e12c..cbf61f26bda2 100644 --- a/.github/workflows/utitcase-paimon-vector.yml +++ b/.github/workflows/utitcase-paimon-vector.yml @@ -21,11 +21,9 @@ name: UTCase and ITCase Paimon Vector on: push: pull_request: - paths-ignore: - - 'docs/**' - - '**/*.md' - - 'paimon-python/**' - - '.github/workflows/paimon-python-checks.yml' + paths: + - 'paimon-vector/**' + - 'pom.xml' env: JDK_VERSION: 11 @@ -57,7 +55,7 @@ jobs: . .github/workflows/utils.sh jvm_timezone=$(random_timezone) echo "JVM timezone is set to $jvm_timezone" - TEST_MODULE="org.apache.paimon:paimon-vector" + TEST_MODULE="org.apache.paimon:paimon-vector" echo "Start testing module: $TEST_MODULE" mvn -T 2C -B test verify -Ppaimon-vector -pl "${TEST_MODULE}" -Duser.timezone=$jvm_timezone echo "All modules tested" diff --git a/.github/workflows/utitcase-spark-3.x.yml b/.github/workflows/utitcase-spark-3.x.yml index 40ae73fab427..0dc98dea838d 100644 --- a/.github/workflows/utitcase-spark-3.x.yml +++ b/.github/workflows/utitcase-spark-3.x.yml @@ -20,12 +20,17 @@ name: UTCase and ITCase Spark 3.x on: push: + paths-ignore: + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' pull_request: paths-ignore: - 'docs/**' - '**/*.md' - 'paimon-python/**' - '.github/workflows/paimon-python-checks.yml' + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' env: JDK_VERSION: 8 diff --git a/.github/workflows/utitcase-spark-4.x.yml b/.github/workflows/utitcase-spark-4.x.yml index b9b3071de552..096ab1a3dc2b 100644 --- a/.github/workflows/utitcase-spark-4.x.yml +++ b/.github/workflows/utitcase-spark-4.x.yml @@ -20,12 +20,17 @@ name: UTCase and ITCase Spark 4.x on: push: + paths-ignore: + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' pull_request: paths-ignore: - 'docs/**' - '**/*.md' - 'paimon-python/**' - '.github/workflows/paimon-python-checks.yml' + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' env: JDK_VERSION: 17 diff --git a/.github/workflows/utitcase.yml b/.github/workflows/utitcase.yml index 3ba3df80fd85..87cf3461a628 100644 --- a/.github/workflows/utitcase.yml +++ b/.github/workflows/utitcase.yml @@ -26,6 +26,8 @@ on: - '**/*.md' - 'paimon-python/**' - '.github/workflows/paimon-python-checks.yml' + - 'paimon-vector/**' + - '.github/workflows/utitcase-paimon-vector.yml' env: JDK_VERSION: 8 diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorIndex.java b/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorIndex.java index 33239b92389e..b0b99d1bab7e 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorIndex.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorIndex.java @@ -23,7 +23,7 @@ import org.apache.lucene.index.VectorSimilarityFunction; /** Vector index for byte vector. */ -public class ByteVectorIndex extends VectorIndex { +public class ByteVectorIndex extends VectorIndex { private final long rowId; private final byte[] vector; @@ -37,11 +37,17 @@ public long rowId() { return rowId; } + @Override + public long dimension() { + return vector().length; + } + @Override public IndexableField indexableField(VectorSimilarityFunction similarityFunction) { return new KnnByteVectorField(VECTOR_FIELD, this.vector, similarityFunction); } + @Override public byte[] vector() { return vector; } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorIndex.java b/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorIndex.java index f2229b49e411..3b2c9ae3c47b 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorIndex.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorIndex.java @@ -23,7 +23,7 @@ import org.apache.lucene.index.VectorSimilarityFunction; /** Vector index for float vector. */ -public class FloatVectorIndex extends VectorIndex { +public class FloatVectorIndex extends VectorIndex { private final long rowId; private final float[] vector; @@ -32,6 +32,7 @@ public FloatVectorIndex(long rowId, float[] vector) { this.vector = vector; } + @Override public float[] vector() { return vector; } @@ -41,6 +42,11 @@ public long rowId() { return rowId; } + @Override + public long dimension() { + return vector().length; + } + @Override public IndexableField indexableField(VectorSimilarityFunction similarityFunction) { return new KnnFloatVectorField(VECTOR_FIELD, this.vector(), similarityFunction); diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/IndexMMapDirectory.java b/paimon-vector/src/main/java/org/apache/paimon/vector/IndexMMapDirectory.java index 4bfb235dca41..2f50081d9ab3 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/IndexMMapDirectory.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/IndexMMapDirectory.java @@ -31,7 +31,7 @@ public class IndexMMapDirectory implements AutoCloseable { public IndexMMapDirectory() throws IOException { this.path = - java.nio.file.Files.createTempDirectory("paimon-vector-index" + UUID.randomUUID()); + java.nio.file.Files.createTempDirectory("paimon-vector-index-" + UUID.randomUUID()); this.mmapDirectory = new MMapDirectory(path); } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java index 495704526551..b7160bc97e6e 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -85,9 +85,19 @@ public GlobalIndexResult search(byte[] query, int k) { @Override public void close() throws IOException { + Throwable firstException = null; + // Close readers for (IndexSearcher searcher : searchers) { - searcher.getIndexReader().close(); + try { + searcher.getIndexReader().close(); + } catch (Throwable t) { + if (firstException == null) { + firstException = t; + } else { + firstException.addSuppressed(t); + } + } } searchers.clear(); @@ -95,11 +105,26 @@ public void close() throws IOException { for (IndexMMapDirectory directory : directories) { try { directory.close(); - } catch (Exception e) { - throw new RuntimeException(e); + } catch (Throwable t) { + if (firstException == null) { + firstException = t; + } else { + firstException.addSuppressed(t); + } } } directories.clear(); + + if (firstException != null) { + if (firstException instanceof IOException) { + throw (IOException) firstException; + } else if (firstException instanceof RuntimeException) { + throw (RuntimeException) firstException; + } else { + throw new RuntimeException( + "Failed to close vector global index reader", firstException); + } + } } private GlobalIndexResult search(Query query, int k) { @@ -148,47 +173,71 @@ private void loadIndices(GlobalIndexFileReader fileReader, List 0) { - int toRead = (int) Math.min(buffer.length, remaining); - readFully(in, buffer, 0, toRead); - output.writeBytes(buffer, 0, toRead); - remaining -= toRead; + try { + int numFiles = readInt(in); + byte[] buffer = new byte[BUFFER_SIZE]; + for (int i = 0; i < numFiles; i++) { + int nameLength = readInt(in); + byte[] nameBytes = new byte[nameLength]; + readFully(in, nameBytes); + String fileName = new String(nameBytes, StandardCharsets.UTF_8); + long fileLength = readLong(in); + try (IndexOutput output = + indexMMapDirectory.directory().createOutput(fileName, null)) { + long remaining = fileLength; + while (remaining > 0) { + int toRead = (int) Math.min(buffer.length, remaining); + readFully(in, buffer, 0, toRead); + output.writeBytes(buffer, 0, toRead); + remaining -= toRead; + } } } + return indexMMapDirectory; + } catch (Exception e) { + try { + indexMMapDirectory.close(); + } catch (Exception closeEx) { + } + if (e instanceof IOException) { + throw (IOException) e; + } else { + throw new IOException("Failed to deserialize directory", e); + } } - - return indexMMapDirectory; } private int readInt(SeekableInputStream in) throws IOException { @@ -219,21 +268,6 @@ private void readFully(SeekableInputStream in, byte[] buffer, int offset, int le } } - private void deleteDirectory(java.nio.file.Path path) throws IOException { - if (java.nio.file.Files.exists(path)) { - java.nio.file.Files.walk(path) - .sorted(java.util.Comparator.reverseOrder()) - .forEach( - p -> { - try { - java.nio.file.Files.delete(p); - } catch (IOException e) { - // Ignore cleanup errors - } - }); - } - } - // Implementation of FunctionVisitor methods @Override public GlobalIndexResult visitIsNotNull(FieldRef fieldRef) { diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java index 31b1e5d62afa..975db6c0e9e0 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -33,10 +33,11 @@ import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; -import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.OutputStream; +import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; @@ -52,12 +53,11 @@ public class VectorGlobalIndexWriter implements GlobalIndexWriter { private final GlobalIndexFileWriter fileWriter; - private final DataType fieldType; private final VectorIndexOptions vectorOptions; private final VectorSimilarityFunction similarityFunction; private final int sizePerIndex; - private final List vectors; + private final List vectorIndices; public VectorGlobalIndexWriter( GlobalIndexFileWriter fileWriter, DataType fieldType, Options options) { @@ -65,78 +65,57 @@ public VectorGlobalIndexWriter( fieldType instanceof ArrayType, "Vector field type must be ARRAY, but was: " + fieldType); this.fileWriter = fileWriter; - this.fieldType = fieldType; - this.vectors = new ArrayList<>(); + this.vectorIndices = new ArrayList<>(); this.vectorOptions = new VectorIndexOptions(options); - this.similarityFunction = parseMetricToLucene(vectorOptions.metric()); + this.similarityFunction = vectorOptions.metric().vectorSimilarityFunction(); this.sizePerIndex = vectorOptions.sizePerIndex(); } @Override public void write(Object key) { + VectorIndex index; if (key instanceof FloatVectorIndex) { - FloatVectorIndex vectorKey = (FloatVectorIndex) key; - float[] vector = vectorKey.vector(); - - checkArgument( - vector.length == vectorOptions.dimension(), - "Vector dimension mismatch: expected " - + vectorOptions.dimension() - + ", but got " - + vector.length); - - vectors.add(vectorKey); + index = (FloatVectorIndex) key; } else if (key instanceof ByteVectorIndex) { - ByteVectorIndex vectorKey = (ByteVectorIndex) key; - byte[] byteVector = vectorKey.vector(); - - checkArgument( - byteVector.length == vectorOptions.dimension(), - "Vector dimension mismatch: expected " - + vectorOptions.dimension() - + ", but got " - + byteVector.length); - - vectors.add(vectorKey); + index = (ByteVectorIndex) key; } else { throw new IllegalArgumentException( "Unsupported index type: " + key.getClass().getName()); } + index.checkDimension(vectorOptions.dimension()); + vectorIndices.add(index); } @Override public List finish() { try { - if (vectors.isEmpty()) { + if (vectorIndices.isEmpty()) { return new ArrayList<>(); } List results = new ArrayList<>(); // Split vectors into batches if size exceeds sizePerIndex - int totalVectors = vectors.size(); + int totalVectors = vectorIndices.size(); int numBatches = (int) Math.ceil((double) totalVectors / sizePerIndex); for (int batchIndex = 0; batchIndex < numBatches; batchIndex++) { int startIdx = batchIndex * sizePerIndex; int endIdx = Math.min(startIdx + sizePerIndex, totalVectors); - List batchVectors = vectors.subList(startIdx, endIdx); - - // Build index - byte[] indexBytes = - buildIndex( - batchVectors, - this.vectorOptions.m(), - this.vectorOptions.efConstruction(), - this.vectorOptions.writeBufferSize()); + List batchVectorIndices = vectorIndices.subList(startIdx, endIdx); - // Write to file String fileName = fileWriter.newFileName(VectorGlobalIndexerFactory.IDENTIFIER); try (OutputStream out = fileWriter.newOutputStream(fileName)) { - out.write(indexBytes); + buildIndex( + batchVectorIndices, + this.vectorOptions.m(), + this.vectorOptions.efConstruction(), + this.vectorOptions.writeBufferSize(), + out); } - long minRowIdInBatch = batchVectors.get(0).rowId(); - long maxRowIdInBatch = batchVectors.get(batchVectors.size() - 1).rowId(); + long minRowIdInBatch = batchVectorIndices.get(0).rowId(); + long maxRowIdInBatch = + batchVectorIndices.get(batchVectorIndices.size() - 1).rowId(); results.add( ResultEntry.of( fileName, null, new Range(minRowIdInBatch, maxRowIdInBatch))); @@ -148,49 +127,27 @@ public List finish() { } } - private VectorSimilarityFunction parseMetricToLucene(String metric) { - switch (metric.toUpperCase()) { - case "COSINE": - return VectorSimilarityFunction.COSINE; - case "DOT_PRODUCT": - return VectorSimilarityFunction.DOT_PRODUCT; - case "EUCLIDEAN": - return VectorSimilarityFunction.EUCLIDEAN; - case "MAX_INNER_PRODUCT": - return VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; - default: - throw new IllegalArgumentException("Unsupported metric: " + metric); - } - } - - private byte[] buildIndex( - List batchVectors, int m, int efConstruction, int writeBufferSize) + private void buildIndex( + List batchVectors, + int m, + int efConstruction, + int writeBufferSize, + OutputStream out) throws IOException { + IndexWriterConfig config = getIndexWriterConfig(m, efConstruction, writeBufferSize); try (IndexMMapDirectory indexMMapDirectory = new IndexMMapDirectory()) { - // Configure index writer - IndexWriterConfig config = getIndexWriterConfig(m, efConstruction, writeBufferSize); - try (IndexWriter writer = new IndexWriter(indexMMapDirectory.directory(), config)) { - // Add each vector as a document for (VectorIndex vectorIndex : batchVectors) { Document doc = new Document(); - - // Add KNN vector field doc.add(vectorIndex.indexableField(similarityFunction)); - - // Store row ID doc.add(vectorIndex.rowIdStoredField()); - writer.addDocument(doc); } - - // Commit changes writer.commit(); } - // Serialize directory to byte array - return serializeDirectory(indexMMapDirectory.directory()); + serializeDirectory(indexMMapDirectory.directory(), out); } catch (Exception e) { throw new RuntimeException(e); } @@ -199,8 +156,7 @@ private byte[] buildIndex( private static IndexWriterConfig getIndexWriterConfig( int m, int efConstruction, int writeBufferSize) { IndexWriterConfig config = new IndexWriterConfig(); - config.setRAMBufferSizeMB( - writeBufferSize); // Configure RAM buffer size based on user settings + config.setRAMBufferSizeMB(writeBufferSize); config.setCodec( new Lucene912Codec(Lucene912Codec.Mode.BEST_SPEED) { @Override @@ -211,57 +167,33 @@ public KnnVectorsFormat getKnnVectorsFormatForField(String field) { return config; } - private byte[] serializeDirectory(Directory directory) throws IOException { - ByteArrayOutputStream baos = new ByteArrayOutputStream(); - - // Write all files from the directory + private void serializeDirectory(Directory directory, OutputStream out) throws IOException { String[] files = directory.listAll(); - - // Write number of files - baos.write(intToBytes(files.length)); + out.write(intToBytes(files.length)); for (String fileName : files) { - // Write file name length and name byte[] nameBytes = fileName.getBytes(StandardCharsets.UTF_8); - baos.write(intToBytes(nameBytes.length)); - baos.write(nameBytes); - - // Write file content length and content + out.write(intToBytes(nameBytes.length)); + out.write(nameBytes); long fileLength = directory.fileLength(fileName); - baos.write(longToBytes(fileLength)); + out.write(ByteBuffer.allocate(8).putLong(fileLength).array()); - try (org.apache.lucene.store.IndexInput input = directory.openInput(fileName, null)) { + try (org.apache.lucene.store.IndexInput input = + directory.openInput(fileName, IOContext.DEFAULT)) { byte[] buffer = new byte[8192]; long remaining = fileLength; while (remaining > 0) { int toRead = (int) Math.min(buffer.length, remaining); input.readBytes(buffer, 0, toRead); - baos.write(buffer, 0, toRead); + out.write(buffer, 0, toRead); remaining -= toRead; } } } - - return baos.toByteArray(); } private byte[] intToBytes(int value) { - return new byte[] { - (byte) (value >>> 24), (byte) (value >>> 16), (byte) (value >>> 8), (byte) value - }; - } - - private byte[] longToBytes(long value) { - return new byte[] { - (byte) (value >>> 56), - (byte) (value >>> 48), - (byte) (value >>> 40), - (byte) (value >>> 32), - (byte) (value >>> 24), - (byte) (value >>> 16), - (byte) (value >>> 8), - (byte) value - }; + return ByteBuffer.allocate(4).putInt(value).array(); } } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndex.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndex.java index 744b65a9c546..0c9aa729ac48 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndex.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndex.java @@ -23,16 +23,29 @@ import org.apache.lucene.index.VectorSimilarityFunction; /** Vector index interface. */ -public abstract class VectorIndex { +public abstract class VectorIndex { public static final String VECTOR_FIELD = "vector"; public static final String ROW_ID_FIELD = "rowId"; public abstract long rowId(); + public abstract long dimension(); + + public abstract T vector(); + public abstract IndexableField indexableField(VectorSimilarityFunction similarityFunction); public StoredField rowIdStoredField() { return new StoredField(ROW_ID_FIELD, rowId()); } + + public void checkDimension(int dimension) { + if (dimension() != dimension) { + throw new IllegalArgumentException( + String.format( + "Vector dimension mismatch: expected %d, but got %d", + dimension(), dimension)); + } + } } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java index 31883e928aca..a91b293aa7c2 100644 --- a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java @@ -31,10 +31,10 @@ public class VectorIndexOptions { .defaultValue(128) .withDescription("The dimension of the vector"); - public static final ConfigOption VECTOR_METRIC = + public static final ConfigOption VECTOR_METRIC = ConfigOptions.key("vector.metric") - .stringType() - .defaultValue("EUCLIDEAN") + .enumType(VectorMetric.class) + .defaultValue(VectorMetric.EUCLIDEAN) .withDescription( "The similarity metric for vector search (COSINE, DOT_PRODUCT, EUCLIDEAN, MAX_INNER_PRODUCT), and EUCLIDEAN is the default"); @@ -62,10 +62,10 @@ public class VectorIndexOptions { ConfigOptions.key("vector.write-buffer-size") .intType() .defaultValue(256) - .withDescription("write buffer size in MB for vector index"); + .withDescription("Write buffer size in MB for vector index"); private final int dimension; - private final String metric; + private final VectorMetric metric; private final int m; private final int efConstruction; private final int sizePerIndex; @@ -87,7 +87,7 @@ public int dimension() { return dimension; } - public String metric() { + public VectorMetric metric() { return metric; } diff --git a/paimon-vector/src/main/java/org/apache/paimon/vector/VectorMetric.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorMetric.java new file mode 100644 index 000000000000..0ae4eaa1def9 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorMetric.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.vector; + +import org.apache.lucene.index.VectorSimilarityFunction; + +/** Enumeration of supported vector similarity metrics. */ +public enum VectorMetric { + /** Cosine similarity metric. */ + COSINE(VectorSimilarityFunction.COSINE), + + /** Dot product similarity metric. */ + DOT_PRODUCT(VectorSimilarityFunction.DOT_PRODUCT), + + /** Euclidean distance metric. */ + EUCLIDEAN(VectorSimilarityFunction.EUCLIDEAN), + + /** Maximum inner product metric. */ + MAX_INNER_PRODUCT(VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT); + + private final VectorSimilarityFunction vectorSimilarityFunction; + + VectorMetric(VectorSimilarityFunction vectorSimilarityFunction) { + this.vectorSimilarityFunction = vectorSimilarityFunction; + } + + /** + * Get the corresponding Lucene vector similarity function. + * + * @return the Lucene VectorSimilarityFunction + */ + public VectorSimilarityFunction vectorSimilarityFunction() { + return vectorSimilarityFunction; + } +} diff --git a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java index d4a528d155aa..3d8167671c13 100644 --- a/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java @@ -121,13 +121,10 @@ public void testDifferentSimilarityFunctions() throws IOException { result.rowRange(), result.meta())); - VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); - - // Verify search works with this metric - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 3); - assertThat(searchResult).isNotNull(); - - reader.close(); + try (VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas)) { + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 3); + assertThat(searchResult).isNotNull(); + } } } @@ -162,13 +159,11 @@ public void testDifferentDimensions() throws IOException { result.rowRange(), result.meta())); - VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); - - // Verify search works with this dimension - GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); - assertThat(searchResult).isNotNull(); - - reader.close(); + try (VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas)) { + // Verify search works with this dimension + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); + assertThat(searchResult).isNotNull(); + } } } @@ -220,18 +215,17 @@ public void testFloatVectorIndexEndToEnd() throws IOException { result.meta())); } - VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas); + try (VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas)) { + GlobalIndexResult result = reader.search(vectors[0], 1); + assertThat(result.results().getLongCardinality()).isEqualTo(1); + assertThat(containsRowId(result, 0)).isTrue(); - GlobalIndexResult result = reader.search(vectors[0], 1); - assertThat(result.results().getLongCardinality()).isEqualTo(1); - assertThat(containsRowId(result, 0)).isTrue(); - - float[] queryVector = new float[] {0.85f, 0.15f}; - result = reader.search(queryVector, 2); - assertThat(result.results().getLongCardinality()).isEqualTo(2); - assertThat(containsRowId(result, 1)).isTrue(); - assertThat(containsRowId(result, 3)).isTrue(); - reader.close(); + float[] queryVector = new float[] {0.85f, 0.15f}; + result = reader.search(queryVector, 2); + assertThat(result.results().getLongCardinality()).isEqualTo(2); + assertThat(containsRowId(result, 1)).isTrue(); + assertThat(containsRowId(result, 3)).isTrue(); + } } private Options createDefaultOptions(int dimension) {