[core] Vector search add offset back (#6984)

Author: leaves12138

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<GlobalIndexResult> visitNotIn(FieldRef fieldRef, List<Object> li
 
     @Override
     public Optional<GlobalIndexResult> visitVectorSearch(VectorSearch vectorSearch) {
-        return applyOffset(wrapped.visitVectorSearch(vectorSearch));
+        return applyOffset(
+                wrapped.visitVectorSearch(vectorSearch.offsetRange(this.offset, this.to)));
     }
 
     private Optional<GlobalIndexResult> applyOffset(Optional<GlobalIndexResult> result) {

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<GlobalIndexResult> visit(GlobalIndexReader visitor) {
         return visitor.visitVectorSearch(this);
     }

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<Range> ranges = vectorSearch.includeRowIds().toRangeList();
+        assertThat(ranges.get(0)).isEqualTo(new Range(40L, 90L));
+    }
+}

paimon-core/src/main/java/org/apache/paimon/globalindex/RowRangeGlobalIndexScanner.java

@@ -137,7 +137,8 @@ private Collection<GlobalIndexReader> createReaders(
                     GlobalIndexReader innerReader =
                             new OffsetGlobalIndexReader(
                                     globalIndexer.createReader(indexFileReadWrite, globalMetas),
-                                    range.from);
+                                    range.from,
+                                    range.to);
                     unionReader.add(innerReader);
                 }