Merge remote-tracking branch 'refs/remotes/origin/main' into gatekeeper/fix-2820-1

# Conflicts:
#	rust/lancedb/src/table/query.rs

python/python/lancedb/_lancedb.pyi

@@ -605,6 +605,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
     )
@@ -798,6 +799,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
 
@@ -819,6 +824,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

@@ -1504,9 +1504,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
         ----------
@@ -3908,6 +3908,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
 
@@ -5799,7 +5800,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)
@@ -5864,6 +5881,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
 
@@ -6354,6 +6372,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

@@ -1908,6 +1908,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

@@ -322,6 +322,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,
@@ -351,6 +352,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
@@ -378,6 +380,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),

rust/lancedb/src/dataloader/permutation/reader.rs

@@ -31,7 +31,7 @@ use lance::io::RecordBatchStream;
 use lance_arrow::RecordBatchExt;
 use lance_core::ROW_ID;
 use lance_core::error::LanceOptionExt;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 /// Reads a permutation of a source table based on row IDs stored in a separate table
@@ -234,7 +234,14 @@ impl PermutationReader {
             .expect_ok()?
             .values();
 
-        let in_list: Vec<Expr> = row_ids.iter().map(|id| lit(*id)).collect();
+        let mut unique_row_ids = HashSet::with_capacity(num_rows);
+        let in_list: Vec<Expr> = row_ids
+            .iter()
+            .copied()
+            .filter(|row_id| unique_row_ids.insert(*row_id))
+            .map(lit)
+            .collect();
+        let num_unique_row_ids = unique_row_ids.len();
 
         let base_query = QueryRequest {
             filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))),
@@ -247,7 +254,7 @@ impl PermutationReader {
             .query(
                 &AnyQuery::Query(base_query),
                 QueryExecutionOptions {
-                    max_batch_length: num_rows as u32,
+                    max_batch_length: num_unique_row_ids as u32,
                     ..Default::default()
                 },
             )
@@ -262,9 +269,9 @@ impl PermutationReader {
             });
         }
 
-        if batches.iter().map(|b| b.num_rows()).sum::<usize>() != num_rows {
+        if batches.iter().map(|b| b.num_rows()).sum::<usize>() != num_unique_row_ids {
             return Err(Error::InvalidInput {
-                message: "Base table returned different number of rows than the number of row IDs"
+                message: "Base table returned a different number of rows than the number of unique row IDs"
                     .to_string(),
             });
         }
@@ -504,6 +511,7 @@ impl PermutationReader {
             let table = Table::from(self.base_table.clone());
             let batches = table
                 .take_offsets(offsets.to_vec())
+                .preserve_order()
                 .select(selection.clone())
                 .execute()
                 .await?
@@ -803,10 +811,10 @@ mod tests {
         .unwrap();
 
         // Take offsets in reverse order and verify returned rows match that order
-        let offsets = vec![5, 3, 1, 0];
+        let offsets = vec![5, 3, 5, 1, 0];
         let batch = reader.take_offsets(&offsets, Select::All).await.unwrap();
 
-        assert_eq!(batch.num_rows(), 4);
+        assert_eq!(batch.num_rows(), 5);
 
         let idx_values = batch
             .column(0)
@@ -820,6 +828,52 @@ mod tests {
         assert_eq!(idx_values, expected);
     }
 
+    #[tokio::test]
+    async fn test_take_offsets_preserves_repeated_rows_in_permutation() {
+        let base_table = lance_datagen::gen_batch()
+            .col("idx", lance_datagen::array::step::<Int32Type>())
+            .into_mem_table("tbl", RowCount::from(5), BatchCount::from(1))
+            .await;
+        let base_row_ids = collect_column::<UInt64Type>(&base_table, "_rowid").await;
+        let permutation_row_ids = vec![
+            base_row_ids[3],
+            base_row_ids[1],
+            base_row_ids[3],
+            base_row_ids[2],
+        ];
+        let permutation_batch = RecordBatch::try_new(
+            Arc::new(Schema::new(vec![
+                Field::new("row_id", DataType::UInt64, false),
+                Field::new(SPLIT_ID_COLUMN, DataType::UInt64, false),
+            ])),
+            vec![
+                Arc::new(UInt64Array::from(permutation_row_ids)),
+                Arc::new(UInt64Array::from(vec![0; 4])),
+            ],
+        )
+        .unwrap();
+        let permutation_table = virtual_table("row_ids", &permutation_batch).await;
+        let reader = PermutationReader::try_from_tables(
+            base_table.base_table().clone(),
+            permutation_table.base_table().clone(),
+            0,
+        )
+        .await
+        .unwrap();
+
+        let batch = reader
+            .take_offsets(&[0, 1, 2, 3], Select::All)
+            .await
+            .unwrap();
+        let idx_values = batch
+            .column(0)
+            .as_primitive::<Int32Type>()
+            .values()
+            .to_vec();
+
+        assert_eq!(idx_values, vec![3, 1, 3, 2]);
+    }
+
     #[tokio::test]
     async fn test_take_offsets_with_column_selection() {
         let (base_table, row_ids_table, row_ids) = setup_permutation_tables(10).await;
@@ -883,17 +937,17 @@ mod tests {
             .unwrap();
 
         // With no permutation table, take_offsets uses the base table directly
-        let offsets = vec![0, 2, 4, 6];
+        let offsets = vec![0, 2, 0, 4, 6];
         let batch = reader.take_offsets(&offsets, Select::All).await.unwrap();
 
-        assert_eq!(batch.num_rows(), 4);
+        assert_eq!(batch.num_rows(), 5);
 
         let idx_values = batch
             .column(0)
             .as_primitive::<Int32Type>()
             .values()
             .to_vec();
-        assert_eq!(idx_values, vec![0, 2, 4, 6]);
+        assert_eq!(idx_values, vec![0, 2, 0, 4, 6]);
     }
 
     #[tokio::test]

rust/lancedb/src/query.rs

@@ -1,21 +1,37 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
+use std::collections::{HashMap, HashSet};
+use std::pin::Pin;
 use std::sync::Arc;
 use std::{future::Future, time::Duration};
 
 use arrow::compute::concat_batches;
-use arrow_array::{Array, Float16Array, Float32Array, Float64Array, RecordBatch, make_array};
+use arrow_array::{
+    Array, Float16Array, Float32Array, Float64Array, RecordBatch, UInt64Array,
+    cast::AsArray,
+    make_array,
+    types::{Int64Type, UInt64Type},
+};
 use arrow_schema::{DataType, SchemaRef};
+use datafusion_common::{DataFusionError, Result as DataFusionResult};
+use datafusion_execution::TaskContext;
 use datafusion_expr::{Expr, col, lit};
-use datafusion_physical_plan::ExecutionPlan;
-use futures::{FutureExt, TryFutureExt, TryStreamExt, stream, try_join};
+use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
+use datafusion_physical_plan::{
+    DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
+    coalesce_partitions::CoalescePartitionsExec,
+    execution_plan::{Boundedness, EmissionType},
+    limit::GlobalLimitExec,
+    stream::RecordBatchStreamAdapter,
+};
+use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt, stream, try_join};
 use half::f16;
 /// Re-export Lance ColumnOrdering type for use in query ordering
 pub use lance::dataset::scanner::ColumnOrdering;
 use lance::dataset::{ROW_ID, scanner::DatasetRecordBatchStream};
 use lance_arrow::RecordBatchExt;
-use lance_datafusion::exec::execute_plan;
+use lance_datafusion::exec::{execute_plan, format_plan as format_analyzed_plan};
 use lance_index::scalar::FullTextSearchQuery;
 use lance_index::scalar::inverted::SCORE_COL;
 use lance_index::vector::DIST_COL;
@@ -825,6 +841,14 @@ pub struct QueryRequest {
     /// Offset of the query.
     pub offset: Option<usize>,
 
+    /// Dataset offsets whose occurrence multiplicity must be restored after
+    /// executing the physical lookup represented by this request.
+    ///
+    /// This is client-side execution metadata used when a [`TakeQuery`] is
+    /// converted into a request. It is not sent to remote services.
+    #[doc(hidden)]
+    pub take_offsets: Option<Vec<u64>>,
+
     /// Apply filter to the returned rows.
     pub filter: Option<QueryFilter>,
 
@@ -893,6 +917,7 @@ impl Default for QueryRequest {
         Self {
             limit: None,
             offset: None,
+            take_offsets: None,
             filter: None,
             filter_error: None,
             full_text_search: None,
@@ -1529,6 +1554,302 @@ impl HasQuery for VectorQuery {
     }
 }
 
+fn take_occurrences(offsets: &[u64]) -> HashMap<u64, usize> {
+    let mut occurrences = HashMap::with_capacity(offsets.len());
+    for offset in offsets {
+        *occurrences.entry(*offset).or_insert(0) += 1;
+    }
+    occurrences
+}
+
+fn restore_take_batch_with_occurrences(
+    batch: RecordBatch,
+    offsets: &[u64],
+    occurrences: &HashMap<u64, usize>,
+    ordering_column: &str,
+    drop_ordering_column: bool,
+    preserve_order: bool,
+) -> Result<RecordBatch> {
+    let actual_offsets = batch
+        .column_by_name(ordering_column)
+        .ok_or_else(|| Error::Schema {
+            message: format!(
+                "take query result did not include ordering column '{ordering_column}'"
+            ),
+        })?;
+    let actual_offsets = match actual_offsets.data_type() {
+        DataType::UInt64 => actual_offsets
+            .as_primitive::<UInt64Type>()
+            .values()
+            .to_vec(),
+        DataType::Int64 => actual_offsets
+            .as_primitive::<Int64Type>()
+            .values()
+            .iter()
+            .map(|offset| {
+                u64::try_from(*offset).map_err(|_| Error::Schema {
+                    message: format!(
+                        "take query ordering column '{ordering_column}' contained a negative offset"
+                    ),
+                })
+            })
+            .collect::<Result<Vec<_>>>()?,
+        data_type => {
+            return Err(Error::Schema {
+                message: format!(
+                    "take query ordering column '{ordering_column}' had unsupported type {data_type}"
+                ),
+            });
+        }
+    };
+
+    let mut desired_order = Vec::with_capacity(offsets.len());
+    if preserve_order {
+        let ordering = actual_offsets
+            .iter()
+            .copied()
+            .enumerate()
+            .map(|(index, offset)| (offset, index as u64))
+            .collect::<HashMap<_, _>>();
+        // Missing offsets retain the filter-based behavior of returning no row.
+        desired_order.extend(
+            offsets
+                .iter()
+                .filter_map(|offset| ordering.get(offset).copied()),
+        );
+    } else {
+        // Public take queries do not guarantee output order. Preserve the lookup's
+        // existing order and only restore the multiplicity of each matching row.
+        for (index, offset) in actual_offsets.iter().enumerate() {
+            if let Some(count) = occurrences.get(offset) {
+                desired_order.extend(std::iter::repeat_n(index as u64, *count));
+            }
+        }
+    }
+
+    let mut ordered_batch = if desired_order.len() == batch.num_rows()
+        && desired_order
+            .iter()
+            .enumerate()
+            .all(|(index, desired)| *desired == index as u64)
+    {
+        batch
+    } else {
+        arrow_select::take::take_record_batch(&batch, &UInt64Array::from(desired_order))?
+    };
+
+    if drop_ordering_column {
+        ordered_batch = ordered_batch.drop_column(ordering_column)?;
+    }
+
+    Ok(ordered_batch)
+}
+
+#[cfg(test)]
+fn restore_take_batch(
+    batch: RecordBatch,
+    offsets: &[u64],
+    ordering_column: &str,
+    drop_ordering_column: bool,
+    preserve_order: bool,
+) -> Result<RecordBatch> {
+    restore_take_batch_with_occurrences(
+        batch,
+        offsets,
+        &take_occurrences(offsets),
+        ordering_column,
+        drop_ordering_column,
+        preserve_order,
+    )
+}
+
+/// Restores the logical offset occurrence sequence above the physical lookup plan.
+///
+/// The lookup plan returns each matching row at most once. For ordinary unordered
+/// takes this operator expands each input batch incrementally and preserves the
+/// lookup's partitioning. The explicitly ordered reader path collects one coalesced
+/// input before restoring requested order. Pagination must remain above this operator
+/// so it applies to occurrences.
+#[derive(Debug)]
+struct TakeRestoreExec {
+    input: Arc<dyn ExecutionPlan>,
+    offsets: Vec<u64>,
+    occurrences: Arc<HashMap<u64, usize>>,
+    ordering_column: String,
+    drop_ordering_column: bool,
+    preserve_order: bool,
+    schema: SchemaRef,
+    properties: Arc<PlanProperties>,
+}
+
+impl TakeRestoreExec {
+    fn try_new(
+        input: Arc<dyn ExecutionPlan>,
+        offsets: Vec<u64>,
+        ordering_column: String,
+        drop_ordering_column: bool,
+        preserve_order: bool,
+    ) -> Result<Self> {
+        let schema = if drop_ordering_column {
+            RecordBatch::new_empty(input.schema())
+                .drop_column(&ordering_column)?
+                .schema()
+        } else {
+            input.schema()
+        };
+        let partition_count = if preserve_order {
+            1
+        } else {
+            input.output_partitioning().partition_count()
+        };
+        let emission_type = if preserve_order {
+            EmissionType::Final
+        } else {
+            EmissionType::Incremental
+        };
+        let properties = Arc::new(PlanProperties::new(
+            EquivalenceProperties::new(schema.clone()),
+            Partitioning::UnknownPartitioning(partition_count),
+            emission_type,
+            Boundedness::Bounded,
+        ));
+
+        Ok(Self {
+            input,
+            occurrences: Arc::new(take_occurrences(&offsets)),
+            offsets,
+            ordering_column,
+            drop_ordering_column,
+            preserve_order,
+            schema,
+            properties,
+        })
+    }
+}
+
+impl DisplayAs for TakeRestoreExec {
+    fn fmt_as(
+        &self,
+        _display_type: DisplayFormatType,
+        formatter: &mut std::fmt::Formatter<'_>,
+    ) -> std::fmt::Result {
+        write!(
+            formatter,
+            "TakeRestoreExec: occurrences={}",
+            self.offsets.len()
+        )
+    }
+}
+
+impl ExecutionPlan for TakeRestoreExec {
+    fn name(&self) -> &str {
+        "TakeRestoreExec"
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.properties
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
+    }
+
+    fn maintains_input_order(&self) -> Vec<bool> {
+        vec![!self.preserve_order]
+    }
+
+    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+        vec![false]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
+        if children.len() != 1 {
+            return Err(DataFusionError::Internal(format!(
+                "TakeRestoreExec expected one child, got {}",
+                children.len()
+            )));
+        }
+        let child = children.into_iter().next().unwrap();
+        let plan = Self::try_new(
+            child,
+            self.offsets.clone(),
+            self.ordering_column.clone(),
+            self.drop_ordering_column,
+            self.preserve_order,
+        )
+        .map_err(|error| DataFusionError::External(Box::new(error)))?;
+        Ok(Arc::new(plan))
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> DataFusionResult<datafusion_physical_plan::SendableRecordBatchStream> {
+        let partition_count = self.input.output_partitioning().partition_count();
+        if partition >= partition_count || (self.preserve_order && partition != 0) {
+            return Err(DataFusionError::Internal(format!(
+                "TakeRestoreExec cannot execute partition {partition}; input has {partition_count} partitions"
+            )));
+        }
+
+        let input = self.input.execute(partition, context)?;
+        let output_schema = self.schema.clone();
+        let offsets = self.offsets.clone();
+        let occurrences = self.occurrences.clone();
+        let ordering_column = self.ordering_column.clone();
+        let drop_ordering_column = self.drop_ordering_column;
+        let preserve_order = self.preserve_order;
+        let stream: Pin<Box<dyn futures::Stream<Item = DataFusionResult<RecordBatch>> + Send>> =
+            if preserve_order {
+                let input_schema = input.schema();
+                Box::pin(stream::once(async move {
+                    let batches = input.try_collect::<Vec<_>>().await?;
+                    let batch = if batches.is_empty() {
+                        RecordBatch::new_empty(input_schema.clone())
+                    } else {
+                        concat_batches(&input_schema, &batches)?
+                    };
+                    restore_take_batch_with_occurrences(
+                        batch,
+                        &offsets,
+                        &occurrences,
+                        &ordering_column,
+                        drop_ordering_column,
+                        true,
+                    )
+                    .map_err(|error| DataFusionError::External(Box::new(error)))
+                }))
+            } else {
+                Box::pin(input.map(move |batch| {
+                    batch.and_then(|batch| {
+                        restore_take_batch_with_occurrences(
+                            batch,
+                            &offsets,
+                            &occurrences,
+                            &ordering_column,
+                            drop_ordering_column,
+                            false,
+                        )
+                        .map_err(|error| DataFusionError::External(Box::new(error)))
+                    })
+                }))
+            };
+
+        Ok(Box::pin(RecordBatchStreamAdapter::new(
+            output_schema,
+            stream,
+        )))
+    }
+
+    fn supports_limit_pushdown(&self) -> bool {
+        false
+    }
+}
+
 /// A builder for LanceDB take queries.
 ///
 /// See [`crate::Table::query`] for more details on queries
@@ -1545,6 +1866,8 @@ impl HasQuery for VectorQuery {
 pub struct TakeQuery {
     parent: Arc<dyn BaseTable>,
     request: QueryRequest,
+    offsets: Option<Vec<u64>>,
+    preserve_order: bool,
 }
 
 impl TakeQuery {
@@ -1552,15 +1875,24 @@ impl TakeQuery {
     ///
     /// See [`crate::Table::take_offsets`] for more details.
     pub fn from_offsets(parent: Arc<dyn BaseTable>, offsets: Vec<u64>) -> Self {
-        let in_list: Vec<Expr> = offsets.iter().map(|o| lit(*o)).collect();
+        let mut seen = HashSet::with_capacity(offsets.len());
+        let in_list: Vec<Expr> = offsets
+            .iter()
+            .copied()
+            .filter(|offset| seen.insert(*offset))
+            .map(lit)
+            .collect();
         Self {
             parent,
             request: QueryRequest {
                 filter: Some(QueryFilter::Datafusion(
                     col("_rowoffset").in_list(in_list, false),
                 )),
+                take_offsets: Some(offsets.clone()),
                 ..Default::default()
             },
+            offsets: Some(offsets),
+            preserve_order: false,
         }
     }
 
@@ -1575,9 +1907,181 @@ impl TakeQuery {
                 filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))),
                 ..Default::default()
             },
+            offsets: None,
+            preserve_order: false,
         }
     }
 
+    /// Preserve the requested offset order when restoring duplicate occurrences.
+    ///
+    /// This is reserved for readers whose API explicitly guarantees ordering.
+    pub(crate) fn preserve_order(mut self) -> Self {
+        debug_assert!(self.offsets.is_some());
+        self.preserve_order = true;
+        self
+    }
+
+    async fn request_with_row_offset(
+        parent: &dyn BaseTable,
+        request: &QueryRequest,
+    ) -> Result<(QueryRequest, String, bool)> {
+        const ROW_OFFSET: &str = "_rowoffset";
+        const INTERNAL_ROW_OFFSET: &str = "__lancedb_take_row_offset";
+
+        let mut request = request.clone();
+        // The physical lookup must not recursively restore occurrences. The
+        // wrapper above this request owns that logical operation.
+        request.take_offsets = None;
+        let (ordering_column, drop_ordering_column) = match &mut request.select {
+            Select::All => {
+                let mut columns = parent
+                    .schema()
+                    .await?
+                    .fields()
+                    .iter()
+                    .map(|field| field.name().clone())
+                    .collect::<Vec<_>>();
+                columns.push(ROW_OFFSET.to_string());
+                request.select = Select::Columns(columns);
+                (ROW_OFFSET.to_string(), true)
+            }
+            Select::Columns(columns) => {
+                if columns.iter().any(|column| column == ROW_OFFSET) {
+                    (ROW_OFFSET.to_string(), false)
+                } else {
+                    columns.push(ROW_OFFSET.to_string());
+                    (ROW_OFFSET.to_string(), true)
+                }
+            }
+            Select::Dynamic(columns) => {
+                let mut ordering_column = INTERNAL_ROW_OFFSET.to_string();
+                while columns.iter().any(|(name, _)| name == &ordering_column) {
+                    ordering_column.push('_');
+                }
+                columns.push((ordering_column.clone(), ROW_OFFSET.to_string()));
+                (ordering_column, true)
+            }
+            Select::Expr(columns) => {
+                let mut ordering_column = INTERNAL_ROW_OFFSET.to_string();
+                while columns.iter().any(|(name, _)| name == &ordering_column) {
+                    ordering_column.push('_');
+                }
+                columns.push((ordering_column.clone(), col(ROW_OFFSET)));
+                (ordering_column, true)
+            }
+        };
+
+        Ok((request, ordering_column, drop_ordering_column))
+    }
+
+    async fn prepare_offsets_lookup(
+        parent: &dyn BaseTable,
+        request: &QueryRequest,
+    ) -> Result<(QueryRequest, String, bool, usize, Option<usize>)> {
+        let (mut request, ordering_column, drop_ordering_column) =
+            Self::request_with_row_offset(parent, request).await?;
+        // The lookup operates on distinct physical rows. Pagination is a logical
+        // operation over occurrences and must be applied only after restoration.
+        let output_offset = request.offset.take().unwrap_or_default();
+        let output_limit = request.limit.take();
+
+        Ok((
+            request,
+            ordering_column,
+            drop_ordering_column,
+            output_offset,
+            output_limit,
+        ))
+    }
+
+    fn wrap_offsets_plan(
+        lookup: Arc<dyn ExecutionPlan>,
+        offsets: &[u64],
+        ordering_column: String,
+        drop_ordering_column: bool,
+        output_offset: usize,
+        output_limit: Option<usize>,
+        preserve_order: bool,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let lookup = if preserve_order {
+            Arc::new(CoalescePartitionsExec::new(lookup)) as Arc<dyn ExecutionPlan>
+        } else {
+            lookup
+        };
+        let restored: Arc<dyn ExecutionPlan> = Arc::new(TakeRestoreExec::try_new(
+            lookup,
+            offsets.to_vec(),
+            ordering_column,
+            drop_ordering_column,
+            preserve_order,
+        )?);
+
+        if output_offset > 0 || output_limit.is_some() {
+            Ok(Arc::new(GlobalLimitExec::new(
+                restored,
+                output_offset,
+                output_limit,
+            )))
+        } else {
+            Ok(restored)
+        }
+    }
+
+    fn wrap_offsets_explanation(
+        lookup: &str,
+        occurrence_count: usize,
+        output_offset: usize,
+        output_limit: Option<usize>,
+        preserve_order: bool,
+    ) -> String {
+        fn indent(plan: &str, spaces: usize) -> String {
+            let indentation = " ".repeat(spaces);
+            plan.lines()
+                .map(|line| format!("{indentation}{line}"))
+                .collect::<Vec<_>>()
+                .join("\n")
+        }
+
+        let restored = if preserve_order {
+            format!(
+                "TakeRestoreExec: occurrences={occurrence_count}\n  CoalescePartitionsExec\n{}",
+                indent(lookup, 4)
+            )
+        } else {
+            format!(
+                "TakeRestoreExec: occurrences={occurrence_count}\n{}",
+                indent(lookup, 2)
+            )
+        };
+
+        if output_offset > 0 || output_limit.is_some() {
+            let fetch = output_limit
+                .map(|limit| limit.to_string())
+                .unwrap_or_else(|| "None".to_string());
+            format!(
+                "GlobalLimitExec: skip={output_offset}, fetch={fetch}\n{}",
+                indent(&restored, 2)
+            )
+        } else {
+            restored
+        }
+    }
+
+    async fn create_offsets_plan(
+        &self,
+        offsets: &[u64],
+        options: QueryExecutionOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        create_take_offsets_plan(
+            self.parent.as_ref(),
+            &self.request,
+            offsets,
+            options,
+            self.preserve_order,
+        )
+        .await
+    }
+
     /// Convert the `TakeQuery` into a `QueryRequest`.
     pub fn into_request(self) -> QueryRequest {
         self.request
@@ -1622,6 +2126,63 @@ impl TakeQuery {
     }
 }
 
+pub(crate) async fn create_take_offsets_plan(
+    parent: &dyn BaseTable,
+    request: &QueryRequest,
+    offsets: &[u64],
+    options: QueryExecutionOptions,
+    preserve_order: bool,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let (request, ordering_column, drop_ordering_column, output_offset, output_limit) =
+        TakeQuery::prepare_offsets_lookup(parent, request).await?;
+    let lookup_options = if preserve_order {
+        options.without_output_batch_length_limit()
+    } else {
+        options
+    };
+    let lookup = parent
+        .create_plan(&AnyQuery::Query(request), lookup_options)
+        .await?;
+
+    TakeQuery::wrap_offsets_plan(
+        lookup,
+        offsets,
+        ordering_column,
+        drop_ordering_column,
+        output_offset,
+        output_limit,
+        preserve_order,
+    )
+}
+
+pub(crate) async fn explain_take_offsets_plan(
+    parent: &dyn BaseTable,
+    request: &QueryRequest,
+    offsets: &[u64],
+    verbose: bool,
+) -> Result<String> {
+    let (request, _, _, output_offset, output_limit) =
+        TakeQuery::prepare_offsets_lookup(parent, request).await?;
+    let lookup = parent
+        .explain_plan(&AnyQuery::Query(request), verbose)
+        .await?;
+    Ok(TakeQuery::wrap_offsets_explanation(
+        &lookup,
+        offsets.len(),
+        output_offset,
+        output_limit,
+        false,
+    ))
+}
+
+pub(crate) async fn prepare_take_offsets_request(
+    parent: &dyn BaseTable,
+    request: &QueryRequest,
+) -> Result<QueryRequest> {
+    let (request, _, _, _, _) = TakeQuery::prepare_offsets_lookup(parent, request).await?;
+    Ok(request)
+}
+
 impl HasQuery for TakeQuery {
     fn mut_query(&mut self) -> &mut QueryRequest {
         &mut self.request
@@ -1630,6 +2191,10 @@ impl HasQuery for TakeQuery {
 
 impl ExecutableQuery for TakeQuery {
     async fn create_plan(&self, options: QueryExecutionOptions) -> Result<Arc<dyn ExecutionPlan>> {
+        if let Some(offsets) = &self.offsets {
+            return self.create_offsets_plan(offsets, options).await;
+        }
+
         let req = AnyQuery::Query(self.request.clone());
         self.parent.clone().create_plan(&req, options).await
     }
@@ -1638,6 +2203,18 @@ impl ExecutableQuery for TakeQuery {
         &self,
         options: QueryExecutionOptions,
     ) -> Result<SendableRecordBatchStream> {
+        if self.offsets.is_some() {
+            let plan = self.create_plan(options.clone()).await?;
+            let inner = execute_plan(plan, Default::default())?;
+            let inner = MaxBatchLengthStream::new_boxed(inner, options.max_batch_length as usize);
+            let inner = if let Some(timeout) = options.timeout {
+                TimeoutStream::new_boxed(inner, timeout)
+            } else {
+                inner
+            };
+            return Ok(DatasetRecordBatchStream::new(inner).into());
+        }
+
         let query = AnyQuery::Query(self.request.clone());
         Ok(SendableRecordBatchStream::from(
             self.parent.clone().query(&query, options).await?,
@@ -1645,11 +2222,51 @@ impl ExecutableQuery for TakeQuery {
     }
 
     async fn explain_plan(&self, verbose: bool) -> Result<String> {
+        if let Some(offsets) = &self.offsets {
+            let (request, _, _, output_offset, output_limit) =
+                Self::prepare_offsets_lookup(self.parent.as_ref(), &self.request).await?;
+            // Ask the backend to explain only the distinct-row lookup. This keeps
+            // remote explanation non-executing while still showing the client-side
+            // operators that create_plan and execution place above that lookup.
+            let lookup = self
+                .parent
+                .explain_plan(&AnyQuery::Query(request), verbose)
+                .await?;
+            return Ok(Self::wrap_offsets_explanation(
+                &lookup,
+                offsets.len(),
+                output_offset,
+                output_limit,
+                self.preserve_order,
+            ));
+        }
+
         let query = AnyQuery::Query(self.request.clone());
         self.parent.explain_plan(&query, verbose).await
     }
 
     async fn analyze_plan_with_options(&self, options: QueryExecutionOptions) -> Result<String> {
+        if self.offsets.is_some() {
+            if self.parent.analyze_plan_is_remote() {
+                let (request, _, _, _, _) =
+                    Self::prepare_offsets_lookup(self.parent.as_ref(), &self.request).await?;
+                // Remote analysis is owned by the service. The current wire
+                // request represents only the distinct-row lookup, so return
+                // the service report unchanged instead of fabricating metrics
+                // for client-side restoration operators.
+                return self
+                    .parent
+                    .analyze_plan(&AnyQuery::Query(request), options)
+                    .await;
+            }
+
+            let plan = self.create_plan(options).await?;
+            execute_plan(plan.clone(), Default::default())?
+                .try_collect::<Vec<_>>()
+                .await?;
+            return Ok(format_analyzed_plan(plan));
+        }
+
         let query = AnyQuery::Query(self.request.clone());
         self.parent.analyze_plan(&query, options).await
     }
@@ -1670,6 +2287,7 @@ mod tests {
         StringArray, cast::AsArray, types::Float32Type,
     };
     use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
+    use datafusion_physical_plan::display::DisplayableExecutionPlan;
     use futures::{StreamExt, TryStreamExt};
     use lance_testing::datagen::{BatchGenerator, IncrementingInt32, RandomVector};
     use rand::seq::IndexedRandom;
@@ -2924,6 +3542,218 @@ mod tests {
         assert_eq!(results[0].num_columns(), 1);
     }
 
+    #[tokio::test]
+    async fn test_take_offsets_preserves_duplicate_multiplicity() {
+        let tmp_dir = tempdir().unwrap();
+        let table = make_test_table(&tmp_dir).await;
+
+        let results = table
+            .take_offsets(vec![5, 1, 5, 17])
+            .select(Select::Columns(vec!["id".to_string()]))
+            .execute_with_options(QueryExecutionOptions {
+                max_batch_length: 2,
+                ..Default::default()
+            })
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+
+        assert_eq!(results.len(), 2);
+        assert!(results.iter().all(|batch| batch.num_columns() == 1));
+        let mut ids = results
+            .iter()
+            .flat_map(|batch| {
+                batch
+                    .column_by_name("id")
+                    .unwrap()
+                    .as_primitive::<Int32Type>()
+                    .values()
+                    .to_vec()
+            })
+            .collect::<Vec<_>>();
+        ids.sort_unstable();
+        assert_eq!(ids, vec![1, 5, 5, 17]);
+    }
+
+    #[tokio::test]
+    async fn test_take_offsets_plan_is_incremental() {
+        let tmp_dir = tempdir().unwrap();
+        let table = make_test_table(&tmp_dir).await;
+
+        let plan = table
+            .take_offsets(vec![5, 1, 17])
+            .create_plan(QueryExecutionOptions {
+                max_batch_length: 1,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        assert_eq!(plan.properties().emission_type, EmissionType::Incremental);
+        let displayed = DisplayableExecutionPlan::new(plan.as_ref())
+            .indent(false)
+            .to_string();
+        assert!(displayed.contains("TakeRestoreExec"));
+        assert!(!displayed.contains("CoalescePartitionsExec"));
+    }
+
+    #[tokio::test]
+    async fn test_take_into_request_preserves_duplicate_multiplicity() {
+        let tmp_dir = tempdir().unwrap();
+        let table = make_test_table(&tmp_dir).await;
+        let request = table.take_offsets(vec![5, 5]).into_request();
+        assert_eq!(request.take_offsets, Some(vec![5, 5]));
+
+        let batches = table
+            .base_table()
+            .query(&AnyQuery::Query(request), QueryExecutionOptions::default())
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+
+        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
+    }
+
+    #[test]
+    fn test_restore_take_batch_only_reorders_when_requested() {
+        let batch = RecordBatch::try_from_iter([
+            (
+                "id",
+                Arc::new(Int32Array::from(vec![17, 5, 1])) as Arc<dyn Array>,
+            ),
+            (
+                "_rowoffset",
+                Arc::new(UInt64Array::from(vec![17, 5, 1])) as Arc<dyn Array>,
+            ),
+        ])
+        .unwrap();
+
+        let restored =
+            restore_take_batch(batch.clone(), &[5, 1, 5, 17], "_rowoffset", true, false).unwrap();
+        assert_eq!(
+            restored
+                .column_by_name("id")
+                .unwrap()
+                .as_primitive::<Int32Type>()
+                .values(),
+            &[17, 5, 5, 1]
+        );
+
+        let ordered = restore_take_batch(batch, &[5, 1, 5, 17], "_rowoffset", true, true).unwrap();
+        assert_eq!(
+            ordered
+                .column_by_name("id")
+                .unwrap()
+                .as_primitive::<Int32Type>()
+                .values(),
+            &[5, 1, 5, 17]
+        );
+    }
+
+    #[tokio::test]
+    async fn test_take_offsets_applies_pagination_after_restoration() {
+        let tmp_dir = tempdir().unwrap();
+        let table = make_test_table(&tmp_dir).await;
+
+        let limited = table
+            .take_offsets(vec![0, 1, 0, 2])
+            .select(Select::Columns(vec!["id".to_string()]))
+            .limit(3)
+            .execute()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let limited = concat_batches(&limited[0].schema(), &limited).unwrap();
+        assert_eq!(limited.num_rows(), 3);
+        assert!(
+            limited
+                .column_by_name("id")
+                .unwrap()
+                .as_primitive::<Int32Type>()
+                .values()
+                .iter()
+                .all(|id| [0, 1, 2].contains(id))
+        );
+
+        let offset = table
+            .take_offsets(vec![5, 1, 5, 17])
+            .select(Select::Columns(vec!["id".to_string()]))
+            .offset(1)
+            .execute()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let offset = concat_batches(&offset[0].schema(), &offset).unwrap();
+        assert_eq!(offset.num_rows(), 3);
+        assert!(
+            offset
+                .column_by_name("id")
+                .unwrap()
+                .as_primitive::<Int32Type>()
+                .values()
+                .iter()
+                .all(|id| [1, 5, 17].contains(id))
+        );
+    }
+
+    #[tokio::test]
+    async fn test_take_offsets_create_plan_restores_occurrences() {
+        let tmp_dir = tempdir().unwrap();
+        let table = make_test_table(&tmp_dir).await;
+        let take = table
+            .take_offsets(vec![5, 1, 5, 17])
+            .select(Select::Columns(vec!["id".to_string()]));
+
+        let plan = take
+            .create_plan(QueryExecutionOptions::default())
+            .await
+            .unwrap();
+        assert_eq!(plan.schema().fields().len(), 1);
+        assert_eq!(plan.schema().field(0).name(), "id");
+        let planned = execute_plan(plan, Default::default())
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let planned = concat_batches(&planned[0].schema(), &planned).unwrap();
+        let mut ids = planned
+            .column_by_name("id")
+            .unwrap()
+            .as_primitive::<Int32Type>()
+            .values()
+            .to_vec();
+        ids.sort_unstable();
+        assert_eq!(ids, vec![1, 5, 5, 17]);
+    }
+
+    #[tokio::test]
+    async fn test_take_offsets_plan_introspection_shows_restoration() {
+        let tmp_dir = tempdir().unwrap();
+        let table = make_test_table(&tmp_dir).await;
+        let take = table
+            .take_offsets(vec![0, 1, 0, 2])
+            .select(Select::Columns(vec!["id".to_string()]))
+            .limit(3);
+
+        let explained = take.explain_plan(false).await.unwrap();
+        assert!(explained.contains("GlobalLimitExec"));
+        assert!(explained.contains("TakeRestoreExec"));
+        assert!(!explained.contains("CoalescePartitionsExec"));
+
+        let analyzed = take.analyze_plan().await.unwrap();
+        assert!(analyzed.contains("GlobalLimitExec"));
+        assert!(analyzed.contains("TakeRestoreExec"));
+        assert!(!analyzed.contains("CoalescePartitionsExec"));
+    }
+
     #[tokio::test]
     async fn test_take_row_ids() {
         let tmp_dir = tempdir().unwrap();

rust/lancedb/src/remote/table.rs

@@ -40,8 +40,8 @@ use crate::table::{
 use crate::table::{AnyQuery, Filter, Predicate, PreprocessingOutput, TableStatistics};
 use crate::utils::background_cache::BackgroundCache;
 use crate::utils::{
-    resolve_arrow_field_path, resolve_arrow_fts_field_path, supported_btree_data_type,
-    supported_vector_data_type,
+    MaxBatchLengthStream, TimeoutStream, resolve_arrow_field_path, resolve_arrow_fts_field_path,
+    supported_btree_data_type, supported_vector_data_type,
 };
 use crate::{DistanceType, Error};
 use crate::{
@@ -2022,6 +2022,9 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
     fn as_any(&self) -> &dyn std::any::Any {
         self
     }
+    fn analyze_plan_is_remote(&self) -> bool {
+        true
+    }
     fn name(&self) -> &str {
         &self.name
     }
@@ -2594,6 +2597,13 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
         query: &AnyQuery,
         options: QueryExecutionOptions,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        if let AnyQuery::Query(request) = query
+            && let Some(offsets) = &request.take_offsets
+        {
+            return crate::query::create_take_offsets_plan(self, request, offsets, options, false)
+                .await;
+        }
+
         let streams = self.execute_query(query, &options).await?;
         if streams.len() == 1 {
             let stream = streams.into_iter().next().unwrap();
@@ -2612,6 +2622,27 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
         query: &AnyQuery,
         options: QueryExecutionOptions,
     ) -> Result<DatasetRecordBatchStream> {
+        if let AnyQuery::Query(request) = query
+            && let Some(offsets) = &request.take_offsets
+        {
+            let plan = crate::query::create_take_offsets_plan(
+                self,
+                request,
+                offsets,
+                options.clone(),
+                false,
+            )
+            .await?;
+            let inner = execute_plan(plan, Default::default())?;
+            let inner = MaxBatchLengthStream::new_boxed(inner, options.max_batch_length as usize);
+            let inner = if let Some(timeout) = options.timeout {
+                TimeoutStream::new_boxed(inner, timeout)
+            } else {
+                inner
+            };
+            return Ok(DatasetRecordBatchStream::new(inner));
+        }
+
         let streams = self.execute_query(query, &options).await?;
 
         if streams.len() == 1 {
@@ -2649,6 +2680,12 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
     }
 
     async fn explain_plan(&self, query: &AnyQuery, verbose: bool) -> Result<String> {
+        if let AnyQuery::Query(request) = query
+            && let Some(offsets) = &request.take_offsets
+        {
+            return crate::query::explain_take_offsets_plan(self, request, offsets, verbose).await;
+        }
+
         let base_request = self
             .client
             .post(&format!("/v1/table/{}/explain_plan/", self.identifier));
@@ -2701,6 +2738,17 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
         query: &AnyQuery,
         options: QueryExecutionOptions,
     ) -> Result<String> {
+        let prepared_query = if let AnyQuery::Query(request) = query
+            && request.take_offsets.is_some()
+        {
+            Some(AnyQuery::Query(
+                crate::query::prepare_take_offsets_request(self, request).await?,
+            ))
+        } else {
+            None
+        };
+        let query = prepared_query.as_ref().unwrap_or(query);
+
         let mut request = self
             .client
             .post(&format!("/v1/table/{}/analyze_plan/", self.identifier));
@@ -3690,7 +3738,7 @@ mod tests {
     };
     use arrow_schema::{DataType, Field, Schema};
     use chrono::{DateTime, Utc};
-    use futures::{StreamExt, TryFutureExt, future::BoxFuture};
+    use futures::{StreamExt, TryFutureExt, TryStreamExt, future::BoxFuture};
     use lance_index::scalar::inverted::{DocumentGranularity, query::MatchQuery};
     use lance_index::scalar::{FullTextSearchQuery, InvertedIndexParams};
     use reqwest::Body;
@@ -5611,6 +5659,114 @@ mod tests {
         assert_eq!(result, "analyzed plan");
     }
 
+    #[tokio::test]
+    async fn test_take_offsets_explain_plan_does_not_execute_query() {
+        let table = Table::new_with_handler("my_table", |request| {
+            assert_eq!(request.method(), "POST");
+            assert_eq!(request.url().path(), "/v1/table/my_table/explain_plan/");
+
+            http::Response::builder()
+                .status(200)
+                .body(r#""RemoteLookupExec""#)
+                .unwrap()
+        });
+
+        let explained = table
+            .take_offsets(vec![0, 1, 0, 2])
+            .select(crate::query::Select::columns(&["id"]))
+            .limit(3)
+            .explain_plan(false)
+            .await
+            .unwrap();
+
+        assert!(explained.contains("GlobalLimitExec"));
+        assert!(explained.contains("TakeRestoreExec"));
+        assert!(!explained.contains("CoalescePartitionsExec"));
+        assert!(explained.contains("RemoteLookupExec"));
+    }
+
+    #[tokio::test]
+    async fn test_converted_take_request_restores_remote_occurrences() {
+        let table = Table::new_with_handler("my_table", |request| {
+            assert_eq!(request.method(), "POST");
+            assert_eq!(request.url().path(), "/v1/table/my_table/query/");
+
+            let body: serde_json::Value =
+                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
+            assert_eq!(body["columns"], json!(["id", "_rowoffset"]));
+
+            let data = RecordBatch::try_new(
+                Arc::new(Schema::new(vec![
+                    Field::new("id", DataType::Int32, false),
+                    Field::new("_rowoffset", DataType::UInt64, false),
+                ])),
+                vec![
+                    Arc::new(Int32Array::from(vec![5])),
+                    Arc::new(arrow_array::UInt64Array::from(vec![5])),
+                ],
+            )
+            .unwrap();
+            http::Response::builder()
+                .status(200)
+                .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
+                .body(write_ipc_file(&data))
+                .unwrap()
+        });
+
+        let request = table
+            .take_offsets(vec![5, 5])
+            .select(crate::query::Select::columns(&["id"]))
+            .into_request();
+        let batches = table
+            .base_table()
+            .query(&AnyQuery::Query(request), QueryExecutionOptions::default())
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+
+        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
+        assert!(
+            batches
+                .iter()
+                .all(|batch| batch.schema().fields().len() == 1)
+        );
+    }
+
+    #[tokio::test]
+    async fn test_take_offsets_analyze_plan_delegates_to_remote() {
+        let table = Table::new_with_handler("my_table", |request| {
+            assert_eq!(request.method(), "POST");
+            assert_eq!(request.url().path(), "/v1/table/my_table/analyze_plan/");
+            assert_eq!(
+                request
+                    .url()
+                    .query_pairs()
+                    .find(|(key, _)| key == "distributed_metrics"),
+                Some(("distributed_metrics".into(), "per_worker".into()))
+            );
+
+            http::Response::builder()
+                .status(200)
+                .body(r#""Remote analyzed plan: worker metrics""#)
+                .unwrap()
+        });
+
+        let analyzed = table
+            .take_offsets(vec![0, 1, 0, 2])
+            .select(crate::query::Select::columns(&["id"]))
+            .limit(3)
+            .analyze_plan_with_options(QueryExecutionOptions {
+                analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics::PerWorker,
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        assert_eq!(analyzed, "Remote analyzed plan: worker metrics");
+    }
+
     #[tokio::test]
     async fn test_query_structured_fts() {
         let table =

rust/lancedb/src/table.rs

@@ -595,6 +595,14 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
         query: &AnyQuery,
         options: QueryExecutionOptions,
     ) -> Result<String>;
+    /// Whether [`BaseTable::analyze_plan`] is provided by a remote service.
+    ///
+    /// Client-side query wrappers use this to preserve backend metrics and
+    /// distributed-analysis options instead of replacing them with a local plan.
+    #[doc(hidden)]
+    fn analyze_plan_is_remote(&self) -> bool {
+        false
+    }
 
     /// Add new records to the table.
     async fn add(&self, add: AddDataBuilder) -> Result<AddResult>;
@@ -1652,9 +1660,9 @@ impl Table {
     /// Offsets are 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
     /// ----------

rust/lancedb/src/table/query.rs

@@ -109,7 +109,7 @@ fn requires_local_namespace_execution(query: &AnyQuery) -> bool {
     // pushing these down would silently ignore the user's setting. For use_lsm that
     // is worse than a tuning miss: MemWAL read routing lives only in `create_plan`,
     // so a pushed-down query would return stale base-only data with no error.
-    if query.base().use_lsm.is_some() {
+    if query.base().use_lsm.is_some() || query.base().take_offsets.is_some() {
         return true;
     }
     matches!(
@@ -153,6 +153,13 @@ pub async fn create_plan(
     options: QueryExecutionOptions,
 ) -> Result<Arc<dyn ExecutionPlan>> {
     let query = query.canonicalized()?;
+    if let AnyQuery::Query(request) = &query
+        && let Some(offsets) = &request.take_offsets
+    {
+        return crate::query::create_take_offsets_plan(table, request, offsets, options, false)
+            .await;
+    }
+
     let query = match query {
         AnyQuery::VectorQuery(query) => query,
         AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query),