From 69ac2158a1b590855c46b2d13c92074383e0ce39 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 5 Nov 2026 16:54:23 +0800 Subject: [PATCH] fix: preserve take_row_ids request order (#4826) `take_row_ids` used an `IN` filter, so the backend could reorder rows and collapse duplicate IDs. That can silently misalign training rows with labels or sample weights. Restore the requested order and duplicate occurrences in the shared Rust take path for native and remote tables, and retain that information through Python query serialization. Missing or deleted IDs are omitted. Explicit Rust `order_by` overrides request order, and pagination applies after restoring occurrences. Empty projections and `_rowid` projection aliases retain their output schema. RTE tables reject unsupported takes with the existing base-table guidance before checking the internal row-ID projection used for ordering. An explicit `use_rte(False)` reads the base table with the same requested order and occurrences. A pre-existing fixture defect also blocked CI: `test_table_add_in_threadpool` responded without reading uploads, allowing a connection reset during response decoding. This PR includes a fix confined to that test's handler: consume each request body before responding and verify one create row plus ten inserted rows. This was not a regression in the take implementation. Production write retries, concurrency and HTTP client code are unchanged. The selective sampling guide scopes the row-ID guarantees to SDK releases containing this fix, preserves the distinct offset-take semantics, and explains that ordered takes collect and combine matching rows before output batching. It recommends smaller ID requests and narrow projections because `max_batch_length` does not cap working memory. This uses the existing REST query protocol and requires a new SDK release; no deployed service was changed. Validation: - Docs follow-up: the changed page compiles with the Mint 4.2.888 MDX compiler; Typos 1.26.0 passes and Vale 3.24.0 has no findings on changed lines. Full-site Mint validation reports four existing snippet-import warnings on the reranking/vector-search pages, identical with the original page restored in a baseline comparison. - The applied threadpool test passed locally after the fixture change. The identical candidate previously passed 10 trials with the ordinary reader and 10 with an unbuffered reader; each trial verifies all uploaded rows. Those checks used macOS arm64/Python 3.14/default features, not the full failing CI environment. The normal new-head macOS Python 3.13/experimental-functions job remains the required platform check. - Python package Ruff format/check passed. Repository-root Ruff still reports 59 existing errors and six formatting differences in unchanged `docs/web-tests` files; none are in changed files. - At `a569d3f0`, the RTE diagnostic regression was reproduced and fixed; 5 targeted cases and strict workspace Clippy passed. The fixture-only follow-up changes no production Rust/Python code, so that evidence is reused. - Earlier validation at `6bcc0878`: 439 Rust query/remote-table tests and 24 Python take/getitems tests passed; the focused 49-test suite also covered offset dataloaders. A local 9,216-row, four-fragment reproduction returned 200 seeded random batches of 64 rows with zero ordering or label-alignment mismatches. The ordering and projection implementation is unchanged by the follow-ups. - No live Enterprise run was performed. Fixes [ENT-3469](https://linear.app/lancedb/issue/ENT-3469). --- docs/web/curation/selective-sampling.mdx | 41 +- python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/query.py | 5 + python/python/lancedb/table.py | 31 +- python/python/tests/test_merge_insert_lsm.py | 25 +- python/python/tests/test_query.py | 47 +- python/python/tests/test_remote_db.py | 59 +- python/src/query.rs | 3 + rust/lancedb/src/query.rs | 604 +++++++++++++++---- rust/lancedb/src/remote/table.rs | 131 +++- rust/lancedb/src/table.rs | 17 +- rust/lancedb/src/table/query.rs | 7 +- rust/lancedb/src/table/query/lsm.rs | 26 +- 13 files changed, 795 insertions(+), 202 deletions(-) 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(()) }