From 9b13e8f72530c8b1950c09685f6920711154a539 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 13 Nov 2026 02:05:39 +0800 Subject: [PATCH] feat(python): materialize remote blobs at the query snapshot (#4973) Python remote queries materialize Blob v2 payloads with `blob_mode="bytes"` through synchronous and asynchronous `to_arrow`, `to_list`, `to_pandas`, and `to_batches`. Deferred reads bind to the actual producing dataset without changing the caller's live table handle, including concurrent writes and same-name DROP/CREATE. Trust only the authoritative `x-lancedb-query-version` plus `x-lancedb-query-snapshot`. A numeric version alone cannot distinguish a recreated table; the ordinary live watermark and user-supplied Arrow metadata cannot establish provenance. An independent immutable handle carries the original identity through every batch, byte-limit split, and large-Blob Range fallback, and requires a verified response echo before accepting payloads. Aliases, NULL/empty results, rerankers, and hybrid legs retain consistent provenance. Missing capability, mixed identities, or a serving node ignoring the precondition fails explicitly; ordinary descriptor queries remain available on older servers. The source docs distinguish Enterprise SQL payload bytes from SDK descriptors. Server dependency: https://github.com/lancedb/sophon/pull/10211. Validation uses retained, independently audited local evidence: 724 Python cases passed (one optional PyAV skip), 67 focused SDK Rust cases passed, 36 server checks passed, and 15 actual REST/SDK scenarios passed across retained runs. These cover sync/async and Arrow file/stream recreation, concurrent overwrite, deletion of a selected version, legacy capabilities, forged IPC metadata, WAL exclusion, and actual mid-batch recreation. The original aggregate experiment exited 101 because its extra harness assumed a Python batch size would split the native result; its 36 server checks and 13 completed scenarios remain valid. The corrected native-batch experiment passed separately for sync and async. No aggregate all-green claim is made. The public batched-range operation has separate cross-partition identity integration in https://github.com/lancedb/sophon/pull/10219 and https://github.com/lancedb/lancedb/pull/4994. This change does not claim to fix that loop. Realtime/MemWAL remains outside Blob v2 GA, without disabling existing writes or compaction. No dependency pins changed. Merge/rollout hold remains: upgrade every serving node before enabling the new SDK capability, verify proxy forwarding of both authoritative headers and the verification echo, and preserve queried versions for later reads. Missing or inconsistent capability fails explicitly; identity does not retain versions or authenticate storage paths. Actual deployed proxy/Cloud/PE behavior and production retention remain unverified. Old SDKs do not gain incarnation protection from a server-only upgrade. Published documentation retains the Early access badge and released Enterprise instructions, including explicit fetches for copying. The new Enterprise `blob_mode="bytes"` guidance, snapshot-bound copy example, future streaming restrictions and GA wording remain source-only MDX drafts until release. The correction distinguishing released Enterprise SQL bytes from SDK descriptors remains visible. Documentation validation checked comment boundaries, visible instructions, Python example syntax and whitespace; the site was not built. This documentation-only revision reuses the implementation evidence above. --- docs/web/management/table.mdx | 102 +++++- docs/web/sql/dql/blob-descriptor.mdx | 8 +- python/python/lancedb/_lancedb.pyi | 7 +- python/python/lancedb/query.py | 138 ++++++-- python/python/tests/test_remote_db.py | 457 +++++++++++++++++++++++-- python/src/table.rs | 16 +- rust/lancedb/src/remote/table.rs | 171 ++++++++- rust/lancedb/src/remote/table/blobs.rs | 30 +- rust/lancedb/src/table.rs | 23 ++ 9 files changed, 876 insertions(+), 76 deletions(-) diff --git a/docs/web/management/table.mdx b/docs/web/management/table.mdx index 23fdbdb4a..031342fbc 100644 --- a/docs/web/management/table.mdx +++ b/docs/web/management/table.mdx @@ -375,8 +375,16 @@ file, up to 64 KiB by default, and larger values are stored outside it. A `blob` column keeps metadata separate from the payload bytes, so queries can filter rows without loading the blobs. -Most queries return descriptors, which contain metadata about the blob instead -of its bytes. You can +{/* Draft for Blob v2 GA. Confirm release status before removing Early access. + +Blob v2 GA supports standard tables. Tables with Realtime or MemWAL enabled +are outside the GA support scope. +*/} + +SDK queries return descriptors by default, which contain metadata about the +blob instead of its bytes. Enterprise SQL queries return payload bytes by +default; use `blob_descriptor(column)` to select metadata instead. With an SDK, +you can [read the bytes in the query results](#multimodal-blob-query-results) or [fetch them by row ID](#multimodal-read-blobs). @@ -500,14 +508,41 @@ the default `"lazy"` mode because it would open `BlobFile` handles. Query with `blob_mode="descriptions"`, then read bytes with `fetch_blobs` or `fetch_blob_files`. +{/* Draft for the Enterprise Python bytes release. Replace the released instructions above only after SDK and serving-node rollout. + +`blob_mode="bytes"` also supports Enterprise queries. The SDK locks all blob +fetches, including streaming batches, to both the version in +`x-lancedb-query-version` and the incarnation/revision identity in +`x-lancedb-query-snapshot`, captured from the same query response. Every payload +request verifies that identity against the actual opened dataset before reading +rows. Concurrent writes therefore preserve the selected payloads. Dropping and +recreating a table, even at the same location and version, cannot substitute the +replacement's bytes: the original snapshot is served or the read fails explicitly. + +Upgrade the SDK and every serving node to support snapshot identities. A server +that reports only a numeric query version, a legacy `x-lancedb-version` watermark, +or Arrow schema metadata does not enable deferred bytes. Missing or unsupported +identities raise `NotImplementedError`; use `"descriptions"` until the upgrade is +complete. Ordinary queries, descriptors, and independent `fetch_blobs` calls keep +working. Old SDKs keep their existing behavior and do not gain this protection +from a server-only upgrade. Realtime/MemWAL and historical manifests without a +native transaction identity do not support deferred bytes. Snapshot identities +do not retain old versions; deleting the selected version still fails explicitly. +In-memory tables and namespace `QueryTable` pushdown do not support bytes mode. +On an Enterprise table, sync `Table.to_pandas()` and `Table.to_arrow()` are not +supported; use `table.search()` instead. A query's `to_pandas()` also raises in +the default `"lazy"` mode because remote lazy handles are not supported. +*/} + ###### Copy a subset of blob rows {#multimodal-copy-subset} -Blob query results describe the stored payload; they are not the bytes to -insert into another table. These APIs return different values: +SDK descriptors identify stored payloads; materialize the bytes before inserting +them into another table. These APIs return different values: | Read path | Value | | --- | --- | -| SQL or an Arrow query in `"descriptions"` mode | Descriptor struct identifying the stored blob. Treat it as opaque. | +| Enterprise SQL | Raw payload bytes by default. | +| An SDK Arrow query in `"descriptions"` mode | Descriptor struct identifying the stored blob. Treat it as opaque. | | `table.fetch_blobs(column, row_ids)` | Arrow `large_binary` array; `.to_pylist()` gives `bytes` or `None` in the requested order. | | A Function parameter annotated `lancedb.BlobReader` | Lazy, seekable reader over the input blob, available inside the Function. | @@ -545,6 +580,50 @@ Empty blobs remain `b""` and null blobs remain `None`. Fetch bounded batches for larger subsets; `fetch_blobs` materializes all requested bytes in memory. For incremental decoding, use `fetch_blob_files` instead. +{/* Draft for the Enterprise Python bytes release. Replace the copy example above when snapshot-bound remote bytes are released. + +For Enterprise tables, prefer server-side `CREATE TABLE AS SELECT` or +`INSERT INTO ... SELECT` to copy blobs without a client round trip: + +```sql +CREATE TABLE image_subset AS SELECT id, image FROM images WHERE id < 100; +-- To append to an existing destination: +INSERT INTO image_subset SELECT id, image FROM images WHERE id >= 100 AND id < 200; +``` + +For SDK copies, query with `blob_mode="bytes"` to bind selection and payload +reads to one snapshot. The following example copies up to 100 rows from an +existing `images` table with an integer `id` and a blob `image` column. + +```python +import lancedb +import pyarrow as pa + +source = db.open_table("images") +hits = ( + source.search() + .where("id < 100") + .select(["id", "image"]) + .limit(100) + .to_arrow(blob_mode="bytes") +) +subset = db.create_table( + "image_subset", + schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]), +) +if hits.num_rows: + subset.add(hits) +``` + +Empty blobs remain `b""` and null blobs remain `None`. Fetch bounded batches +for larger subsets; `fetch_blobs` materializes all requested bytes in memory. +For incremental decoding, use `fetch_blob_files` instead. + +Remote bytes mode requires a server that supports snapshot-bound Blob reads. +Older servers and Realtime/MemWAL tables can still return descriptors, but +bytes mode raises an explicit unsupported-operation error before fetching payloads. +*/} + `blob_mode="bytes"` materializes top-level and struct-nested blob columns. It raises `ValueError` when the selected columns include blobs inside a list. Select those columns with `blob_mode="descriptions"`, or leave @@ -579,6 +658,19 @@ overwrite the table, use a separate reader checked out to a fixed version before querying. Otherwise, an overwrite can invalidate the row IDs used to fetch bytes and raise `ValueError`. +{/* Draft for the Enterprise Python bytes release. Replace the streaming restrictions above after release. + +Bytes mode has the same restrictions as `to_arrow(blob_mode="bytes")`, +including no blobs inside lists, namespace `QueryTable` pushdown, or query +shapes that cannot preserve their row results. Enterprise queries require a +valid `x-lancedb-query-version` and `x-lancedb-query-snapshot` response headers. +Legacy v1 bytes are limited to +eligible local plain scans. For local queries, if concurrent writers can +overwrite the table, use a separate reader checked out to a fixed version +before querying. Otherwise, an overwrite can invalidate the row IDs used to +fetch bytes and raise `ValueError`. +*/} + ###### Optional pandas examples {#multimodal-blob-pandas} The lazy-handle calls below require pylance. The query example expects `tbl_async`, diff --git a/docs/web/sql/dql/blob-descriptor.mdx b/docs/web/sql/dql/blob-descriptor.mdx index 5021f2d20..7ba83311c 100644 --- a/docs/web/sql/dql/blob-descriptor.mdx +++ b/docs/web/sql/dql/blob-descriptor.mdx @@ -15,10 +15,10 @@ SELECT id, blob_descriptor(video) AS descriptor FROM videos ``` -A [BlobV2](/management/table#managed-blobs) column's bytes are not fetched by -a plain `SELECT`: that already returns `large_binary`, the same as any other -binary column. `blob_descriptor` instead returns a small struct describing -where the payload lives, without reading it. Pass one row's descriptor to +A plain `SELECT` of a [BlobV2](/management/table#managed-blobs) column returns +its payload as `large_binary`. `blob_descriptor` instead returns a small +struct describing where the payload lives, without reading the payload. +Pass one row's descriptor to the SDK's blob file reader to open that payload directly: diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 15234532d..956806cde 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -749,7 +749,12 @@ class Table: def take_row_ids(self, row_ids: list[int]) -> TakeQuery: ... async def blob_columns(self) -> list[str]: ... async def fetch_blobs( - self, column: str, row_ids: list[int] + self, + column: str, + row_ids: list[int], + *, + version: Optional[int] = None, + snapshot_identity: Optional[str] = None, ) -> pa.LargeBinaryArray: ... async def fetch_blob_ranges( self, diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index bab40be95..f99aaeb73 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio +import re from abc import ABC, abstractmethod from concurrent.futures import ThreadPoolExecutor from datetime import timedelta @@ -194,20 +195,57 @@ def _unsupported_blob_pandas_error(reason: str) -> RuntimeError: ) -def _supports_local_blob_fetch( +def _supports_blob_fetch( async_table: "AsyncTable", pushdown_operations: Iterable[str] ) -> bool: - # Remote tables and QueryTable pushdown answer queries on a server, so a - # fetch through this handle could read another version. - return async_table._inner._is_native() and "QueryTable" not in pushdown_operations + # Remote queries carry their producing snapshot in response metadata. + # Namespace QueryTable pushdown has no equivalent snapshot contract. + return ( + not async_table._inner._is_native() or "QueryTable" not in pushdown_operations + ) + + +def _remote_blob_snapshot(schema: pa.Schema) -> tuple[int, str]: + value = (schema.metadata or {}).get(b"lancedb:query_version", b"") + if not value.isdigit() or not 0 < int(value) < 2**64: + raise NotImplementedError( + "blob_mode='bytes' requires a stable table snapshot: the query response " + "has a missing, invalid, or inconsistent x-lancedb-query-version header. " + "Upgrade the server and use a table without Realtime/MemWAL, or use " + "blob_mode='descriptions'." + ) + identity = (schema.metadata or {}).get(b"lancedb:query_snapshot", b"") + if not re.fullmatch(rb"v1:[0-9a-f]{64}", identity): + raise NotImplementedError( + "blob_mode='bytes' requires a trustworthy query snapshot identity: " + "missing or unsupported x-lancedb-query-snapshot. Upgrade the server " + "on every serving node, or use blob_mode='descriptions'." + ) + return int(value), identity.decode("ascii") + + +def _merge_query_snapshot(result: pa.Table, *sources: pa.Table) -> pa.Table: + """Preserve a snapshot only when every query leg used the same identity.""" + keys = (b"lancedb:query_version", b"lancedb:query_snapshot") + snapshots = { + tuple((source.schema.metadata or {}).get(key) for key in keys) + for source in sources + } + metadata = dict(result.schema.metadata or {}) + for key in keys: + metadata.pop(key, None) + snapshot = next(iter(snapshots)) if len(snapshots) == 1 else None + if snapshot is not None and all(value is not None for value in snapshot): + metadata.update(zip(keys, snapshot)) + return result.replace_schema_metadata(metadata) def _blob_bytes_without_snapshot_error() -> NotImplementedError: return NotImplementedError( "blob_mode='bytes' cannot safely materialize blob columns without a " - "stable table snapshot. LanceDB Enterprise and namespace QueryTable pushdown " - "run the query on a server, so the blob fetch could read another table " - "version. Use blob_mode='descriptions' to inspect descriptors. To read " + "stable table snapshot. Namespace QueryTable pushdown does not bind " + "the query and blob fetches to one table version. " + "Use blob_mode='descriptions' to inspect descriptors. To read " "payloads, select row IDs with with_row_id(True), then call " "table.fetch_blobs(column, row_ids) or table.fetch_blob_files(column, " "row_ids). Keep the source unchanged between selection and fetching." @@ -1297,9 +1335,9 @@ class LanceQueryBuilder(ABC): columns without pylance. Other local blob v2 queries return descriptors in ``"lazy"`` and ``"descriptions"`` modes, or payloads in ``"bytes"`` mode, through :meth:`to_arrow`. - LanceDB Enterprise blob queries support only ``"descriptions"`` until a - stable - table snapshot can bind the query and blob fetches. + Enterprise queries support ``"bytes"`` at the query response's + snapshot and ``"descriptions"`` without payload reads. Remote + ``"lazy"`` handles are not supported. **kwargs Forwarded to pyarrow.Table.to_pandas after query execution and optional flattening. @@ -1307,6 +1345,10 @@ class LanceQueryBuilder(ABC): validate_blob_mode(blob_mode) if _is_remote_query_table(self._table): blob_sources = self._validated_blob_sources(blob_mode) + if blob_mode == "bytes": + return flatten_columns( + self.to_arrow(timeout=timeout, blob_mode=blob_mode), flatten + ).to_pandas(**kwargs) return _finish_remote_pandas( self.to_arrow(timeout=timeout), blob_mode=blob_mode, @@ -1629,7 +1671,8 @@ class LanceQueryBuilder(ABC): _query_is_plain_blob_scan(query) and getattr(self, "_reranker", None) is None and not _query_needs_lance_scanner(self) - and _supports_local_blob_fetch( + and self._table._table._inner._is_native() + and _supports_blob_fetch( self._table._table, getattr(self._table, "_pushdown_operations", ()), ) @@ -1650,11 +1693,12 @@ class LanceQueryBuilder(ABC): if blob_mode != "bytes": return blob_sources reject_unfetchable_blob_columns(schema, self._columns, blob_sources) - if blob_sources and not _supports_local_blob_fetch( + if blob_sources and not _supports_blob_fetch( self._table._table, getattr(self._table, "_pushdown_operations", ()) ): raise _blob_bytes_without_snapshot_error() - reject_unmatched_blob_outputs(self._blob_output_schema(), blob_sources) + if not _is_remote_query_table(self._table): + reject_unmatched_blob_outputs(self._blob_output_schema(), blob_sources) return blob_sources def _blob_output_schema(self) -> pa.Schema: @@ -1667,9 +1711,20 @@ class LanceQueryBuilder(ABC): return tbl if blob_mode == "descriptions": return strip_auto_row_ids(tbl, blob_sources) - return replace_blob_columns_with_bytes_sync( - tbl, blob_sources, self._table.fetch_blobs - ) + reject_unmatched_blob_outputs(tbl.schema, blob_sources) + if _is_remote_query_table(self._table) and blob_sources: + version, identity = _remote_blob_snapshot(tbl.schema) + + async def fetch_at_version(column, row_ids): + return await self._table._table._inner.fetch_blobs( + column, row_ids, version=version, snapshot_identity=identity + ) + + def fetch_blobs(column, row_ids): + return LOOP.run(fetch_at_version(column, row_ids)) + else: + fetch_blobs = self._table.fetch_blobs + return replace_blob_columns_with_bytes_sync(tbl, blob_sources, fetch_blobs) def _finalize_blob_query_table( self, @@ -2386,6 +2441,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): rs_table = result_set.read_all() result_set = self._reranker.rerank_vector(self._str_query, rs_table) check_reranker_result(result_set) + result_set = _merge_query_snapshot(result_set, rs_table) # convert result_set back to RecordBatchReader result_set = pa.RecordBatchReader.from_batches( result_set.schema, result_set.to_batches() @@ -2599,8 +2655,10 @@ class LanceFtsQueryBuilder(LanceQueryBuilder): query = self._query_for_scan(blob_sources) results = self._table._execute_query(query, timeout=timeout).read_all() if self._reranker is not None: + source = results results = self._reranker.rerank_fts(self._query, results) check_reranker_result(results) + results = _merge_query_snapshot(results, source) return self._finalize_blob_query_table( results, blob_mode=blob_mode, blob_sources=blob_sources ) @@ -2847,6 +2905,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): with_row_ids=True, offset=self._offset, ) + results = _merge_query_snapshot(results, fts_results, vector_results) return self._finish_hybrid_results( results, blob_mode=blob_mode, blob_sources=blob_sources ) @@ -3559,9 +3618,8 @@ class AsyncQueryBase(object): _query_is_plain_blob_scan(query) and getattr(self, "_reranker", None) is None and not _query_needs_lance_scanner(self) - and _supports_local_blob_fetch( - self._table, self._table._pushdown_operations - ) + and self._table._inner._is_native() + and _supports_blob_fetch(self._table, self._table._pushdown_operations) ) def select(self, columns: Union[List[str], dict[str, str]]) -> Self: @@ -3766,11 +3824,14 @@ class AsyncQueryBase(object): if blob_mode != "bytes": return blob_sources reject_unfetchable_blob_columns(schema, projection, blob_sources) - if blob_sources and not _supports_local_blob_fetch( + if blob_sources and not _supports_blob_fetch( self._table, self._table._pushdown_operations ): raise _blob_bytes_without_snapshot_error() - reject_unmatched_blob_outputs(await self._blob_output_schema(), blob_sources) + if self._table._inner._is_native(): + reject_unmatched_blob_outputs( + await self._blob_output_schema(), blob_sources + ) return blob_sources async def _blob_output_schema(self) -> pa.Schema: @@ -3783,9 +3844,21 @@ class AsyncQueryBase(object): return tbl if blob_mode == "descriptions": return strip_auto_row_ids(tbl, blob_sources) - return await replace_blob_columns_with_bytes( - tbl, blob_sources, self._table.fetch_blobs - ) + reject_unmatched_blob_outputs(tbl.schema, blob_sources) + if ( + self._table is not None + and not self._table._inner._is_native() + and blob_sources + ): + version, identity = _remote_blob_snapshot(tbl.schema) + + async def fetch_blobs(column, row_ids): + return await self._table._inner.fetch_blobs( + column, row_ids, version=version, snapshot_identity=identity + ) + else: + fetch_blobs = self._table.fetch_blobs + return await replace_blob_columns_with_bytes(tbl, blob_sources, fetch_blobs) async def to_arrow( self, @@ -3811,7 +3884,8 @@ class AsyncQueryBase(object): ``fetch_blobs`` reads. ``"descriptions"`` returns them without it. ``"bytes"`` returns payload bytes for supported row-addressable blob columns and raises for blobs inside lists. - LanceDB Enterprise and namespace query pushdown refuse ``"bytes"``. + Enterprise queries fetch bytes at the version reported by the query + response. Namespace query pushdown refuses ``"bytes"``. """ blob_sources = await self._validated_blob_sources(blob_mode) if blob_mode == "bytes" and blob_sources: @@ -3893,9 +3967,9 @@ class AsyncQueryBase(object): columns without pylance. Other local blob v2 queries return descriptors in ``"lazy"`` and ``"descriptions"`` modes, or payloads in ``"bytes"`` mode, through :meth:`to_arrow`. - LanceDB Enterprise blob queries support only ``"descriptions"`` until a - stable - table snapshot can bind the query and blob fetches. + Enterprise queries support ``"bytes"`` at the query response's + snapshot and ``"descriptions"`` without payload reads. Remote + ``"lazy"`` handles are not supported. **kwargs Forwarded to pyarrow.Table.to_pandas after query execution and optional flattening. @@ -3903,6 +3977,9 @@ class AsyncQueryBase(object): validate_blob_mode(blob_mode) if self._table is not None and not self._table._inner._is_native(): await self._validated_blob_sources(blob_mode) + if blob_mode == "bytes": + tbl = await self.to_arrow(timeout=timeout, blob_mode=blob_mode) + return flatten_columns(tbl, flatten).to_pandas(**kwargs) return await self._remote_to_pandas( flatten=flatten, timeout=timeout, blob_mode=blob_mode, **kwargs ) @@ -4508,7 +4585,9 @@ class AsyncFTSQuery(AsyncStandardQuery): ) results = pa.Table.from_batches(await reader.read_all(), reader.schema) if self._reranker: + source = results results = self._reranker.rerank_fts(self.get_query(), results) + results = _merge_query_snapshot(results, source) return AsyncRecordBatchReader(results, max_batch_length=max_batch_length) @@ -4835,7 +4914,9 @@ class AsyncVectorQuery(AsyncStandardQuery, AsyncVectorQueryBase): ) results = pa.Table.from_batches(await reader.read_all(), reader.schema) if self._reranker: + source = results results = self._reranker.rerank_vector(self._query_string, results) + results = _merge_query_snapshot(results, source) return AsyncRecordBatchReader(results, max_batch_length=max_batch_length) @@ -5027,6 +5108,7 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): ): result = result.drop(["_rowid"]) + result = _merge_query_snapshot(result, fts_results, vector_results) return AsyncRecordBatchReader(result, max_batch_length=max_batch_length) async def explain_plan(self, verbose: Optional[bool] = False): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 448fd6c39..561a0f900 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -5462,14 +5462,34 @@ def blob_descriptor_take_table(): @contextlib.contextmanager -def blob_remote_table(*, server_version=Version("0.5.0"), location=None): +def blob_remote_table( + *, + server_version=Version("0.5.0"), + location=None, + version_header="1", + framing="file", + requests=None, + state=None, + empty=False, + schema_version=None, + search_scores=False, + snapshot_supported=True, + identity_header="v1:" + "a" * 64, + verify_identity=True, +): + if state is None: + state = {"version": 1} + def handler(request): if request.path == "/v1/table/test/describe/": + content_len = int(request.headers.get("Content-Length", 0)) + body = json.loads(request.rfile.read(content_len)) if content_len else {} request.send_response(200) request.send_header("Content-Type", "application/json") request.send_header("phalanx-version", str(server_version)) request.end_headers() response = dict(BLOB_DESCRIBE_RESPONSE) + response["version"] = body.get("version") or state["version"] if location is not None: response["location"] = str(location) request.wfile.write(json.dumps(response).encode()) @@ -5502,9 +5522,22 @@ def blob_remote_table(*, server_version=Version("0.5.0"), location=None): elif request.path == "/v1/table/test/query/": content_len = int(request.headers.get("Content-Length", 0)) body = json.loads(request.rfile.read(content_len)) + if requests is not None: + requests.append(("query", body)) columns = body.get("columns") - assert columns in (None, ["id", "image"], ["id"], ["image"]) - if columns == ["id"]: + assert columns in ( + None, + ["id", "image"], + ["id"], + ["image"], + {"payload": "image"}, + ) + if isinstance(columns, dict): + response_table = blob_query_response_table().select(["image", "_rowid"]) + response_table = response_table.rename_columns(["payload", "_rowid"]) + if not body.get("with_row_id", False): + response_table = response_table.drop(["_rowid"]) + elif columns == ["id"]: response_table = blob_query_response_table().select(["id"]) else: if columns == ["image"]: @@ -5521,25 +5554,86 @@ def blob_remote_table(*, server_version=Version("0.5.0"), location=None): response_table = blob_query_response_table() if not body.get("with_row_id", False): response_table = response_table.drop(["_rowid"]) + if search_scores and not body.get("filter", "").startswith("_rowid IN"): + if body.get("vector"): + response_table = response_table.append_column( + "_distance", pa.array([0.1, 0.2, 0.3], type=pa.float32()) + ) + elif body.get("full_text_query"): + response_table = response_table.append_column( + "_score", pa.array([3.0, 2.0, 1.0], type=pa.float32()) + ) + if empty: + response_table = response_table.slice(0, 0) + if schema_version is not None: + response_table = response_table.replace_schema_metadata( + { + b"lancedb:query_version": schema_version.encode(), + b"lancedb:query_snapshot": ("v1:" + "f" * 64).encode(), + } + ) request.send_response(200) - request.send_header("Content-Type", "application/vnd.apache.arrow.file") - request.send_header("x-lancedb-version", "1") + request.send_header( + "Content-Type", f"application/vnd.apache.arrow.{framing}" + ) + header = ( + version_header(body) if callable(version_header) else version_header + ) + if header is not None: + request.send_header("x-lancedb-version", header) + if snapshot_supported: + request.send_header("x-lancedb-query-version", header) + identity = ( + identity_header(body) + if callable(identity_header) + else identity_header + ) + if identity is not None: + request.send_header("x-lancedb-query-snapshot", identity) request.end_headers() - with pa.ipc.new_file(request.wfile, response_table.schema) as writer: - writer.write_table(response_table) + writer_factory = pa.ipc.new_file if framing == "file" else pa.ipc.new_stream + with writer_factory(request.wfile, response_table.schema) as writer: + writer.write_table(response_table, max_chunksize=1) elif request.path == "/v1/table/test/fetch_blobs/": content_len = int(request.headers.get("Content-Length", 0)) body = json.loads(request.rfile.read(content_len)) assert body["column"] == "image" - assert body["row_ids"] == [10, 20, 30] + if requests is not None: + requests.append(("fetch_blobs", body)) + payloads = {10: b"alpha", 20: None, 30: b"gamma", 40: b""} + if state["version"] > 1 and body.get("version") != 1: + payloads = {row: b"new-version" for row in payloads} response_table = pa.table( - {"image": pa.array([b"alpha", None, b"gamma"], type=pa.large_binary())} + { + "image": pa.array( + [payloads[row] for row in body["row_ids"]], + type=pa.large_binary(), + ) + } ) request.send_response(200) request.send_header("Content-Type", "application/vnd.apache.arrow.stream") + if verify_identity and request.headers.get("x-lancedb-query-snapshot"): + request.send_header( + "x-lancedb-query-snapshot", + request.headers["x-lancedb-query-snapshot"], + ) request.end_headers() with pa.ipc.new_stream(request.wfile, response_table.schema) as writer: writer.write_table(response_table) + elif request.path == "/v1/table/test/insert/": + if request.headers.get("Transfer-Encoding") == "chunked": + while size := int(request.rfile.readline().split(b";")[0], 16): + request.rfile.read(size) + assert request.rfile.read(2) == b"\r\n" + assert request.rfile.read(2) == b"\r\n" + else: + request.rfile.read(int(request.headers.get("Content-Length", 0))) + state["version"] += 1 + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps({"version": state["version"]}).encode()) else: request.send_response(404) request.end_headers() @@ -5548,6 +5642,222 @@ def blob_remote_table(*, server_version=Version("0.5.0"), location=None): yield db.open_table("test") +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("framing", ["file", "stream"]) +@pytest.mark.parametrize("method", ["to_arrow", "to_list", "to_pandas", "to_batches"]) +async def test_remote_blob_bytes_conversion(use_async, framing, method): + if method == "to_pandas": + pytest.importorskip("pandas") + requests = [] + with blob_remote_table(framing=framing, requests=requests) as remote_table: + table = remote_table._table if use_async else remote_table + query = table.query() if use_async else table.search() + query = query.select(["id", "image"]) + if method == "to_batches": + kwargs = {"max_batch_length" if use_async else "batch_size": 1} + reader = query.to_batches(blob_mode="bytes", **kwargs) + if use_async: + reader = await reader + batches = [batch async for batch in reader] + else: + batches = list(reader) + assert all(batch.num_rows == 1 for batch in batches) + result = pa.Table.from_batches(batches, schema=reader.schema).to_pylist() + else: + result = getattr(query, method)(blob_mode="bytes") + if use_async: + result = await result + if method == "to_arrow": + assert result.schema.field("image").type == pa.large_binary() + result = result.to_pylist() + elif method == "to_pandas": + result = result.to_dict("records") + assert result == [ + {"id": 1, "image": b"alpha"}, + {"id": 2, "image": None}, + {"id": 3, "image": b"gamma"}, + ] + fetches = [body for path, body in requests if path == "fetch_blobs"] + assert fetches + assert all(body["version"] == 1 for body in fetches) + assert [row for body in fetches for row in body["row_ids"]] == [10, 20, 30] + descriptor_queries = [ + body + for path, body in requests + if path == "query" and body.get("filter", "").startswith("_rowid IN") + ] + assert all(body["version"] == 1 for body in descriptor_queries) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("framing", ["file", "stream"]) +async def test_remote_blob_bytes_keeps_query_snapshot_during_writes(use_async, framing): + requests = [] + state = {"version": 1} + with blob_remote_table( + framing=framing, requests=requests, state=state + ) as remote_table: + table = remote_table._table if use_async else remote_table + query = table.query() if use_async else table.search() + kwargs = {"max_batch_length" if use_async else "batch_size": 1} + reader = query.select(["id", "image"]).to_batches(blob_mode="bytes", **kwargs) + if use_async: + reader = await reader + # The query has selected v1; commit v2 before consuming any payloads. + write = table.add([{"id": 4, "image": b"new-version"}]) + if use_async: + await write + version = await table.version() if use_async else table.version + assert version == 2 + if use_async: + rows = [row async for batch in reader for row in batch.to_pylist()] + else: + rows = [row for batch in reader for row in batch.to_pylist()] + assert [row["image"] for row in rows] == [b"alpha", None, b"gamma"] + fetches = [body for path, body in requests if path == "fetch_blobs"] + assert all(body["version"] == 1 for body in fetches) + # Snapshot fetches must not check out the caller's live handle. + version = await table.version() if use_async else table.version + assert version == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("version_header", [None, "invalid", "0", "-1", str(2**64)]) +async def test_remote_blob_bytes_invalid_snapshot(use_async, version_header): + requests = [] + with blob_remote_table( + version_header=version_header, requests=requests, schema_version="1" + ) as remote_table: + table = remote_table._table if use_async else remote_table + query = table.query() if use_async else table.search() + with pytest.raises(NotImplementedError, match="x-lancedb-query-version"): + result = query.to_arrow(blob_mode="bytes") + if use_async: + await result + assert not any(path == "fetch_blobs" for path, _ in requests) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("framing", ["file", "stream"]) +async def test_remote_blob_bytes_empty_results(use_async, framing): + requests = [] + with blob_remote_table( + empty=True, framing=framing, requests=requests + ) as remote_table: + table = remote_table._table if use_async else remote_table + query = table.query() if use_async else table.search() + result = query.to_arrow(blob_mode="bytes") + if use_async: + result = await result + assert result.num_rows == 0 + assert result.schema.field("image").type == pa.large_binary() + assert result.column_names == ["id", "image"] + assert not any(path == "fetch_blobs" for path, _ in requests) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("kind", ["vector", "fts", "hybrid"]) +@pytest.mark.parametrize("rerank", [False, True]) +async def test_remote_blob_bytes_search_queries(use_async, kind, rerank): + with blob_remote_table(search_scores=True) as remote_table: + if use_async: + query = remote_table._table.query() + if kind in ("vector", "hybrid"): + query = query.nearest_to([1.0, 2.0]) + if kind in ("fts", "hybrid"): + query = query.nearest_to_text("hello") + elif kind == "vector": + query = remote_table.search([1.0, 2.0], vector_column_name="vector") + elif kind == "fts": + query = remote_table.search("hello", query_type="fts") + else: + query = ( + remote_table.search(query_type="hybrid", vector_column_name="vector") + .vector([1.0, 2.0]) + .text("hello") + ) + if rerank: + from lancedb.rerankers import RRFReranker + + class StripMetadataReranker(RRFReranker): + def rerank_vector(self, query, vector_results): + return vector_results.append_column( + "_relevance_score", pa.array([3.0, 2.0, 1.0]) + ).replace_schema_metadata(None) + + def rerank_fts(self, query, fts_results): + return self.rerank_vector(query, fts_results) + + def rerank_hybrid(self, query, vector_results, fts_results): + return ( + super() + .rerank_hybrid(query, vector_results, fts_results) + .replace_schema_metadata(None) + ) + + kwargs = {"query_string": "hello"} if kind == "vector" else {} + query = query.rerank(StripMetadataReranker(), **kwargs) + result = query.select(["id", "image"]).to_arrow(blob_mode="bytes") + if use_async: + result = await result + assert result["image"].to_pylist() == [b"alpha", None, b"gamma"] + assert "_rowid" not in result.column_names + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_remote_blob_bytes_hybrid_rejects_different_snapshots(use_async): + def header(body): + return "2" if body.get("full_text_query") else "1" + + requests = [] + with blob_remote_table( + search_scores=True, version_header=header, requests=requests + ) as remote_table: + if use_async: + query = ( + remote_table._table.query() + .nearest_to([1.0, 2.0]) + .nearest_to_text("hello") + ) + else: + query = ( + remote_table.search(query_type="hybrid", vector_column_name="vector") + .vector([1.0, 2.0]) + .text("hello") + ) + with pytest.raises( + NotImplementedError, match="inconsistent x-lancedb-query-version" + ): + result = query.select(["id", "image"]).to_arrow(blob_mode="bytes") + if use_async: + await result + assert not any(path == "fetch_blobs" for path, _ in requests) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("with_row_id", [False, True]) +async def test_remote_blob_bytes_alias_and_explicit_row_ids(use_async, with_row_id): + with blob_remote_table() as remote_table: + query = remote_table._table.query() if use_async else remote_table.search() + query = query.select({"payload": "image"}) + if with_row_id: + query = query.with_row_id() if use_async else query.with_row_id(True) + result = query.to_arrow(blob_mode="bytes") + if use_async: + result = await result + assert result["payload"].to_pylist() == [b"alpha", None, b"gamma"] + assert ("_rowid" in result.column_names) == with_row_id + if with_row_id: + assert result["_rowid"].to_pylist() == [10, 20, 30] + + def test_remote_blob_columns_and_fetch(): with blob_remote_table() as table: assert table.blob_columns() == ["image"] @@ -5569,7 +5879,7 @@ def test_sync_remote_blob_pandas_avoids_storage_access( raise AssertionError("remote storage must not be opened") monkeypatch.setattr(lance, "dataset", fail_dataset) - with blob_remote_table(location=location) as table: + with blob_remote_table(location=location, version_header=None) as table: monkeypatch.setattr(table, "to_lance", fail_dataset) query = table.search().select(["id", "image"]) if blob_mode == "descriptions": @@ -5582,8 +5892,11 @@ def test_sync_remote_blob_pandas_avoids_storage_access( NotImplementedError, match="stable table snapshot" ) as exc: query.to_pandas(blob_mode=blob_mode) - assert "fetch_blobs(column, row_ids)" in str(exc.value) - assert "with_row_id(True)" in str(exc.value) + if blob_mode == "bytes": + assert "x-lancedb-query-version" in str(exc.value) + else: + assert "fetch_blobs(column, row_ids)" in str(exc.value) + assert "with_row_id(True)" in str(exc.value) @pytest.mark.asyncio @@ -5667,7 +5980,7 @@ async def test_async_remote_blob_descriptions_to_pandas_uses_server(monkeypatch) @pytest.mark.asyncio @pytest.mark.parametrize("blob_mode", ["bytes", "lazy"]) async def test_async_remote_blob_pandas_rejects_unpinned_fetch(monkeypatch, blob_mode): - with blob_remote_table() as remote_table: + with blob_remote_table(version_header=None) as remote_table: table = remote_table._table async def fail_uri(): @@ -5692,8 +6005,8 @@ async def test_async_remote_blob_pandas_rejects_unpinned_fetch(monkeypatch, blob @pytest.mark.parametrize("method", ["to_arrow", "to_list"]) -def test_remote_search_bytes_refused_before_query(monkeypatch, method): - with blob_remote_table() as table: +def test_remote_search_bytes_requires_response_version(method): + with blob_remote_table(version_header=None) as table: query = table.search().select(["id", "image"]) descriptions = getattr(query, method)(blob_mode="descriptions") images = ( @@ -5703,19 +6016,15 @@ def test_remote_search_bytes_refused_before_query(monkeypatch, method): ) assert "_lance_row_id" not in images[0] - def fail_query(*args, **kwargs): - raise AssertionError("bytes must be refused before the query runs") - - monkeypatch.setattr(table, "_execute_query", fail_query) - with pytest.raises(NotImplementedError, match="stable table snapshot"): + with pytest.raises(NotImplementedError, match="x-lancedb-query-version"): getattr(query, method)(blob_mode="bytes") with pytest.raises(ValueError, match="blob_mode must be one of"): table.to_pandas(blob_mode="handles") @pytest.mark.asyncio -async def test_async_remote_to_arrow_bytes_refused_before_query(monkeypatch): - with blob_remote_table() as remote_table: +async def test_async_remote_to_arrow_bytes_requires_response_version(monkeypatch): + with blob_remote_table(version_header=None) as remote_table: table = remote_table._table async def fail_fetch(*args, **kwargs): @@ -5723,8 +6032,7 @@ async def test_async_remote_to_arrow_bytes_refused_before_query(monkeypatch): monkeypatch.setattr(table, "fetch_blobs", fail_fetch) query = table.query().select(["id", "image"]) - monkeypatch.setattr(query, "to_batches", fail_fetch) - with pytest.raises(NotImplementedError, match="stable table snapshot"): + with pytest.raises(NotImplementedError, match="x-lancedb-query-version"): await query.to_arrow(blob_mode="bytes") with pytest.raises(NotImplementedError, match="stable table snapshot"): await table.to_arrow(blob_mode="bytes") @@ -5821,7 +6129,7 @@ def test_remote_blob_query_to_pandas(): @pytest.mark.parametrize("blob_mode", ["default", "lazy", "bytes"]) def test_remote_blob_query_to_pandas_rejects_unpinned_fetch(monkeypatch, blob_mode): - with blob_remote_table() as table: + with blob_remote_table(version_header=None) as table: def fail_fetch(*args, **kwargs): raise AssertionError("blob fetch must not use an unpinned remote table") @@ -6884,3 +7192,102 @@ def test_remote_optimize_rejects_explicit_pruning_before_submitting(): table.optimize(delete_unverified=True) assert server.tickets == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("framing", ["file", "stream"]) +async def test_remote_blob_bytes_rejects_legacy_watermark(use_async, framing): + requests = [] + with blob_remote_table( + framing=framing, + snapshot_supported=False, + version_header="1", + schema_version="1", + requests=requests, + ) as remote_table: + table = remote_table._table if use_async else remote_table + query = table.query() if use_async else table.search() + descriptors = query.to_arrow(blob_mode="descriptions") + if use_async: + descriptors = await descriptors + assert descriptors.num_rows == 3 + assert pa.types.is_struct(descriptors.schema.field("image").type) + query = table.query() if use_async else table.search() + with pytest.raises(NotImplementedError, match="x-lancedb-query-version"): + result = query.to_arrow(blob_mode="bytes") + if use_async: + await result + assert not any(path == "fetch_blobs" for path, _ in requests) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("identity_header", [None, "", "v2:" + "a" * 64, "v1:broken"]) +async def test_remote_blob_bytes_requires_trustworthy_identity( + use_async, identity_header +): + requests = [] + with blob_remote_table( + identity_header=identity_header, schema_version="1", requests=requests + ) as remote_table: + table = remote_table._table if use_async else remote_table + query = table.query() if use_async else table.search() + # Ordinary descriptors still work, even with forged Arrow schema metadata. + result = query.to_arrow(blob_mode="descriptions") + if use_async: + result = await result + assert result.num_rows == 3 + with pytest.raises(NotImplementedError, match="x-lancedb-query-snapshot"): + result = query.to_arrow(blob_mode="bytes") + if use_async: + await result + assert not any(path == "fetch_blobs" for path, _ in requests) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_remote_blob_bytes_rejects_unverified_serving_node(use_async): + with blob_remote_table(verify_identity=False) as remote_table: + table = remote_table._table if use_async else remote_table + query = table.query() if use_async else table.search() + with pytest.raises( + NotImplementedError, match="did not verify the query snapshot identity" + ): + result = query.to_arrow(blob_mode="bytes") + if use_async: + await result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_remote_blob_bytes_hybrid_rejects_same_version_different_incarnations( + use_async, +): + requests = [] + + def identity(body): + return "v1:" + ("b" if body.get("full_text_query") else "a") * 64 + + with blob_remote_table( + search_scores=True, identity_header=identity, requests=requests + ) as remote_table: + if use_async: + query = ( + remote_table._table.query() + .nearest_to([1.0, 2.0]) + .nearest_to_text("hello") + ) + else: + query = ( + remote_table.search(query_type="hybrid", vector_column_name="vector") + .vector([1.0, 2.0]) + .text("hello") + ) + with pytest.raises( + NotImplementedError, match="inconsistent x-lancedb-query-version" + ): + result = query.select(["id", "image"]).to_arrow(blob_mode="bytes") + if use_async: + await result + assert not any(path == "fetch_blobs" for path, _ in requests) diff --git a/python/src/table.rs b/python/src/table.rs index e76bae00a..31cf8d25c 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1978,14 +1978,28 @@ impl Table { } /// Read blob bytes for `row_ids` from blob v2 column `column`. - #[pyo3(signature = (column, row_ids))] + #[pyo3(signature = (column, row_ids, *, version=None, snapshot_identity=None))] pub fn fetch_blobs( self_: PyRef<'_, Self>, column: String, row_ids: Vec, + version: Option, + snapshot_identity: Option, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { + let inner = match (version, snapshot_identity) { + (None, None) => inner, + (Some(version), Some(identity)) => inner + .query_snapshot_with_identity(version, &identity) + .await + .infer_error()?, + _ => { + return Err(PyValueError::new_err( + "deferred blob reads require both query version and snapshot identity", + )); + } + }; let blobs: LargeBinaryArray = inner .fetch_blobs(column.as_str(), &row_ids) .await diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index cc495a83f..d560a4802 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -95,6 +95,34 @@ const MIN_VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-ve const MIN_TIMESTAMP_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-timestamp"); const MIN_READ_VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-read-version"); const VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-version"); +const QUERY_SNAPSHOT_HEADER: &str = "x-lancedb-query-snapshot"; + +fn valid_query_snapshot_identity(value: &str) -> bool { + value.strip_prefix("v1:").is_some_and(|digest| { + digest.len() == 64 + && digest + .bytes() + .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) + }) +} + +fn check_query_snapshot_response( + expected: Option<&str>, + response: &reqwest::Response, +) -> Result<()> { + if let Some(expected) = expected + && response + .headers() + .get(QUERY_SNAPSHOT_HEADER) + .and_then(|value| value.to_str().ok()) + != Some(expected) + { + return Err(Error::NotSupported { message: "Blob response did not verify the query snapshot identity; upgrade every serving node".into() }); + } + Ok(()) +} + +const QUERY_VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-query-version"); /// Asks for stream framing, which is decoded as it arrives. Servers that do not /// stream query results ignore this and answer with file framing. const QUERY_ACCEPT: &str = @@ -613,6 +641,7 @@ pub struct RemoteTable { /// Stamped onto every branch-accepting request so reads and writes resolve /// on the branch's own version chain rather than main's. branch: Option, + expected_query_snapshot: Option, } impl std::fmt::Debug for RemoteTable { @@ -855,6 +884,7 @@ impl RemoteTable { wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, + expected_query_snapshot: None, } } @@ -891,6 +921,7 @@ impl RemoteTable { wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch, + expected_query_snapshot: None, } } @@ -1170,20 +1201,37 @@ impl RemoteTable { .and_then(|value| value.to_str().ok()) .map(str::to_owned); let framing = resolve_arrow_ipc_framing(content_type.as_deref(), request_id)?; + let table_version = response + .headers() + .get(&VERSION_HEADER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .filter(|version| *version > 0); - match framing { + // Legacy version headers are freshness watermarks, not proof that the + // query and a later Blob read use the same immutable dataset. + let query_version = response + .headers() + .get(&QUERY_VERSION_HEADER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .filter(|version| *version > 0); + + let query_identity = response + .headers() + .get(QUERY_SNAPSHOT_HEADER) + .and_then(|value| value.to_str().ok()) + .filter(|value| valid_query_snapshot_identity(value)) + .map(str::to_owned); + + let stream = match framing { ArrowIpcFraming::Stream => { - let table_version = response - .headers() - .get(&VERSION_HEADER) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.parse().ok()); ipc_stream::read_ipc_stream( response.bytes_stream(), request_id.to_string(), table_version, ) - .await + .await? } ArrowIpcFraming::File => { // File framing keeps its footer at the end, so the whole body is @@ -1192,9 +1240,21 @@ impl RemoteTable { let reader = FileReader::try_new(Cursor::new(body), None)?; let schema = reader.schema(); let stream = futures::stream::iter(reader).map_err(DataFusionError::from); - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + Box::pin(RecordBatchStreamAdapter::new(schema, stream)) as SendableRecordBatchStream } + }; + // Bind subsequent blob reads to the response's snapshot, never to a + // version sampled from the live table before or after the query. + let mut metadata = stream.schema().metadata().clone(); + metadata.remove("lancedb:query_version"); + metadata.remove("lancedb:query_snapshot"); + if let Some(identity) = query_identity { + metadata.insert("lancedb:query_snapshot".into(), identity); } + if let Some(version) = query_version { + metadata.insert("lancedb:query_version".into(), version.to_string()); + } + Ok(with_stream_metadata(stream, metadata)) } fn apply_query_params( @@ -1620,9 +1680,37 @@ impl RemoteTable { streams.await? }; + if let Some(first) = streams.first() { + let schema = first.schema(); + if streams.iter().skip(1).any(|stream| { + stream.schema().metadata().get("lancedb:query_version") + != schema.metadata().get("lancedb:query_version") + }) { + return Err(Error::Runtime { + message: "Remote query responses refer to different table snapshots; retry the query against a fixed version".into(), + }); + } + } + + let identities_agree = streams.first().is_none_or(|first| { + streams.iter().skip(1).all(|stream| { + stream.schema().metadata().get("lancedb:query_snapshot") + == first.schema().metadata().get("lancedb:query_snapshot") + }) + }); Ok(streams .into_iter() .map(|stream| { + // Ordinary descriptor queries remain usable during an upgrade. + // Deferred bytes require every leg's identity, so never preserve + // the first leg's capability when another leg lacks or differs. + let stream = if identities_agree { + stream + } else { + let mut metadata = stream.schema().metadata().clone(); + metadata.remove("lancedb:query_snapshot"); + with_stream_metadata(stream, metadata) + }; let stream = restore_projection_order(stream, &query.base().select); match deadline { Some((deadline, timeout)) => { @@ -1673,6 +1761,26 @@ impl RemoteTable { } } +fn with_stream_metadata( + stream: SendableRecordBatchStream, + metadata: HashMap, +) -> SendableRecordBatchStream { + let schema = Arc::new( + stream + .schema() + .as_ref() + .clone() + .with_metadata(metadata.clone()), + ); + let batches = stream.map(move |batch| { + batch.map(|mut batch| { + batch.schema_metadata_mut().clone_from(&metadata); + batch + }) + }); + Box::pin(RecordBatchStreamAdapter::new(schema, batches)) +} + fn restore_projection_order( stream: Pin>, select: &Select, @@ -1921,6 +2029,7 @@ mod test_utils { wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, + expected_query_snapshot: None, } } @@ -1948,6 +2057,7 @@ mod test_utils { wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, + expected_query_snapshot: None, } } @@ -1984,6 +2094,7 @@ mod test_utils { wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, + expected_query_snapshot: None, } } } @@ -2570,6 +2681,23 @@ impl BaseTable for RemoteTable { } } } + async fn query_snapshot_with_identity( + &self, + version: u64, + identity: &str, + ) -> Result> { + if version == 0 || !valid_query_snapshot_identity(identity) { + return Err(Error::InvalidInput { + message: "invalid or unsupported query snapshot identity".into(), + }); + } + let mut snapshot = self.with_branch(self.branch.clone()); + *snapshot.version.write().await = Some(version); + snapshot.reset_freshness(None, true); + snapshot.expected_query_snapshot = Some(identity.to_owned()); + Ok(Arc::new(snapshot)) + } + async fn query_snapshot(&self) -> Result> { let description = self.describe().await?; let TableDescription { @@ -6704,8 +6832,15 @@ mod tests { assert_eq!(request_sizes.lock().unwrap().as_slice(), &[1024, 1]); } + #[rstest::rstest] + #[case(false)] + #[case(true)] #[tokio::test] - async fn test_fetch_blobs_splits_byte_limited_requests_and_reads_large_blob_by_range() { + async fn test_fetch_blobs_splits_byte_limited_requests_and_reads_large_blob_by_range( + #[case] guarded: bool, + ) { + const IDENTITY: &str = + "v1:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; // Simulate a lower byte cap so this test exercises the same 400 response // without allocating 64 MiB of blob data. let requests = Arc::new(std::sync::Mutex::new(Vec::new())); @@ -6714,8 +6849,15 @@ mod tests { "my_table", semver::Version::new(0, 5, 0), move |request| { + if guarded { + assert_eq!(request.headers()[QUERY_SNAPSHOT_HEADER], IDENTITY); + } let path = request.url().path(); if path == "/v1/table/my_table/describe/" { + assert!( + !guarded, + "a later describe cannot establish the producing identity" + ); return http::Response::builder() .status(200) .body(br#"{"version":42,"schema":{"fields":[]}}"#.to_vec()) @@ -6766,6 +6908,7 @@ mod tests { return http::Response::builder() .status(200) .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .header(QUERY_SNAPSHOT_HEADER, IDENTITY) .body(write_ipc_stream_uncompressed(&batch)) .unwrap(); } @@ -6781,12 +6924,14 @@ mod tests { match range { "bytes=0-0" => http::Response::builder() .status(206) + .header(QUERY_SNAPSHOT_HEADER, IDENTITY) .header(reqwest::header::CONTENT_RANGE, "bytes 0-0/10") .header(VERSION_HEADER, "42") .body(b"0".to_vec()) .unwrap(), "bytes=0-9" => http::Response::builder() .status(206) + .header(QUERY_SNAPSHOT_HEADER, IDENTITY) .header(reqwest::header::CONTENT_RANGE, "bytes 0-9/10") .body(b"0123456789".to_vec()) .unwrap(), @@ -6795,6 +6940,14 @@ mod tests { }, ); + let table = if guarded { + table + .query_snapshot_with_identity(42, IDENTITY) + .await + .unwrap() + } else { + table + }; let blobs = table .fetch_blobs("image", &[10, 20, 30, 40, 20]) .await diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index 345b731c5..a0051c955 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -33,8 +33,8 @@ use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; use crate::table::{AnyQuery, BaseTable}; use super::{ - FreshnessHeaders, FreshnessState, ReadSnapshot, RemoteTable, VERSION_HEADER, - freshness_headers_snapshot, + FreshnessHeaders, FreshnessState, QUERY_SNAPSHOT_HEADER, ReadSnapshot, RemoteTable, + VERSION_HEADER, check_query_snapshot_response, freshness_headers_snapshot, }; // The Cloud route rejects larger row-id lists. Its separate 64 MiB byte limit @@ -93,6 +93,7 @@ struct TableBlobRangeRequester { version: Option, branch: Option, source_uri: Option, + expected_query_snapshot: Option, source_commit_micros: Option, row_is_address: bool, freshness: Arc>, @@ -125,6 +126,9 @@ impl BlobRangeRequester for TableBlobRangeRequester { .apply(self.client.get(&self.path)) .header(header::RANGE, range_header) }; + if let Some(identity) = &self.expected_query_snapshot { + request = request.header(QUERY_SNAPSHOT_HEADER, identity); + } if let Some(version) = version.or(self.version) { request = request.query(&[("version", version)]); } @@ -140,6 +144,10 @@ impl BlobRangeRequester for TableBlobRangeRequester { request = request.query(&[("source_commit_micros", timestamp)]); } let (request_id, response) = self.client.send_with_retry(request, None, true).await?; + if response.status().is_success() || response.status() == StatusCode::RANGE_NOT_SATISFIABLE + { + check_query_snapshot_response(self.expected_query_snapshot.as_deref(), &response)?; + } // Preserve 416 size-probe responses so the caller can detect empty blobs. if response.status() == StatusCode::RANGE_NOT_SATISFIABLE && matches!(mode, RangeRequestMode::SizeProbe) @@ -531,7 +539,11 @@ impl RemoteTable { // A dotted path needs client-side validation because the server may // report an internal type mismatch when the path crosses a list. Use // an already cached schema for top-level list fields as well. - if column.contains('.') || self.schema_cache.try_get().is_some() { + // Guarded readers must validate the producing identity on the server + // before consulting any schema resolved from a later logical binding. + if self.expected_query_snapshot.is_none() + && (column.contains('.') || self.schema_cache.try_get().is_some()) + { let schema = self.schema().await?; crate::blob::reject_list_blob_column(schema.as_ref(), column)?; } @@ -539,6 +551,9 @@ impl RemoteTable { } async fn explain_list_blob_bad_request(&self, column: &str, error: Error) -> Error { + if self.expected_query_snapshot.is_some() { + return error; + } let is_bad_request = match &error { Error::Http { status_code: Some(code), @@ -648,9 +663,17 @@ impl RemoteTable { .client .post(&format!("/v1/table/{}/fetch_blobs/", self.identifier)) .json(&body); + let request = if let Some(identity) = &self.expected_query_snapshot { + request.header(QUERY_SNAPSHOT_HEADER, identity) + } else { + request + }; let (request_id, response) = self .send_with_freshness(request, true, read_snapshot.freshness) .await?; + if response.status().is_success() { + check_query_snapshot_response(self.expected_query_snapshot.as_deref(), &response)?; + } let mut stream = match self.read_arrow_response(&request_id, response).await { Ok(stream) => stream, Err(error) => return Err(self.explain_list_blob_bad_request(column, error).await), @@ -1023,6 +1046,7 @@ impl RemoteTable { version: read_snapshot.version, branch: source.map_or_else(|| self.branch.clone(), |source| source.branch.clone()), source_uri: source.map(|source| source.uri.clone()), + expected_query_snapshot: self.expected_query_snapshot.clone(), source_commit_micros: source.map(|source| source.committed_at_micros), row_is_address: source.is_some_and(|source| source.row_is_address), freshness: Arc::new(std::sync::Mutex::new(read_snapshot.freshness_state)), diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 18a5b7ddf..db68ab290 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -832,6 +832,16 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "blob_columns is not supported on this table type".into(), }) } + /// Bind a remote payload reader to an identity from its producing query. + async fn query_snapshot_with_identity( + &self, + _version: u64, + _identity: &str, + ) -> Result> { + Err(Error::NotSupported { + message: "query snapshot identities are not supported on this table type".into(), + }) + } /// Materialize blob bytes for the given row ids. See [`Table::fetch_blobs`]. async fn fetch_blobs(&self, _column: &str, _row_ids: &[u64]) -> Result { Err(Error::NotSupported { @@ -1352,6 +1362,19 @@ impl Table { }) } + /// Create an independent remote handle constrained by a query response. + #[doc(hidden)] + pub async fn query_snapshot_with_identity(&self, version: u64, identity: &str) -> Result { + Ok(Self { + inner: self + .inner + .query_snapshot_with_identity(version, identity) + .await?, + database: self.database.clone(), + embedding_registry: self.embedding_registry.clone(), + }) + } + /// Count the number of rows in this dataset. /// /// # Arguments