fix: preserve duplicate take offsets (#4024)

## Summary
- preserve repeated table offsets without adding a public ordering
guarantee
- retain exact requested ordering in identity and persisted permutations
- cover local, projected, multi-batch, and mocked-remote query paths

## Root cause
Take queries lowered offsets to a set-like IN predicate and discarded
repeated occurrences. Persisted permutation loading also compared the
distinct base-table result count with the requested occurrence count,
rejecting repeated row IDs before its existing reordering step could
expand them.

## Fix
The shared take-query path now deduplicates the predicate for efficient
lookup, requests row-offset metadata internally, and expands each
matching row to the requested multiplicity in backend result order. An
internal opt-in keeps exact requested order for identity
PermutationReader reads, while persisted permutations continue using
their existing ordering map.

## Validation
- cargo test --quiet --features remote --tests
- cargo check --quiet --features remote --tests --examples
- cargo clippy --quiet --features remote --tests --examples
- targeted Python local and mocked-remote regression tests
- exact issue reproduction

Fixes #2820

<!-- lance-gatekeeper-fix:v1 agent=75acf840afa6f4be4bff98b567b504bd
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>

python/python/lancedb/_lancedb.pyi

@@ -607,6 +607,7 @@ class FullTextQuery:
 class PyQueryRequest:
     limit: Optional[int]
     offset: Optional[int]
+    take_offsets: Optional[List[int]]
     filter: Optional[Union[str, bytes]]
     full_text_search: Optional[FullTextQuery]
     select: Optional[Union[str, List[str]]]

python/python/lancedb/query.py

@@ -109,6 +109,7 @@ def _query_is_plain_scan(query: Query) -> bool:
     return (
         query.vector is None
         and query.full_text_query is None
+        and query.take_offsets is None
         and not query.postfilter
         and not query.order_by
     )
@@ -804,6 +805,10 @@ class Query(pydantic.BaseModel):
     # offset to start fetching results from
     offset: Optional[int] = None
 
+    # Dataset offsets whose duplicate occurrences must be restored after lookup.
+    # This is populated when a take query is converted to this serializable form.
+    take_offsets: Optional[List[int]] = None
+
     # if true, will only search the indexed data
     fast_search: Optional[bool] = None
 
@@ -825,6 +830,7 @@ class Query(pydantic.BaseModel):
         query = cls()
         query.limit = req.limit
         query.offset = req.offset
+        query.take_offsets = req.take_offsets
         query.filter = req.filter
         query.full_text_query = req.full_text_search
         query.columns = req.select

python/python/lancedb/table.py

@@ -1678,9 +1678,9 @@ class Table(ABC):
         Offsets are mostly useful for sampling as the set of all valid offsets is easily
         known in advance to be [0, len(table)).
 
-        No guarantees are made regarding the order in which results are returned.  If
-        you desire an output order that matches the order of the given offsets, you will
-        need to add the row offset column to the output and align it yourself.
+        No guarantees are made regarding the order in which results are returned.
+        Repeated offsets produce repeated rows, which makes this method suitable for
+        sampling with replacement.
 
         Parameters
         ----------
@@ -4090,6 +4090,7 @@ class LanceTable(Table):
             )
             and not self._route_pushdown_to_rust
             and self.current_branch() is None
+            and query.take_offsets is None
         ):
             from lancedb.namespace import _execute_server_side_query
 
@@ -5983,7 +5984,23 @@ class AsyncTable:
 
     def _sync_query_to_async(
         self, query: Query
-    ) -> AsyncHybridQuery | AsyncFTSQuery | AsyncVectorQuery | AsyncQuery:
+    ) -> (
+        AsyncHybridQuery
+        | AsyncFTSQuery
+        | AsyncVectorQuery
+        | AsyncQuery
+        | AsyncTakeQuery
+    ):
+        if query.take_offsets is not None:
+            take_query = self.take_offsets(query.take_offsets)
+            if query.columns:
+                take_query = take_query.select(query.columns)
+            if query.use_lsm is not None:
+                take_query = take_query.use_lsm(query.use_lsm)
+            if query.with_row_id:
+                take_query = take_query.with_row_id()
+            return take_query
+
         async_query = self.query()
         if query.limit is not None:
             async_query = async_query.limit(query.limit)
@@ -6048,6 +6065,7 @@ class AsyncTable:
                 self._namespace_client, self._pushdown_operations
             )
             and not self._route_pushdown_to_rust
+            and query.take_offsets is None
         ):
             from lancedb.namespace import _execute_server_side_query
 
@@ -6545,6 +6563,9 @@ class AsyncTable:
         Offsets are mostly useful for sampling as the set of all valid offsets is easily
         known in advance to be [0, len(table)).
 
+        No guarantees are made regarding the order in which results are returned.
+        Repeated offsets produce repeated rows.
+
         Parameters
         ----------
         offsets: list[int]

python/python/tests/test_query.py

@@ -1923,6 +1923,21 @@ def test_take_queries(tmp_path):
         17,
     ]
 
+    # Duplicate offsets are occurrences, not set members. Ordering is unspecified.
+    assert sorted(table.take_offsets([5, 2, 5, 17]).to_pandas()["idx"].to_list()) == [
+        2,
+        5,
+        5,
+        17,
+    ]
+
+    # Converting a take builder to its serializable query representation must
+    # retain occurrence metadata and execute with the same multiplicity.
+    query = table.take_offsets([5, 2, 5, 17]).select(["idx"]).to_query_object()
+    assert query.take_offsets == [5, 2, 5, 17]
+    converted = table._execute_query(query).read_all()
+    assert sorted(converted["idx"].to_pylist()) == [2, 5, 5, 17]
+
     # Take by row id
     assert list(
         sorted(table.take_row_ids([5, 2, 17]).to_pandas()["idx"].to_list())

python/python/tests/test_remote_db.py

@@ -479,24 +479,49 @@ def test_remote_permutation_is_picklable():
                 match = re.search(
                     r"_rowoffset\s+in\s+\((.*?)\)", body["filter"], re.IGNORECASE
                 )
-                offsets = [int(o.strip()) for o in match.group(1).split(",")]
+                offsets = list(
+                    dict.fromkeys(int(o.strip()) for o in match.group(1).split(","))
+                )
             else:
                 offsets = list(range(len(rows)))
-            table = pa.table({"a": [rows[offset] for offset in offsets]})
+            columns = body.get("columns") or ["a"]
+            table = pa.table(
+                {
+                    column: (
+                        [rows[offset] for offset in offsets]
+                        if column == "a"
+                        else offsets
+                    )
+                    for column in columns
+                }
+            )
 
             request.send_response(200)
             request.send_header("Content-Type", "application/vnd.apache.arrow.file")
             request.end_headers()
             with pa.ipc.new_file(request.wfile, schema=table.schema) as writer:
-                writer.write_table(table)
+                writer.write_table(table, max_chunksize=2)
         else:
             request.send_response(404)
             request.end_headers()
 
     with mock_lancedb_connection(handler) as db:
-        permutation = Permutation.identity(db.open_table("test"))
+        table = db.open_table("test")
+        assert table.take_offsets([0, 2, 0, 4]).to_list() == [
+            {"a": 0},
+            {"a": 0},
+            {"a": 2},
+            {"a": 4},
+        ]
+
+        permutation = Permutation.identity(table)
         restored = pickle.loads(pickle.dumps(permutation))
-        assert restored.__getitems__([0, 2, 4]) == [{"a": 0}, {"a": 2}, {"a": 4}]
+        assert restored.__getitems__([0, 2, 0, 4]) == [
+            {"a": 0},
+            {"a": 2},
+            {"a": 0},
+            {"a": 4},
+        ]
 
 
 def test_create_table_exist_ok():

python/src/query.rs

@@ -323,6 +323,7 @@ impl<'py> IntoPyObject<'py> for PyQueryVectors {
 pub struct PyQueryRequest {
     pub limit: Option<usize>,
     pub offset: Option<usize>,
+    pub take_offsets: Option<Vec<u64>>,
     pub filter: Option<PyQueryFilter>,
     pub full_text_search: Option<PyLanceDB<FtsQuery>>,
     pub select: PySelect,
@@ -353,6 +354,7 @@ impl From<AnyQuery> for PyQueryRequest {
             AnyQuery::Query(query_request) => Self {
                 limit: query_request.limit,
                 offset: query_request.offset,
+                take_offsets: query_request.take_offsets,
                 filter: query_request.filter.map(PyQueryFilter),
                 full_text_search: query_request
                     .full_text_search
@@ -381,6 +383,7 @@ impl From<AnyQuery> for PyQueryRequest {
             AnyQuery::VectorQuery(vector_query) => Self {
                 limit: vector_query.base.limit,
                 offset: vector_query.base.offset,
+                take_offsets: vector_query.base.take_offsets,
                 filter: vector_query.base.filter.map(PyQueryFilter),
                 full_text_search: None,
                 select_source_columns: PySelect::source_columns(&vector_query.base.select),