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