diff --git a/paimon-common/src/main/java/org/apache/paimon/globalindex/OffsetGlobalIndexReader.java b/paimon-common/src/main/java/org/apache/paimon/globalindex/OffsetGlobalIndexReader.java index ecd5b3fd66f6..5c3afecc7c44 100644 --- a/paimon-common/src/main/java/org/apache/paimon/globalindex/OffsetGlobalIndexReader.java +++ b/paimon-common/src/main/java/org/apache/paimon/globalindex/OffsetGlobalIndexReader.java @@ -33,10 +33,12 @@ public class OffsetGlobalIndexReader implements GlobalIndexReader { private final GlobalIndexReader wrapped; private final long offset; + private final long to; - public OffsetGlobalIndexReader(GlobalIndexReader wrapped, long offset) { + public OffsetGlobalIndexReader(GlobalIndexReader wrapped, long offset, long to) { this.wrapped = wrapped; this.offset = offset; + this.to = to; } @Override @@ -111,7 +113,8 @@ public Optional visitNotIn(FieldRef fieldRef, List li @Override public Optional visitVectorSearch(VectorSearch vectorSearch) { - return applyOffset(wrapped.visitVectorSearch(vectorSearch)); + return applyOffset( + wrapped.visitVectorSearch(vectorSearch.offsetRange(this.offset, this.to))); } private Optional applyOffset(Optional result) { diff --git a/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java b/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java index a6042ec7e931..8ed13e043bea 100644 --- a/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java +++ b/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java @@ -20,6 +20,7 @@ import org.apache.paimon.globalindex.GlobalIndexReader; import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.utils.Range; import org.apache.paimon.utils.RoaringNavigableMap64; import javax.annotation.Nullable; @@ -76,6 +77,22 @@ public VectorSearch withIncludeRowIds(RoaringNavigableMap64 includeRowIds) { return this; } + public VectorSearch offsetRange(long from, long to) { + if (includeRowIds != null) { + RoaringNavigableMap64 range = new RoaringNavigableMap64(); + range.addRange(new Range(from, to)); + RoaringNavigableMap64 and64 = RoaringNavigableMap64.and(range, includeRowIds); + final RoaringNavigableMap64 roaringNavigableMap64Offset = new RoaringNavigableMap64(); + for (long rowId : and64) { + roaringNavigableMap64Offset.add(rowId - from); + } + VectorSearch target = new VectorSearch(vector, limit, fieldName); + target.withIncludeRowIds(roaringNavigableMap64Offset); + return target; + } + return this; + } + public Optional visit(GlobalIndexReader visitor) { return visitor.visitVectorSearch(this); } diff --git a/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java b/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java new file mode 100644 index 000000000000..735874ce848d --- /dev/null +++ b/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java @@ -0,0 +1,52 @@ +/* + * 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.predicate; + +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringNavigableMap64; + +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Test vector search. */ +public class VectorSearchTest { + + @Test + public void testVectorSearchOffset() { + 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} + }; + + VectorSearch vectorSearch = new VectorSearch(vectors[0], 1, "test"); + + RoaringNavigableMap64 includeRowIds = new RoaringNavigableMap64(); + includeRowIds.addRange(new Range(100L, 200L)); + vectorSearch.withIncludeRowIds(includeRowIds); + + vectorSearch = vectorSearch.offsetRange(60, 150); + + List ranges = vectorSearch.includeRowIds().toRangeList(); + assertThat(ranges.get(0)).isEqualTo(new Range(40L, 90L)); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/globalindex/RowRangeGlobalIndexScanner.java b/paimon-core/src/main/java/org/apache/paimon/globalindex/RowRangeGlobalIndexScanner.java index fc870570274f..f378669652e1 100644 --- a/paimon-core/src/main/java/org/apache/paimon/globalindex/RowRangeGlobalIndexScanner.java +++ b/paimon-core/src/main/java/org/apache/paimon/globalindex/RowRangeGlobalIndexScanner.java @@ -137,7 +137,8 @@ private Collection createReaders( GlobalIndexReader innerReader = new OffsetGlobalIndexReader( globalIndexer.createReader(indexFileReadWrite, globalMetas), - range.from); + range.from, + range.to); unionReader.add(innerReader); }