diff --git a/docs/web/curation/selective-sampling.mdx b/docs/web/curation/selective-sampling.mdx index 7e6d941b4..f284abd22 100644 --- a/docs/web/curation/selective-sampling.mdx +++ b/docs/web/curation/selective-sampling.mdx @@ -57,16 +57,37 @@ let rows_by_offset = tbl.take_offsets(vec![0, 1, 2]).execute().await?; Both methods return a query builder; call `.to_arrow()`, `.to_list()` or -similar to execute it. Repeated *offsets* produce repeated rows, so -`take_offsets` also supports sampling with replacement -- `take_row_ids` -does not: a row ID repeated in the input list still returns that row only -once, since the lookup is a set membership test. Neither method guarantees -result order matches the order of the input list; for `take_row_ids`, pass -`with_row_id=True` to add `_rowid` to the output and align the results -yourself if you need to. `take_offsets` tracks `_rowoffset` internally to -restore repeated rows, but drops that column from the output again -afterward unless you explicitly `select` it -- add `_rowoffset` to your -selection when you need it to align a `take_offsets` result yourself. +similar to execute it. + + +The guarantees for row IDs below apply to SDK releases containing +[the ordering fix in #4826](https://github.com/lancedb/lancedb/pull/4826). +Earlier releases may collapse repeated IDs and return rows in a different +order. + + +`take_row_ids` returns matching rows in the order of the input IDs and +preserves repeated IDs in both LanceDB Embedded and LanceDB Enterprise. For example, +`take_row_ids([1005, 1001, 1005])` returns row 1005, row 1001, then row 1005 +again if both IDs exist. Missing or deleted IDs are omitted. This supports +sampling with replacement. In Rust, an explicit nonempty +`QueryBase::order_by` list overrides the requested order; limits and offsets +apply after restoring repeated rows. + +`take_offsets` also preserves repeated offsets, but does not guarantee +input order. It tracks `_rowoffset` internally to restore repeated rows +and drops that column from the output unless you explicitly `select` it. +Add `_rowoffset` to your selection when you need to align a `take_offsets` +result with the input offsets. + + +Ordered takes by row ID collect all matching batches, combine them into one +batch and reorder the rows before returning output batches. Even +`to_batches(max_batch_length=1)` only limits the returned batch size; it +does not bound working memory. For large requests or rows with large +column values, split the input IDs into smaller requests and use +`.select(...)` to fetch only the columns you need. + Offsets are mostly useful for sampling: the valid range is always `[0, len(table))`, so you can draw a random sample of offsets without diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index ed62b6b1d..f17856f3c 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -897,6 +897,7 @@ class PyQueryRequest: limit: Optional[int] offset: Optional[int] take_offsets: Optional[List[int]] + take_row_ids: 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 5185bf2a8..ce313c610 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -236,6 +236,7 @@ def _query_is_plain_scan(query: Query) -> bool: query.vector is None and query.full_text_query is None and query.take_offsets is None + and query.take_row_ids is None and not query.postfilter and not query.order_by ) @@ -1009,6 +1010,9 @@ class Query(pydantic.BaseModel): # This is populated when a take query is converted to this serializable form. take_offsets: Optional[List[int]] = None + # Row IDs whose requested order and duplicate occurrences must be restored. + take_row_ids: Optional[List[int]] = None + # if true, will only search the indexed data fast_search: Optional[bool] = None @@ -1034,6 +1038,7 @@ class Query(pydantic.BaseModel): query.limit = req.limit query.offset = req.offset query.take_offsets = req.take_offsets + query.take_row_ids = req.take_row_ids 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 4724776a1..59876a8c3 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -132,6 +132,7 @@ def _requires_local_namespace_execution(query: Query) -> bool: # The namespace QueryTable request does not support these query options. return ( query.take_offsets is not None + or query.take_row_ids is not None or query.search_effort is not None or query.approx_mode is not None ) @@ -2556,9 +2557,8 @@ class Table(ABC): Row ids are not stable and are relative to the current version of the table. They can change due to compaction and updates. - 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 ids, you will - need to add the row id column to the output and align it yourself. + Results follow the requested row-id order. Repeated IDs produce repeated + rows; missing or deleted IDs are omitted. This also applies to remote tables. Unlike offsets, row ids are not 0-indexed and no assumptions should be made about the possible range of row ids. In order to use this method you must @@ -7256,10 +7256,24 @@ class AsyncTable: | 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.take_offsets is not None or query.take_row_ids is not None: + if query.take_offsets is not None and query.take_row_ids is not None: + raise ValueError("Only one of take_offsets and take_row_ids may be set") + take_query = ( + self.take_row_ids(query.take_row_ids) + if query.take_row_ids is not None + else self.take_offsets(query.take_offsets) + ) + if query.columns is not None: + columns = query.columns + # Native requests serialize dynamic projections as (name, SQL) pairs. + if ( + isinstance(columns, list) + and columns + and isinstance(columns[0], tuple) + ): + columns = dict(columns) + take_query = take_query.select(columns) if query.use_rte is not None: take_query = take_query.use_rte(query.use_rte) if query.with_row_id: @@ -8119,6 +8133,9 @@ class AsyncTable: """ Take a list of row ids from the table. + Results follow the requested row-id order. Repeated IDs produce repeated + rows; missing or deleted IDs are omitted. This also applies to remote tables. + Row ids are not stable and are relative to the current version of the table. They can change due to compaction and updates. diff --git a/python/python/tests/test_merge_insert_lsm.py b/python/python/tests/test_merge_insert_lsm.py index c7d69d779..f5280383c 100644 --- a/python/python/tests/test_merge_insert_lsm.py +++ b/python/python/tests/test_merge_insert_lsm.py @@ -278,7 +278,7 @@ def test_lsm_read_unsupported_shape_errors_without_use_rte_false(tmp_path): # with_row_id is unsupported by the LSM scanner; on an RTE table the default # (auto-routed) read hard-errors instead of silently reading a stale base. - with pytest.raises(Exception): + with pytest.raises(Exception, match=r"with_row_id.*use_rte\(false\)"): table.search().with_row_id(True).to_arrow() # use_rte(False) is the escape hatch: it reads the base table only. @@ -563,19 +563,28 @@ def test_lsm_read_time_travel_errors(tmp_path): assert sorted(base["id"].to_pylist()) == [1, 2, 3] -def test_lsm_read_take_row_ids_errors(tmp_path): +@pytest.mark.parametrize("with_row_id", [False, True]) +@pytest.mark.parametrize("use_rte", [None, True]) +def test_lsm_read_take_row_ids_errors(tmp_path, with_row_id, use_rte): db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0)) table = db.create_table("t", _reader([1, 2, 3])) table.set_unenforced_primary_key("id") table.enable_rte(RTEWriteSpec.unsharded()) _lsm_upsert(table, [4]) - # take-by-row-id auto-routes through the LSM scanner, which has no stable _rowid, - # so it hard-errors instead of failing with an opaque column-not-found error. - with pytest.raises(Exception, match="row id"): - table.take_row_ids([0, 1]).to_arrow() + # Automatic and explicit RTE routing must explain the unsupported take, + # including when row IDs are only requested internally for ordering. + take = table.take_row_ids([1, 0, 1, 2**64 - 1], with_row_id=with_row_id) + if use_rte is not None: + take = take.use_rte(use_rte) + with pytest.raises(Exception, match=r"take by row id.*use_rte\(false\)"): + take.to_arrow() # use_rte(False) is the escape hatch: it reads the base table. - base = table.take_row_ids([0, 1]).use_rte(False).to_arrow() - assert base.num_rows == 2 + base = take.use_rte(False).to_arrow() + assert base["id"].to_pylist() == [2, 1, 2] + assert base["value"].to_pylist() == [1, 0, 1] + assert ("_rowid" in base.column_names) == with_row_id + if with_row_id: + assert base["_rowid"].to_pylist() == [1, 0, 1] def test_lsm_read_fts_postfilter_errors(tmp_path): diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index a4ac3851d..d39f01a39 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -2417,13 +2417,36 @@ def test_take_queries(tmp_path): 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()) - ) == [ - 2, - 5, - 17, - ] + assert table.take_row_ids([5, 2, 17]).to_pandas()["idx"].to_list() == [5, 2, 17] + + +@pytest.mark.parametrize("selection", [None, [], ["idx"], {"_rowid": "100 - idx"}]) +@pytest.mark.parametrize("empty", [False, True]) +def test_take_row_ids_order_and_roundtrip(tmp_path, selection, empty): + table = lancedb.connect(tmp_path).create_table("test", pa.table({"idx": [0, 1]})) + table.add(pa.table({"idx": [2, 3]})) + rows = table.search().with_row_id(True).to_arrow().to_pylist() + row_ids = {row["idx"]: row["_rowid"] for row in rows} + table.delete("idx = 1") + requested = ( + [] if empty else [row_ids[3], row_ids[1], row_ids[3], 2**64 - 1, row_ids[0]] + ) + take = table.take_row_ids(requested) + if selection is not None: + take = take.select(selection) + actual = take.to_arrow() + assert actual.num_rows == (0 if empty else 3) + if selection == {"_rowid": "100 - idx"}: + assert actual.column_names == ["_rowid"] + assert actual["_rowid"].to_pylist() == ([] if empty else [97, 97, 100]) + elif selection == []: + assert actual.num_columns == 0 + else: + assert actual.column_names == ["idx"] + assert actual["idx"].to_pylist() == ([] if empty else [3, 3, 0]) + query = take.to_query_object() + assert query.take_row_ids == requested + assert table._execute_query(query).read_all() == actual def test_take_queries_with_row_id(tmp_path): @@ -2432,6 +2455,7 @@ def test_take_queries_with_row_id(tmp_path): by_row_id = table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow() assert "_rowid" in by_row_id.column_names + assert by_row_id["_rowid"].to_pylist() == [5, 2, 17] by_offset = table.take_offsets([5, 2, 17], with_row_id=True).to_arrow() assert "_rowid" in by_offset.column_names @@ -2446,6 +2470,7 @@ async def test_take_queries_with_row_id_async(tmp_path): by_row_id = await table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow() assert "_rowid" in by_row_id.column_names + assert by_row_id["_rowid"].to_pylist() == [5, 2, 17] by_offset = await table.take_offsets([5, 2, 17], with_row_id=True).to_arrow() assert "_rowid" in by_offset.column_names @@ -2472,16 +2497,12 @@ def test_take_queries_to_batches(tmp_path): rs = list(table.take_row_ids([5, 2, 17]).to_batches()) assert all(isinstance(b, pa.RecordBatch) for b in rs) assert sum(b.num_rows for b in rs) == 3 - assert sorted(v for b in rs for v in b.column("idx").to_pylist()) == [2, 5, 17] + assert [v for b in rs for v in b.column("idx").to_pylist()] == [5, 2, 17] # Take with select projection → to_batches preserves the projection rs = list(table.take_row_ids([5, 2, 17]).select(["label"]).to_batches()) assert all(b.schema.names == ["label"] for b in rs) - assert sorted(v for b in rs for v in b.column("label").to_pylist()) == [ - "17", - "2", - "5", - ] + assert [v for b in rs for v in b.column("label").to_pylist()] == ["5", "2", "17"] def test_getitems(tmp_path): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 8178aec62..cd2261f04 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1422,7 +1422,26 @@ def test_table_to_pandas_not_supported(): def test_table_add_in_threadpool(): + received = [] + def handler(request): + # Consume uploads before responding: closing with unread bytes can reset + # the connection while the client is still reading the response. + if request.headers.get("Transfer-Encoding") == "chunked": + body = bytearray() + while size := int(request.rfile.readline().split(b";")[0], 16): + body.extend(request.rfile.read(size)) + assert request.rfile.read(2) == b"\r\n" + assert request.rfile.read(2) == b"\r\n" + else: + body = request.rfile.read(int(request.headers.get("Content-Length", 0))) + if request.path in ( + "/v1/table/test/create/?mode=create", + "/v1/table/test/insert/", + ): + received.append( + (request.path, pa.ipc.open_stream(body).read_all().to_pylist()) + ) if request.path == "/v1/table/test/insert/": request.send_response(200) request.end_headers() @@ -1461,6 +1480,9 @@ def test_table_add_in_threadpool(): for future in futures: future.result() + assert [rows for path, rows in received if "/create/" in path] == [[{"id": 1}]] + assert [rows for path, rows in received if "/insert/" in path] == [[{"id": 1}]] * 10 + def test_table_create_indices(): # Track received index creation requests to validate name parameter @@ -3232,8 +3254,41 @@ def test_remote_take_row_ids_with_row_id(): assert bodies[-1]["with_row_id"] is True with query_test_table(handler) as table: - table.take_row_ids([0, 1]).to_arrow() - assert "with_row_id" not in bodies[-1] + result = table.take_row_ids([0, 1]).to_arrow() + assert result.column_names == ["id"] + assert bodies[-1]["with_row_id"] is True + + +@pytest.mark.parametrize("with_row_id", [False, True]) +@pytest.mark.parametrize("selection", [None, [], ["id"], ["_rowid", "id"]]) +def test_remote_take_row_ids_order_and_roundtrip(with_row_id, selection): + def handler(body): + assert body["with_row_id"] is True + assert "_rowid IN" in body["filter"] + assert "take_row_ids" not in body + assert "_rowid" not in body.get("columns", []) + data = pa.table( + { + "id": pa.chunked_array([[1, 5], [17]]), + "_rowid": pa.chunked_array([[1, 5], [17]], type=pa.uint64()), + } + ) + return data.select(body.get("columns", ["id"]) + ["_rowid"]) + + with query_test_table(handler) as table: + take = table.take_row_ids([17, 5, 2**64 - 1, 17, 1], with_row_id=with_row_id) + if selection is not None: + take = take.select(selection) + result = take.to_batches(max_batch_length=1).read_all() + assert result.num_rows == 4 + expected_columns = ["id"] if selection is None else list(selection) + if with_row_id: + expected_columns = [name for name in expected_columns if name != "_rowid"] + expected_columns.append("_rowid") + assert result.column_names == expected_columns + for column in expected_columns: + assert result[column].to_pylist() == [17, 5, 17, 1] + assert table._execute_query(take.to_query_object()).read_all() == result def test_query_sync_minimal(): diff --git a/python/src/query.rs b/python/src/query.rs index bda4b7783..2097dac69 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -355,6 +355,7 @@ pub struct PyQueryRequest { pub limit: Option, pub offset: Option, pub take_offsets: Option>, + pub take_row_ids: Option>, pub filter: Option, pub full_text_search: Option>, pub select: PySelect, @@ -387,6 +388,7 @@ impl From for PyQueryRequest { limit: query_request.limit, offset: query_request.offset, take_offsets: query_request.take_offsets, + take_row_ids: query_request.take_row_ids, filter: query_request.filter.map(PyQueryFilter), full_text_search: query_request .full_text_search @@ -419,6 +421,7 @@ impl From for PyQueryRequest { limit: vector_query.base.limit, offset: vector_query.base.offset, take_offsets: vector_query.base.take_offsets, + take_row_ids: vector_query.base.take_row_ids, 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/query.rs b/rust/lancedb/src/query.rs index 10914f5f2..866bfbc0d 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -17,12 +17,15 @@ 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_expr::{EquivalenceProperties, Partitioning}; +use datafusion_physical_expr::{ + EquivalenceProperties, Partitioning, expressions::col as physical_col, +}; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, coalesce_partitions::CoalescePartitionsExec, execution_plan::{Boundedness, EmissionType}, limit::GlobalLimitExec, + projection::ProjectionExec, stream::RecordBatchStreamAdapter, }; use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt, stream, try_join}; @@ -850,6 +853,11 @@ pub struct QueryRequest { #[doc(hidden)] pub take_offsets: Option>, + /// Row IDs whose requested order and multiplicity are restored after lookup. + /// This is client-side metadata, not part of the remote query protocol. + #[doc(hidden)] + pub take_row_ids: Option>, + /// Apply filter to the returned rows. pub filter: Option, @@ -919,6 +927,7 @@ impl Default for QueryRequest { limit: None, offset: None, take_offsets: None, + take_row_ids: None, filter: None, filter_error: None, full_text_search: None, @@ -936,6 +945,10 @@ impl Default for QueryRequest { } impl QueryRequest { + pub(crate) fn has_take(&self) -> bool { + self.take_offsets.is_some() || self.take_row_ids.is_some() + } + /// Add a filter, combining it with any existing filter using a logical AND. /// /// If the new filter cannot be combined with the existing one (because they @@ -1676,42 +1689,42 @@ 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; +fn take_occurrences(indices: &[u64]) -> HashMap { + let mut occurrences = HashMap::with_capacity(indices.len()); + for index in indices { + *occurrences.entry(*index).or_insert(0) += 1; } occurrences } fn restore_take_batch_with_occurrences( batch: RecordBatch, - offsets: &[u64], + indices: &[u64], occurrences: &HashMap, ordering_column: &str, drop_ordering_column: bool, preserve_order: bool, ) -> Result { - let actual_offsets = batch + let actual_indices = 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 + let actual_indices = match actual_indices.data_type() { + DataType::UInt64 => actual_indices .as_primitive::() .values() .to_vec(), - DataType::Int64 => actual_offsets + DataType::Int64 => actual_indices .as_primitive::() .values() .iter() - .map(|offset| { - u64::try_from(*offset).map_err(|_| Error::Schema { + .map(|index| { + u64::try_from(*index).map_err(|_| Error::Schema { message: format!( - "take query ordering column '{ordering_column}' contained a negative offset" + "take query ordering column '{ordering_column}' contained a negative index" ), }) }) @@ -1725,26 +1738,26 @@ fn restore_take_batch_with_occurrences( } }; - let mut desired_order = Vec::with_capacity(offsets.len()); + let mut desired_order = Vec::with_capacity(indices.len()); if preserve_order { - let ordering = actual_offsets + let ordering = actual_indices .iter() .copied() .enumerate() - .map(|(index, offset)| (offset, index as u64)) + .map(|(position, index)| (index, position as u64)) .collect::>(); - // Missing offsets retain the filter-based behavior of returning no row. + // Missing indices retain the filter-based behavior of returning no row. desired_order.extend( - offsets + indices .iter() - .filter_map(|offset| ordering.get(offset).copied()), + .filter_map(|index| ordering.get(index).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)); + // Unordered offset takes and explicitly sorted queries retain the lookup's + // order while restoring the multiplicity of each matching row. + for (position, index) in actual_indices.iter().enumerate() { + if let Some(count) = occurrences.get(index) { + desired_order.extend(std::iter::repeat_n(position as u64, *count)); } } } @@ -1761,7 +1774,15 @@ fn restore_take_batch_with_occurrences( }; if drop_ordering_column { - ordered_batch = ordered_batch.drop_column(ordering_column)?; + // Project preserves the row count even when no user columns were selected. + let columns = ordered_batch + .schema() + .fields() + .iter() + .enumerate() + .filter_map(|(i, field)| (field.name() != ordering_column).then_some(i)) + .collect::>(); + ordered_batch = ordered_batch.project(&columns)?; } Ok(ordered_batch) @@ -1770,32 +1791,32 @@ fn restore_take_batch_with_occurrences( #[cfg(test)] fn restore_take_batch( batch: RecordBatch, - offsets: &[u64], + indices: &[u64], ordering_column: &str, drop_ordering_column: bool, preserve_order: bool, ) -> Result { restore_take_batch_with_occurrences( batch, - offsets, - &take_occurrences(offsets), + indices, + &take_occurrences(indices), ordering_column, drop_ordering_column, preserve_order, ) } -/// Restores the logical offset occurrence sequence above the physical lookup plan. +/// Restores the logical index 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 +/// lookup's partitioning. Ordered takes collect 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, + indices: Vec, occurrences: Arc>, ordering_column: String, drop_ordering_column: bool, @@ -1807,15 +1828,20 @@ struct TakeRestoreExec { impl TakeRestoreExec { fn try_new( input: Arc, - offsets: Vec, + indices: 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() + let schema = input.schema(); + let columns = schema + .fields() + .iter() + .enumerate() + .filter_map(|(i, field)| (field.name() != &ordering_column).then_some(i)) + .collect::>(); + Arc::new(schema.project(&columns)?) } else { input.schema() }; @@ -1838,8 +1864,8 @@ impl TakeRestoreExec { Ok(Self { input, - occurrences: Arc::new(take_occurrences(&offsets)), - offsets, + occurrences: Arc::new(take_occurrences(&indices)), + indices, ordering_column, drop_ordering_column, preserve_order, @@ -1858,7 +1884,7 @@ impl DisplayAs for TakeRestoreExec { write!( formatter, "TakeRestoreExec: occurrences={}", - self.offsets.len() + self.indices.len() ) } } @@ -1907,7 +1933,7 @@ impl ExecutionPlan for TakeRestoreExec { let child = children.into_iter().next().unwrap(); let plan = Self::try_new( child, - self.offsets.clone(), + self.indices.clone(), self.ordering_column.clone(), self.drop_ordering_column, self.preserve_order, @@ -1930,7 +1956,7 @@ impl ExecutionPlan for TakeRestoreExec { let input = self.input.execute(partition, context)?; let output_schema = self.schema.clone(); - let offsets = self.offsets.clone(); + let indices = self.indices.clone(); let occurrences = self.occurrences.clone(); let ordering_column = self.ordering_column.clone(); let drop_ordering_column = self.drop_ordering_column; @@ -1947,7 +1973,7 @@ impl ExecutionPlan for TakeRestoreExec { }; restore_take_batch_with_occurrences( batch, - &offsets, + &indices, &occurrences, &ordering_column, drop_ordering_column, @@ -1960,7 +1986,7 @@ impl ExecutionPlan for TakeRestoreExec { batch.and_then(|batch| { restore_take_batch_with_occurrences( batch, - &offsets, + &indices, &occurrences, &ordering_column, drop_ordering_column, @@ -1998,10 +2024,20 @@ impl ExecutionPlan for TakeRestoreExec { pub struct TakeQuery { parent: Arc, request: QueryRequest, - offsets: Option>, preserve_order: bool, } +struct TakeLookup { + request: QueryRequest, + indices: Vec, + ordering_column: String, + drop_ordering_column: bool, + output_offset: usize, + output_limit: Option, + preserve_order: bool, + output_projection: Option>, +} + impl TakeQuery { /// Create a new `TakeQuery` that will return rows at the given offsets. /// @@ -2020,10 +2056,9 @@ impl TakeQuery { filter: Some(QueryFilter::Datafusion( col("_rowoffset").in_list(in_list, false), )), - take_offsets: Some(offsets.clone()), + take_offsets: Some(offsets), ..Default::default() }, - offsets: Some(offsets), preserve_order: false, } } @@ -2032,14 +2067,20 @@ impl TakeQuery { /// /// See [`crate::Table::take_row_ids`] for more details. pub fn from_row_ids(parent: Arc, row_ids: Vec) -> Self { - let in_list: Vec = row_ids.iter().map(|id| lit(*id)).collect(); + let mut seen = HashSet::with_capacity(row_ids.len()); + let in_list = row_ids + .iter() + .copied() + .filter(|id| seen.insert(*id)) + .map(lit) + .collect(); Self { parent, request: QueryRequest { filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))), + take_row_ids: Some(row_ids), ..Default::default() }, - offsets: None, preserve_order: false, } } @@ -2048,7 +2089,7 @@ impl TakeQuery { /// /// This is reserved for readers whose API explicitly guarantees ordering. pub(crate) fn preserve_order(mut self) -> Self { - debug_assert!(self.offsets.is_some()); + debug_assert!(self.request.take_offsets.is_some()); self.preserve_order = true; self } @@ -2106,29 +2147,124 @@ impl TakeQuery { Ok((request, ordering_column, drop_ordering_column)) } - async fn prepare_offsets_lookup( + async fn prepare_take_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?; + preserve_order: bool, + ) -> Result { + let (indices, by_row_id) = match (&request.take_offsets, &request.take_row_ids) { + (Some(indices), None) => (indices.clone(), false), + (None, Some(indices)) => (indices.clone(), true), + _ => { + return Err(Error::InvalidInput { + message: "take lookup requires exactly one of take_offsets or take_row_ids" + .to_string(), + }); + } + }; + let preserve_order = + preserve_order || (by_row_id && request.order_by.as_ref().is_none_or(Vec::is_empty)); + let (mut request, ordering_column, drop_ordering_column, output_projection) = if by_row_id { + let mut request = request.clone(); + let with_row_id = request.with_row_id; + let mut names = match &request.select { + Select::All => None, + Select::Columns(names) => Some(names.clone()), + Select::Dynamic(pairs) => { + Some(pairs.iter().map(|(name, _)| name.clone()).collect()) + } + Select::Expr(pairs) => Some(pairs.iter().map(|(name, _)| name.clone()).collect()), + }; + let selected_row_id = names + .as_ref() + .is_some_and(|names| names.iter().any(|name| name == ROW_ID)); + // with_row_id appends the metadata column, but Lance leaves an + // existing projection alias named _rowid untouched. Rename that + // alias during lookup so it cannot be mistaken for the real ID. + let alias = if selected_row_id + && matches!(request.select, Select::Dynamic(_) | Select::Expr(_)) + { + let mut alias = "__lancedb_take_row_id".to_string(); + while names.as_ref().is_some_and(|names| names.contains(&alias)) { + alias.push('_'); + } + Some(alias) + } else { + None + }; + match &mut request.select { + // Remote tables expose metadata through with_row_id, not + // through the ordinary selected-column list. + Select::Columns(columns) => columns.retain(|name| name != ROW_ID), + Select::Dynamic(pairs) => { + if let Some(alias) = &alias { + for (name, _) in pairs.iter_mut().filter(|(name, _)| name == ROW_ID) { + *name = alias.clone(); + } + } + } + Select::Expr(pairs) => { + if let Some(alias) = &alias { + for (name, _) in pairs.iter_mut().filter(|(name, _)| name == ROW_ID) { + *name = alias.clone(); + } + } + } + Select::All => {} + } + if with_row_id && let Some(names) = &mut names { + if let Some(position) = names.iter().position(|name| name == ROW_ID) { + names.remove(position); + } + names.push(ROW_ID.to_string()); + } + let output_projection = if selected_row_id { + names.map(|names| { + names + .into_iter() + .map(|name| { + let input_name = if name == ROW_ID { + alias.as_ref().unwrap_or(&name).clone() + } else { + name.clone() + }; + (input_name, name) + }) + .collect() + }) + } else { + None + }; + let drop_row_id = alias.is_some() || (!with_row_id && !selected_row_id); + request.with_row_id = true; + (request, ROW_ID.to_string(), drop_row_id, output_projection) + } else { + let (request, column, drop_column) = + Self::request_with_row_offset(parent, request).await?; + (request, column, drop_column, None) + }; + request.take_offsets = None; + request.take_row_ids = None; // 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(( + Ok(TakeLookup { request, + indices, ordering_column, drop_ordering_column, output_offset, output_limit, - )) + preserve_order, + output_projection, + }) } - fn wrap_offsets_plan( + fn wrap_take_plan( lookup: Arc, - offsets: &[u64], + indices: &[u64], ordering_column: String, drop_ordering_column: bool, output_offset: usize, @@ -2142,7 +2278,7 @@ impl TakeQuery { }; let restored: Arc = Arc::new(TakeRestoreExec::try_new( lookup, - offsets.to_vec(), + indices.to_vec(), ordering_column, drop_ordering_column, preserve_order, @@ -2159,7 +2295,7 @@ impl TakeQuery { } } - fn wrap_offsets_explanation( + fn wrap_take_explanation( lookup: &str, occurrence_count: usize, output_offset: usize, @@ -2199,21 +2335,6 @@ impl TakeQuery { } } - 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 @@ -2258,15 +2379,22 @@ impl TakeQuery { } } -pub(crate) async fn create_take_offsets_plan( +pub(crate) async fn create_take_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 TakeLookup { + request, + indices, + ordering_column, + drop_ordering_column, + output_offset, + output_limit, + preserve_order, + output_projection, + } = TakeQuery::prepare_take_lookup(parent, request, preserve_order).await?; let lookup_options = if preserve_order { options.without_output_batch_length_limit() } else { @@ -2276,43 +2404,73 @@ pub(crate) async fn create_take_offsets_plan( .create_plan(&AnyQuery::Query(request), lookup_options) .await?; - TakeQuery::wrap_offsets_plan( + let plan = TakeQuery::wrap_take_plan( lookup, - offsets, + &indices, ordering_column, drop_ordering_column, output_offset, output_limit, preserve_order, - ) + )?; + if let Some(projection) = output_projection { + let schema = plan.schema(); + let expressions = projection + .into_iter() + .map(|(input, output)| Ok((physical_col(&input, &schema)?, output))) + .collect::>>()?; + Ok(Arc::new(ProjectionExec::try_new(expressions, plan)?)) + } else { + Ok(plan) + } } -pub(crate) async fn explain_take_offsets_plan( +pub(crate) async fn explain_take_plan( parent: &dyn BaseTable, request: &QueryRequest, - offsets: &[u64], verbose: bool, + preserve_order: bool, ) -> Result { - let (request, _, _, output_offset, output_limit) = - TakeQuery::prepare_offsets_lookup(parent, request).await?; + let TakeLookup { + request, + indices, + output_offset, + output_limit, + preserve_order, + output_projection, + .. + } = TakeQuery::prepare_take_lookup(parent, request, preserve_order).await?; let lookup = parent .explain_plan(&AnyQuery::Query(request), verbose) .await?; - Ok(TakeQuery::wrap_offsets_explanation( + let explanation = TakeQuery::wrap_take_explanation( &lookup, - offsets.len(), + indices.len(), output_offset, output_limit, - false, - )) + preserve_order, + ); + if output_projection.is_some() { + Ok(format!( + "ProjectionExec\n{}", + explanation + .lines() + .map(|line| format!(" {line}")) + .collect::>() + .join("\n") + )) + } else { + Ok(explanation) + } } -pub(crate) async fn prepare_take_offsets_request( +pub(crate) async fn prepare_take_request( parent: &dyn BaseTable, request: &QueryRequest, ) -> Result { - let (request, _, _, _, _) = TakeQuery::prepare_offsets_lookup(parent, request).await?; - Ok(request) + Ok(TakeQuery::prepare_take_lookup(parent, request, false) + .await? + .request) } impl HasQuery for TakeQuery { @@ -2323,8 +2481,14 @@ 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; + if self.request.has_take() { + return create_take_plan( + self.parent.as_ref(), + &self.request, + options, + self.preserve_order, + ) + .await; } let req = AnyQuery::Query(self.request.clone()); @@ -2335,7 +2499,7 @@ impl ExecutableQuery for TakeQuery { &self, options: QueryExecutionOptions, ) -> Result { - if self.offsets.is_some() { + if self.request.has_take() { 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); @@ -2354,23 +2518,14 @@ 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, + if self.request.has_take() { + return explain_take_plan( + self.parent.as_ref(), + &self.request, + verbose, self.preserve_order, - )); + ) + .await; } let query = AnyQuery::Query(self.request.clone()); @@ -2378,10 +2533,9 @@ impl ExecutableQuery for TakeQuery { } async fn analyze_plan_with_options(&self, options: QueryExecutionOptions) -> Result { - if self.offsets.is_some() { + if self.request.has_take() { if self.parent.analyze_plan_is_remote() { - let (request, _, _, _, _) = - Self::prepare_offsets_lookup(self.parent.as_ref(), &self.request).await?; + let request = prepare_take_request(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 @@ -4200,13 +4354,17 @@ mod tests { assert!(!analyzed.contains("CoalescePartitionsExec")); } + #[rstest::rstest] + #[case::unsorted(vec![5, 1, 17], vec![5, 1, 17])] + #[case::duplicate_and_missing(vec![17, 5, 17, u64::MAX, 1], vec![17, 5, 17, 1])] + #[case::empty(vec![], vec![])] #[tokio::test] - async fn test_take_row_ids() { + async fn test_take_row_ids(#[case] requested: Vec, #[case] expected: Vec) { let tmp_dir = tempdir().unwrap(); let table = make_test_table(&tmp_dir).await; let results = table - .take_row_ids(vec![5, 1, 17]) + .take_row_ids(requested) .execute() .await .unwrap() @@ -4215,18 +4373,226 @@ mod tests { .unwrap(); assert_eq!(results.len(), 1); - assert_eq!(results[0].num_rows(), 3); + assert_eq!(results[0].num_rows(), expected.len()); assert_eq!(results[0].num_columns(), 2); - let mut ids = results[0] + let ids = results[0] .column_by_name("id") .unwrap() .as_primitive::() .values() .to_vec(); - ids.sort(); + assert_eq!(ids, expected); + } - assert_eq!(ids, vec![1, 5, 17]); + #[rstest::rstest] + #[case::stable(true)] + #[case::physical(false)] + #[tokio::test] + async fn test_take_row_ids_deleted_and_pagination(#[case] stable: bool) { + let tmp_dir = tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let batch = arrow_array::record_batch!( + ("id", Int32, [0, 1, 2, 3]), + ("label", Utf8, ["a", "b", "c", "d"]) + ) + .unwrap(); + let table = conn + .create_table("rows", vec![batch]) + .write_options(crate::table::WriteOptions { + lance_write_params: Some(lance::dataset::WriteParams { + max_rows_per_file: 2, + max_rows_per_group: 2, + enable_stable_row_ids: stable, + ..Default::default() + }), + }) + .execute() + .await + .unwrap(); + let batches = table + .query() + .select(Select::columns(&["id"])) + .with_row_id() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); + let row_ids: HashMap<_, _> = batch["id"] + .as_primitive::() + .values() + .iter() + .copied() + .zip( + batch[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied(), + ) + .collect(); + table.delete("id = 1").await.unwrap(); + let take = table + .take_row_ids(vec![ + row_ids[&3], + row_ids[&1], + row_ids[&3], + u64::MAX, + row_ids[&0], + ]) + .select(Select::columns(&["label"])) + .with_row_id() + .offset(1) + .limit(2); + let batches = take + .execute_with_options(QueryExecutionOptions { + max_batch_length: 1, + ..Default::default() + }) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(batches.iter().all(|batch| batch.num_rows() <= 1)); + let actual = concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!( + actual["label"] + .as_string::() + .iter() + .collect::>(), + vec![Some("d"), Some("a")] + ); + assert_eq!( + actual[ROW_ID] + .as_primitive::() + .values() + .as_ref(), + &[row_ids[&3], row_ids[&0]] + ); + let request = take.clone().into_request(); + let direct = table + .base_table() + .query(&AnyQuery::Query(request), QueryExecutionOptions::default()) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + concat_batches(&direct[0].schema(), &direct).unwrap(), + actual + ); + let plan = take + .create_plan(QueryExecutionOptions::default()) + .await + .unwrap(); + let planned = execute_plan(plan, Default::default()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + concat_batches(&planned[0].schema(), &planned).unwrap(), + actual + ); + } + + #[rstest::rstest] + #[case::empty(Select::Columns(vec![]), vec![], vec![])] + #[case::explicit_row_id(Select::columns(&[ROW_ID, "id"]), vec![ROW_ID, "id"], vec![5, 1, 17])] + #[case::dynamic_alias(Select::dynamic(&[(ROW_ID, "20 - id"), ("__lancedb_take_row_id", "id")]), vec![ROW_ID, "__lancedb_take_row_id"], vec![15, 19, 3])] + #[case::expr_alias(Select::expr_projection(&[(ROW_ID, lit(20) - col("id"))]), vec![ROW_ID], vec![15, 19, 3])] + #[tokio::test] + async fn test_take_row_ids_projection( + #[case] selection: Select, + #[case] names: Vec<&str>, + #[case] values: Vec, + #[values(false, true)] with_row_id: bool, + #[values(false, true)] order_by: bool, + ) { + let tmp_dir = tempdir().unwrap(); + let table = make_test_table(&tmp_dir).await; + let mut take = table.take_row_ids(vec![5, 1, 17]).select(selection); + if with_row_id { + take = take.with_row_id(); + } + if order_by { + take = take.order_by(Some(vec![ColumnOrdering::asc_nulls_first( + "id".to_string(), + )])); + } + let batches = take + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let actual = concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!(actual.num_rows(), 3); + let mut expected_names = names.clone(); + if with_row_id { + expected_names.retain(|name| *name != ROW_ID); + expected_names.push(ROW_ID); + } + assert_eq!( + actual + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + expected_names + ); + if !values.is_empty() { + let row_ids = actual.column_by_name(ROW_ID).unwrap(); + if row_ids.data_type() == &DataType::UInt64 { + assert_eq!( + row_ids.as_primitive::().values().as_ref(), + if order_by { &[1, 5, 17] } else { &[5, 1, 17] } + ); + } else { + let mut expected = values; + if order_by { + expected.swap(0, 1); + } + assert_eq!( + row_ids.as_primitive::().values().as_ref(), + expected.as_slice() + ); + } + } + } + + #[tokio::test] + async fn test_take_row_ids_explicit_order_by() { + let tmp_dir = tempdir().unwrap(); + let table = make_test_table(&tmp_dir).await; + let batches = table + .take_row_ids(vec![17, 5, 1, 5]) + .order_by(Some(vec![ColumnOrdering::asc_nulls_first( + "id".to_string(), + )])) + .offset(1) + .limit(2) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let actual = concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!( + actual["id"].as_primitive::().values().as_ref(), + &[5, 5] + ); } } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index af0946600..6534f5519 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3013,10 +3013,9 @@ impl BaseTable for RemoteTable { options: QueryExecutionOptions, ) -> Result> { if let AnyQuery::Query(request) = query - && let Some(offsets) = &request.take_offsets + && request.has_take() { - return crate::query::create_take_offsets_plan(self, request, offsets, options, false) - .await; + return crate::query::create_take_plan(self, request, options, false).await; } let streams = self.execute_query(query, &options).await?; @@ -3038,16 +3037,10 @@ impl BaseTable for RemoteTable { options: QueryExecutionOptions, ) -> Result { if let AnyQuery::Query(request) = query - && let Some(offsets) = &request.take_offsets + && request.has_take() { - let plan = crate::query::create_take_offsets_plan( - self, - request, - offsets, - options.clone(), - false, - ) - .await?; + let plan = + crate::query::create_take_plan(self, request, 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 { @@ -3108,9 +3101,9 @@ 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 + && request.has_take() { - return crate::query::explain_take_offsets_plan(self, request, offsets, verbose).await; + return crate::query::explain_take_plan(self, request, verbose, false).await; } let base_request = self @@ -3166,10 +3159,10 @@ impl BaseTable for RemoteTable { options: QueryExecutionOptions, ) -> Result { let prepared_query = if let AnyQuery::Query(request) = query - && request.take_offsets.is_some() + && request.has_take() { Some(AnyQuery::Query( - crate::query::prepare_take_offsets_request(self, request).await?, + crate::query::prepare_take_request(self, request).await?, )) } else { None @@ -7741,8 +7734,11 @@ mod tests { assert_eq!(result, "analyzed plan"); } + #[rstest::rstest] + #[case::offsets(false)] + #[case::row_ids(true)] #[tokio::test] - async fn test_take_offsets_explain_plan_does_not_execute_query() { + async fn test_take_explain_plan_does_not_execute_query(#[case] by_row_id: bool) { 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/"); @@ -7753,8 +7749,12 @@ mod tests { .unwrap() }); - let explained = table - .take_offsets(vec![0, 1, 0, 2]) + let take = if by_row_id { + table.take_row_ids(vec![0, 1, 0, 2]) + } else { + table.take_offsets(vec![0, 1, 0, 2]) + }; + let explained = take .select(crate::query::Select::columns(&["id"])) .limit(3) .explain_plan(false) @@ -7763,7 +7763,7 @@ mod tests { assert!(explained.contains("GlobalLimitExec")); assert!(explained.contains("TakeRestoreExec")); - assert!(!explained.contains("CoalescePartitionsExec")); + assert_eq!(explained.contains("CoalescePartitionsExec"), by_row_id); assert!(explained.contains("RemoteLookupExec")); } @@ -7816,8 +7816,11 @@ mod tests { ); } + #[rstest::rstest] + #[case::offsets(false)] + #[case::row_ids(true)] #[tokio::test] - async fn test_take_offsets_analyze_plan_delegates_to_remote() { + async fn test_take_analyze_plan_delegates_to_remote(#[case] by_row_id: bool) { 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/"); @@ -7835,8 +7838,12 @@ mod tests { .unwrap() }); - let analyzed = table - .take_offsets(vec![0, 1, 0, 2]) + let take = if by_row_id { + table.take_row_ids(vec![0, 1, 0, 2]) + } else { + table.take_offsets(vec![0, 1, 0, 2]) + }; + let analyzed = take .select(crate::query::Select::columns(&["id"])) .limit(3) .analyze_plan_with_options(QueryExecutionOptions { @@ -7849,6 +7856,84 @@ mod tests { assert_eq!(analyzed, "Remote analyzed plan: worker metrics"); } + #[rstest::rstest] + #[tokio::test] + async fn test_take_row_ids_restores_remote_order( + #[values(false, true)] with_row_id: bool, + #[values(false, true)] converted: bool, + ) { + let table = Table::new_with_handler("my_table", |request| { + 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"])); + assert_eq!(body["with_row_id"], true); + assert_eq!(body["k"], json!(isize::MAX as usize)); + assert!(body.get("offset").is_none()); + assert!(body["filter"].as_str().unwrap().contains("_rowid IN")); + assert!(body.get("take_row_ids").is_none()); + + // The server can return the physical rows in any order and batches. + let data = arrow_array::record_batch!( + ("id", Int32, [1, 5, 17]), + ("_rowid", UInt64, [1, 5, 17]) + ) + .unwrap(); + let mut response = Vec::new(); + { + let mut writer = + arrow_ipc::writer::FileWriter::try_new(&mut response, &data.schema()).unwrap(); + writer.write(&data.slice(0, 2)).unwrap(); + writer.write(&data.slice(2, 1)).unwrap(); + writer.finish().unwrap(); + } + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE) + .body(response) + .unwrap() + }); + let mut take = table + .take_row_ids(vec![17, 5, u64::MAX, 17, 1]) + .select(crate::query::Select::columns(&["id"])) + .offset(1) + .limit(3); + if with_row_id { + take = take.with_row_id(); + } + let options = QueryExecutionOptions { + max_batch_length: 1, + ..Default::default() + }; + let batches = if converted { + table + .base_table() + .query(&AnyQuery::Query(take.into_request()), options) + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + } else { + take.execute_with_options(options) + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + }; + assert!(batches.iter().all(|batch| batch.num_rows() <= 1)); + let actual = arrow::compute::concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!( + actual["id"] + .as_primitive::() + .values() + .as_ref(), + &[5, 17, 1] + ); + assert_eq!(actual.column_by_name("_rowid").is_some(), with_row_id); + } + #[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 1bd3a711a..34846d33e 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1864,9 +1864,20 @@ impl Table { /// There is an ongoing effort to make row ids stable which is tracked at /// https://github.com/lancedb/lancedb/issues/1120 /// - /// 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 ids, you will need - /// to add the row id column to the output and align it yourself. + /// Results follow the requested row-id order, including repeated IDs. Missing or + /// deleted IDs are omitted. An explicit nonempty [`QueryBase::order_by`] overrides + /// the requested order. Limits and offsets apply after restoring occurrences. + /// These guarantees also apply to remote tables. + /// + /// ```rust + /// use lancedb::query::ExecutableQuery; + /// # async fn example(table: &lancedb::Table, first: u64, second: u64) -> lancedb::Result<()> { + /// // If both IDs exist, return second, first, then second again. + /// let rows = table.take_row_ids(vec![second, first, second]).execute().await?; + /// # Ok(()) + /// # } + /// ``` + /// /// Parameters /// ---------- /// row_ids: list[int] diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 8a72872e8..e9bcd380d 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -121,7 +121,7 @@ fn requires_local_namespace_execution(query: &AnyQuery) -> bool { // The namespace QueryTable request cannot represent approx_mode, search_effort or use_rte. // For use_rte that is worse than a tuning miss: RTE read routing lives only // in `create_plan`, so a pushed-down query would return stale base-only data. - if query.base().use_rte.is_some() || query.base().take_offsets.is_some() { + if query.base().use_rte.is_some() || query.base().has_take() { return true; } matches!( @@ -164,10 +164,9 @@ pub async fn create_plan( ) -> Result> { let query = query.canonicalized()?; if let AnyQuery::Query(request) = &query - && let Some(offsets) = &request.take_offsets + && request.has_take() { - return crate::query::create_take_offsets_plan(table, request, offsets, options, false) - .await; + return crate::query::create_take_plan(table, request, options, false).await; } let query = match query { diff --git a/rust/lancedb/src/table/query/lsm.rs b/rust/lancedb/src/table/query/lsm.rs index d348a6668..400df6c82 100644 --- a/rust/lancedb/src/table/query/lsm.rs +++ b/rust/lancedb/src/table/query/lsm.rs @@ -161,6 +161,19 @@ fn reject_unsupported(query: &VectorQueryRequest) -> Result<()> { if !query.query_vector.is_empty() && query.base.full_text_search.is_some() { return unsupported("hybrid (vector + full-text) search"); } + // Take lookups add with_row_id internally to restore requested order. Report + // the unsupported take before that implementation detail. Neither _rowid nor + // _rowoffset is exposed by the LSM scanner, which only exposes _rowaddr. + if let Some(QueryFilter::Datafusion(expr)) = &query.base.filter + && expr + .column_refs() + .iter() + .any(|c| c.name == "_rowid" || c.name == "_rowoffset") + { + return unsupported( + "take by row id or row offset (the LSM scanner has no stable _rowid / _rowoffset)", + ); + } if query.base.with_row_id { return unsupported("with_row_id (the LSM scanner exposes _rowaddr, not a stable _rowid)"); } @@ -204,19 +217,6 @@ fn reject_unsupported(query: &VectorQueryRequest) -> Result<()> { if let Some(QueryFilter::Substrait(_)) = &query.base.filter { return unsupported("Substrait filters"); } - // Take-by-row-id / row-offset queries carry a `_rowid` / `_rowoffset` filter, - // columns the LSM scanner never exposes (only `_rowaddr`); reject with guidance - // rather than failing deep in datafusion with a column-not-found error. - if let Some(QueryFilter::Datafusion(expr)) = &query.base.filter - && expr - .column_refs() - .iter() - .any(|c| c.name == "_rowid" || c.name == "_rowoffset") - { - return unsupported( - "take by row id or row offset (the LSM scanner has no stable _rowid / _rowoffset)", - ); - } Ok(()) }