fix: preserve streaming take conversions

python/python/lancedb/_lancedb.pyi

@@ -604,6 +604,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
     )
@@ -775,6 +776,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
 
@@ -796,6 +801,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

@@ -3873,6 +3873,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
 
@@ -5759,7 +5760,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)
@@ -5824,6 +5841,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
 

python/python/tests/test_query.py

@@ -1899,6 +1899,13 @@ def test_take_queries(tmp_path):
         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/src/query.rs

@@ -289,6 +289,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,
@@ -318,6 +319,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
@@ -345,6 +347,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: PySelect(vector_query.base.select),