From 19232f9c50aa988e4c97a147f70f75feca2c097a Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 20:08:47 +0800 Subject: [PATCH] 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 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/query.py | 6 + python/python/lancedb/table.py | 29 +- python/python/tests/test_query.py | 15 + python/python/tests/test_remote_db.py | 35 +- python/src/query.rs | 3 + .../src/dataloader/permutation/reader.rs | 74 +- rust/lancedb/src/query.rs | 840 +++++++++++++++++- rust/lancedb/src/remote/table.rs | 162 +++- rust/lancedb/src/table.rs | 14 +- rust/lancedb/src/table/query.rs | 9 +- 11 files changed, 1157 insertions(+), 31 deletions(-) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 05ece3043..2b2691139 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/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]]] diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 451384ad1..c76e9e7db 100644 --- a/python/python/lancedb/query.py +++ b/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 diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 36c5b727d..287dff1f6 100644 --- a/python/python/lancedb/table.py +++ b/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] diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index ff62b2b51..6fbe0689b 100644 --- a/python/python/tests/test_query.py +++ b/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()) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index ab0df386d..01e2cc4c5 100644 --- a/python/python/tests/test_remote_db.py +++ b/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(): diff --git a/python/src/query.rs b/python/src/query.rs index 38153729f..ef71939f2 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -323,6 +323,7 @@ impl<'py> IntoPyObject<'py> for PyQueryVectors { pub struct PyQueryRequest { pub limit: Option, pub offset: Option, + pub take_offsets: Option>, pub filter: Option, pub full_text_search: Option>, pub select: PySelect, @@ -353,6 +354,7 @@ impl From 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 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), diff --git a/rust/lancedb/src/dataloader/permutation/reader.rs b/rust/lancedb/src/dataloader/permutation/reader.rs index 9757dc552..015c3a17d 100644 --- a/rust/lancedb/src/dataloader/permutation/reader.rs +++ b/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 = row_ids.iter().map(|id| lit(*id)).collect(); + let mut unique_row_ids = HashSet::with_capacity(num_rows); + let in_list: Vec = 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::() != num_rows { + if batches.iter().map(|b| b.num_rows()).sum::() != 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::()) + .into_mem_table("tbl", RowCount::from(5), BatchCount::from(1)) + .await; + let base_row_ids = collect_column::(&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::() + .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::() .values() .to_vec(); - assert_eq!(idx_values, vec![0, 2, 4, 6]); + assert_eq!(idx_values, vec![0, 2, 0, 4, 6]); } #[tokio::test] diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index cd346f42e..5b889cada 100644 --- a/rust/lancedb/src/query.rs +++ b/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, + /// 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>, + /// Apply filter to the returned rows. pub filter: Option, @@ -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 { + 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, + ordering_column: &str, + drop_ordering_column: bool, + preserve_order: bool, +) -> Result { + 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::() + .values() + .to_vec(), + DataType::Int64 => actual_offsets + .as_primitive::() + .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::>>()?, + 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::>(); + // 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 { + 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, + offsets: Vec, + occurrences: Arc>, + ordering_column: String, + drop_ordering_column: bool, + preserve_order: bool, + schema: SchemaRef, + properties: Arc, +} + +impl TakeRestoreExec { + fn try_new( + input: Arc, + offsets: Vec, + ordering_column: String, + drop_ordering_column: bool, + preserve_order: bool, + ) -> Result { + 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 { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn maintains_input_order(&self) -> Vec { + vec![!self.preserve_order] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + 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, + ) -> DataFusionResult { + 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> + Send>> = + if preserve_order { + let input_schema = input.schema(); + Box::pin(stream::once(async move { + let batches = input.try_collect::>().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, request: QueryRequest, + offsets: Option>, + 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, offsets: Vec) -> Self { - let in_list: Vec = offsets.iter().map(|o| lit(*o)).collect(); + let mut seen = HashSet::with_capacity(offsets.len()); + let in_list: Vec = 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::>(); + 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)> { + 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, + offsets: &[u64], + ordering_column: String, + drop_ordering_column: bool, + output_offset: usize, + output_limit: Option, + preserve_order: bool, + ) -> Result> { + let lookup = if preserve_order { + Arc::new(CoalescePartitionsExec::new(lookup)) as Arc + } else { + lookup + }; + let restored: Arc = 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, + preserve_order: bool, + ) -> String { + fn indent(plan: &str, spaces: usize) -> String { + let indentation = " ".repeat(spaces); + plan.lines() + .map(|line| format!("{indentation}{line}")) + .collect::>() + .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> { + 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> { + 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 { + 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 { + 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> { + 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 { + 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 { + 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 { + 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::>() + .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::>() + .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::() + .values() + .to_vec() + }) + .collect::>(); + 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::>() + .await + .unwrap(); + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 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, + ), + ( + "_rowoffset", + Arc::new(UInt64Array::from(vec![17, 5, 1])) as Arc, + ), + ]) + .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::() + .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::() + .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::>() + .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::() + .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::>() + .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::() + .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::>() + .await + .unwrap(); + let planned = concat_batches(&planned[0].schema(), &planned).unwrap(); + let mut ids = planned + .column_by_name("id") + .unwrap() + .as_primitive::() + .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(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index d372a6f56..daea028e9 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/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 BaseTable for RemoteTable { 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 BaseTable for RemoteTable { query: &AnyQuery, options: QueryExecutionOptions, ) -> Result> { + 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 BaseTable for RemoteTable { query: &AnyQuery, options: QueryExecutionOptions, ) -> Result { + 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 BaseTable for RemoteTable { } async fn explain_plan(&self, query: &AnyQuery, verbose: bool) -> Result { + 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 BaseTable for RemoteTable { query: &AnyQuery, options: QueryExecutionOptions, ) -> Result { + 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::>() + .await + .unwrap(); + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 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 = diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index efc705e3c..4602d35ed 100644 --- a/rust/lancedb/src/table.rs +++ b/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; + /// 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; @@ -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 /// ---------- diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 6f6bbf372..d611fe8a5 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -110,7 +110,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!( @@ -154,6 +154,13 @@ pub async fn create_plan( options: QueryExecutionOptions, ) -> Result> { 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),