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 new file mode 100644 index 000000000000..cbf61f26bda2 --- /dev/null +++ b/.github/workflows/utitcase-paimon-vector.yml @@ -0,0 +1,63 @@ +################################################################################ +# 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: + - 'paimon-vector/**' + - 'pom.xml' + +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/.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/pom.xml b/paimon-vector/pom.xml new file mode 100644 index 000000000000..11c4d96afc76 --- /dev/null +++ b/paimon-vector/pom.xml @@ -0,0 +1,124 @@ + + + + 4.0.0 + + + paimon-parent + org.apache.paimon + 1.4-SNAPSHOT + + + paimon-vector + Paimon : Vector + + + 9.12.3 + 11 + + + + + org.apache.paimon + paimon-common + ${project.version} + + + + org.apache.lucene + lucene-core + ${lucene.version} + + + + org.apache.lucene + lucene-codecs + ${lucene.version} + + + + + + org.apache.paimon + paimon-common + ${project.version} + test-jar + test + + + + org.apache.paimon + paimon-test-utils + ${project.version} + 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/ByteVectorIndex.java b/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorIndex.java new file mode 100644 index 000000000000..b0b99d1bab7e --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/ByteVectorIndex.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.lucene.document.KnnByteVectorField; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.VectorSimilarityFunction; + +/** Vector index for byte vector. */ +public class ByteVectorIndex extends VectorIndex { + private final long rowId; + private final byte[] vector; + + public ByteVectorIndex(long rowId, byte[] vector) { + this.rowId = rowId; + this.vector = vector; + } + + @Override + 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 new file mode 100644 index 000000000000..3b2c9ae3c47b --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/FloatVectorIndex.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.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.VectorSimilarityFunction; + +/** Vector index for float vector. */ +public class FloatVectorIndex extends VectorIndex { + private final long rowId; + private final float[] vector; + + public FloatVectorIndex(long rowId, float[] vector) { + this.rowId = rowId; + this.vector = vector; + } + + @Override + public float[] vector() { + return vector; + } + + @Override + 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 new file mode 100644 index 000000000000..2f50081d9ab3 --- /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 new file mode 100644 index 000000000000..b7160bc97e6e --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexReader.java @@ -0,0 +1,341 @@ +/* + * 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.predicate.FieldRef; +import org.apache.paimon.utils.RoaringNavigableMap64; + +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.KnnByteVectorQuery; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.IndexOutput; + +import java.io.IOException; +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; + +/** + * Vector global index reader using Apache Lucene 9.x. + * + *

This implementation uses Lucene's native KnnFloatVectorQuery with HNSW graph for efficient + * approximate nearest neighbor search. + */ +public class VectorGlobalIndexReader implements GlobalIndexReader { + + private static final int BUFFER_SIZE = 8192; // 8KB buffer for streaming + + private final List searchers; + private final List directories; + + public VectorGlobalIndexReader(GlobalIndexFileReader fileReader, List files) + throws IOException { + this.searchers = new ArrayList<>(); + this.directories = 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) { + KnnFloatVectorQuery knnQuery = new KnnFloatVectorQuery(VectorIndex.VECTOR_FIELD, query, k); + return search(knnQuery, k); + } + + public GlobalIndexResult search(byte[] query, int k) { + KnnByteVectorQuery knnQuery = new KnnByteVectorQuery(VectorIndex.VECTOR_FIELD, query, k); + return search(knnQuery, k); + } + + @Override + public void close() throws IOException { + Throwable firstException = null; + + // Close readers + for (IndexSearcher searcher : searchers) { + try { + searcher.getIndexReader().close(); + } catch (Throwable t) { + if (firstException == null) { + firstException = t; + } else { + firstException.addSuppressed(t); + } + } + } + searchers.clear(); + + // Close directories + for (IndexMMapDirectory directory : directories) { + try { + directory.close(); + } 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) { + PriorityQueue topK = + new PriorityQueue<>(Comparator.comparingDouble(sr -> sr.score)); + for (IndexSearcher searcher : searchers) { + try { + TopDocs topDocs = searcher.search(query, k); + StoredFields storedFields = searcher.storedFields(); + Set fieldsToLoad = Set.of(VectorIndex.ROW_ID_FIELD); + for (org.apache.lucene.search.ScoreDoc scoreDoc : topDocs.scoreDocs) { + Document doc = storedFields.document(scoreDoc.doc, fieldsToLoad); + long rowId = doc.getField(VectorIndex.ROW_ID_FIELD).numericValue().longValue(); + 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) { + try (SeekableInputStream in = fileReader.getInputStream(meta.fileName())) { + IndexMMapDirectory directory = null; + IndexReader reader = null; + boolean success = false; + try { + directory = deserializeDirectory(in); + reader = DirectoryReader.open(directory.directory()); + IndexSearcher searcher = new IndexSearcher(reader); + directories.add(directory); + searchers.add(searcher); + success = true; + } finally { + if (!success) { + if (reader != null) { + try { + reader.close(); + } catch (IOException e) { + } + } + if (directory != null) { + try { + directory.close(); + } catch (Exception e) { + throw new IOException("Failed to close directory", e); + } + } + } + } + } + } + } + + private IndexMMapDirectory deserializeDirectory(SeekableInputStream in) throws IOException { + IndexMMapDirectory indexMMapDirectory = new IndexMMapDirectory(); + 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); + } + } + } + + 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; + } + } + + // Implementation of FunctionVisitor methods + @Override + public GlobalIndexResult visitIsNotNull(FieldRef fieldRef) { + throw new UnsupportedOperationException( + "Vector index does not support isNotNull predicate"); + } + + @Override + public GlobalIndexResult visitIsNull(FieldRef fieldRef) { + throw new UnsupportedOperationException("Vector index does not support isNull predicate"); + } + + @Override + public GlobalIndexResult visitStartsWith(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException( + "Vector index does not support startsWith predicate"); + } + + @Override + public GlobalIndexResult visitEndsWith(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support endsWith predicate"); + } + + @Override + public GlobalIndexResult visitContains(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support contains predicate"); + } + + @Override + public GlobalIndexResult visitLessThan(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support lessThan predicate"); + } + + @Override + public GlobalIndexResult visitGreaterOrEqual(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException( + "Vector index does not support greaterOrEqual predicate"); + } + + @Override + public GlobalIndexResult visitNotEqual(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support notEqual predicate"); + } + + @Override + public GlobalIndexResult visitLessOrEqual(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException( + "Vector index does not support lessOrEqual predicate"); + } + + @Override + public GlobalIndexResult visitEqual(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException("Vector index does not support equal predicate"); + } + + @Override + public GlobalIndexResult visitGreaterThan(FieldRef fieldRef, Object literal) { + throw new UnsupportedOperationException( + "Vector index does not support greaterThan predicate"); + } + + @Override + public GlobalIndexResult visitIn(FieldRef fieldRef, List literals) { + throw new UnsupportedOperationException("Vector index does not support in predicate"); + } + + @Override + public GlobalIndexResult visitNotIn(FieldRef fieldRef, List literals) { + throw new UnsupportedOperationException("Vector index does not support notIn predicate"); + } +} 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..975db6c0e9e0 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorGlobalIndexWriter.java @@ -0,0 +1,199 @@ +/* + * 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.Range; + +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; +import org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat; +import org.apache.lucene.document.Document; +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.IOContext; + +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; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** + * Vector global index writer using Apache Lucene 9.x. + * + *

This implementation uses Lucene's native KnnFloatVectorField with HNSW algorithm for efficient + * approximate nearest neighbor search. + */ +public class VectorGlobalIndexWriter implements GlobalIndexWriter { + + private final GlobalIndexFileWriter fileWriter; + private final VectorIndexOptions vectorOptions; + private final VectorSimilarityFunction similarityFunction; + private final int sizePerIndex; + + private final List vectorIndices; + + 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.vectorIndices = new ArrayList<>(); + this.vectorOptions = new VectorIndexOptions(options); + this.similarityFunction = vectorOptions.metric().vectorSimilarityFunction(); + this.sizePerIndex = vectorOptions.sizePerIndex(); + } + + @Override + public void write(Object key) { + VectorIndex index; + if (key instanceof FloatVectorIndex) { + index = (FloatVectorIndex) key; + } else if (key instanceof ByteVectorIndex) { + 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 (vectorIndices.isEmpty()) { + return new ArrayList<>(); + } + + List results = new ArrayList<>(); + + // Split vectors into batches if size exceeds sizePerIndex + 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 batchVectorIndices = vectorIndices.subList(startIdx, endIdx); + + String fileName = fileWriter.newFileName(VectorGlobalIndexerFactory.IDENTIFIER); + try (OutputStream out = fileWriter.newOutputStream(fileName)) { + buildIndex( + batchVectorIndices, + this.vectorOptions.m(), + this.vectorOptions.efConstruction(), + this.vectorOptions.writeBufferSize(), + out); + } + long minRowIdInBatch = batchVectorIndices.get(0).rowId(); + long maxRowIdInBatch = + batchVectorIndices.get(batchVectorIndices.size() - 1).rowId(); + results.add( + ResultEntry.of( + fileName, null, new Range(minRowIdInBatch, maxRowIdInBatch))); + } + + return results; + } catch (IOException e) { + throw new RuntimeException("Failed to write vector global index", e); + } + } + + 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()) { + try (IndexWriter writer = new IndexWriter(indexMMapDirectory.directory(), config)) { + for (VectorIndex vectorIndex : batchVectors) { + Document doc = new Document(); + doc.add(vectorIndex.indexableField(similarityFunction)); + doc.add(vectorIndex.rowIdStoredField()); + writer.addDocument(doc); + } + writer.commit(); + } + + serializeDirectory(indexMMapDirectory.directory(), out); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static IndexWriterConfig getIndexWriterConfig( + int m, int efConstruction, int writeBufferSize) { + IndexWriterConfig config = new IndexWriterConfig(); + config.setRAMBufferSizeMB(writeBufferSize); + config.setCodec( + new Lucene912Codec(Lucene912Codec.Mode.BEST_SPEED) { + @Override + public KnnVectorsFormat getKnnVectorsFormatForField(String field) { + return new Lucene99HnswScalarQuantizedVectorsFormat(m, efConstruction); + } + }); + return config; + } + + private void serializeDirectory(Directory directory, OutputStream out) throws IOException { + String[] files = directory.listAll(); + out.write(intToBytes(files.length)); + + for (String fileName : files) { + byte[] nameBytes = fileName.getBytes(StandardCharsets.UTF_8); + out.write(intToBytes(nameBytes.length)); + out.write(nameBytes); + long fileLength = directory.fileLength(fileName); + out.write(ByteBuffer.allocate(8).putLong(fileLength).array()); + + 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); + out.write(buffer, 0, toRead); + remaining -= toRead; + } + } + } + } + + private byte[] intToBytes(int value) { + return ByteBuffer.allocate(4).putInt(value).array(); + } +} 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..2b10306a9cd8 --- /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. */ +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); + } +} 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..b76864f30f6c --- /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. */ +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/java/org/apache/paimon/vector/VectorIndex.java b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndex.java new file mode 100644 index 000000000000..0c9aa729ac48 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndex.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.document.StoredField; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.VectorSimilarityFunction; + +/** Vector index interface. */ +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 new file mode 100644 index 000000000000..a91b293aa7c2 --- /dev/null +++ b/paimon-vector/src/main/java/org/apache/paimon/vector/VectorIndexOptions.java @@ -0,0 +1,109 @@ +/* + * 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") + .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"); + + 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"); + + 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 VectorMetric 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() { + return dimension; + } + + public VectorMetric metric() { + return metric; + } + + public int m() { + return m; + } + + public int efConstruction() { + return efConstruction; + } + + public int sizePerIndex() { + return sizePerIndex; + } + + public int writeBufferSize() { + return writeBufferSize; + } +} 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/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/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..3d8167671c13 --- /dev/null +++ b/paimon-vector/src/test/java/org/apache/paimon/vector/VectorGlobalIndexTest.java @@ -0,0 +1,269 @@ +/* + * 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.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}. */ +public class VectorGlobalIndexTest { + + @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 testDifferentSimilarityFunctions() throws IOException { + int dimension = 32; + int numVectors = 20; + + String[] metrics = {"COSINE", "DOT_PRODUCT", "EUCLIDEAN"}; + + for (String metric : metrics) { + Options options = createDefaultOptions(dimension); + options.setString("vector.metric", metric); + + 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 FloatVectorIndex(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())); + + try (VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas)) { + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 3); + assertThat(searchResult).isNotNull(); + } + } + } + + @Test + public void testDifferentDimensions() throws IOException { + int[] dimensions = {8, 32, 128, 256}; + + for (int dimension : dimensions) { + Options options = createDefaultOptions(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 FloatVectorIndex(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())); + + try (VectorGlobalIndexReader reader = new VectorGlobalIndexReader(fileReader, metas)) { + // Verify search works with this dimension + GlobalIndexResult searchResult = reader.search(testVectors.get(0), 5); + assertThat(searchResult).isNotNull(); + } + } + } + + @Test + public void testDimensionMismatch() throws IOException { + Options options = createDefaultOptions(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 FloatVectorIndex(0, wrongDimVector))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("dimension mismatch"); + } + + @Test + public void testFloatVectorIndexEndToEnd() throws IOException { + int dimension = 2; + Options options = createDefaultOptions(dimension); + options.setInteger("vector.size-per-index", 3); + + 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); + for (int i = 0; i < vectors.length; i++) { + writer.write(new FloatVectorIndex(i, vectors[i])); + } + + List results = writer.finish(); + assertThat(results).hasSize(2); + + GlobalIndexFileReader fileReader = createFileReader(indexPath); + List metas = new ArrayList<>(); + for (GlobalIndexWriter.ResultEntry result : results) { + metas.add( + new GlobalIndexIOMeta( + result.fileName(), + fileIO.getFileSize(new Path(indexPath, result.fileName())), + result.rowRange(), + result.meta())); + } + + 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(); + + 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) { + Options options = new Options(); + 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); + 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 + 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 boolean containsRowId(GlobalIndexResult result, long rowId) { + List resultIds = new ArrayList<>(); + result.results().iterator().forEachRemaining(resultIds::add); + return resultIds.contains(rowId); + } +} diff --git a/pom.xml b/pom.xml index 6a08bbea9be2..612d9106aa07 100644 --- a/pom.xml +++ b/pom.xml @@ -510,6 +510,16 @@ under the License. 11 + + + paimon-vector + + paimon-vector + + + 11 + +