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).

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]]]

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

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.
 

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):

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):

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():

python/src/query.rs

@@ -355,6 +355,7 @@ pub struct PyQueryRequest {
     pub limit: Option<usize>,
     pub offset: Option<usize>,
     pub take_offsets: Option<Vec<u64>>,
+    pub take_row_ids: Option<Vec<u64>>,
     pub filter: Option<PyQueryFilter>,
     pub full_text_search: Option<PyLanceDB<FtsQuery>>,
     pub select: PySelect,
@@ -387,6 +388,7 @@ impl From<AnyQuery> 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<AnyQuery> 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),