diff --git a/docs/web/management/table.mdx b/docs/web/management/table.mdx index 4e87bcbb1..182a80ec3 100644 --- a/docs/web/management/table.mdx +++ b/docs/web/management/table.mdx @@ -447,8 +447,8 @@ above: In the Python SDK, `Table.to_pandas()`, `Table.to_arrow()`, and the query builder -methods `to_pandas()`, `to_arrow()`, and `to_list()` accept a `blob_mode` -argument: +methods `to_pandas()`, `to_arrow()`, `to_list()`, and `to_batches()` accept a +`blob_mode` argument. Sync and async query builders support the same modes: | `blob_mode` | `to_pandas()` | `to_arrow()` and `to_list()` | | --- | --- | --- | @@ -456,10 +456,11 @@ argument: | `"descriptions"` | Returns descriptor dicts with `kind`, `position`, `size`, `blob_id`, and `blob_uri`. | Returns the 5-field descriptor without `_lance_row_id`. | | `"bytes"` | Materializes each blob as `bytes` in the DataFrame. | Materializes each blob as `large_binary` in Arrow or `bytes` in Python dicts. | -`to_pandas(blob_mode="lazy")` and `to_pandas(blob_mode="bytes")` require -pylance. Without it they raise `ImportError`. `to_arrow()`, -`to_list(blob_mode="bytes")`, `to_pandas(blob_mode="descriptions")`, and the -fetch methods do not. +For supported local v2 blob projections, `to_pandas(blob_mode="bytes")` and +`to_pandas(blob_mode="descriptions")` use native Arrow conversion without +pylance. Pandas conversions that need the Lance scanner, including plain-scan +`"lazy"` handles and legacy blobs, still require pylance. `to_arrow()`, +`to_list()`, `to_batches()`, and the fetch methods do not require pylance. `blob_mode="bytes"` requires a local table snapshot, so remote and in-memory tables raise `NotImplementedError`. On a remote table, `Table.to_pandas()` and @@ -475,9 +476,34 @@ them out of the projection. Extra keyword arguments on `to_pandas()`, such as `split_blocks` and `self_destruct`, are forwarded to PyArrow. +###### Stream blob results {#multimodal-blob-batches} + +`to_batches(blob_mode="bytes")` fetches blob v2 payloads one batch at a time +and returns them as Arrow `large_binary` values. Iterate over the reader and +call `batch.to_pandas()` to convert each batch to a DataFrame without collecting +all payloads at once. + +For sync queries, call `query.to_batches(batch_size=64, blob_mode="bytes")` +and iterate with `for`. For async queries, call +`await query.to_batches(max_batch_length=64, blob_mode="bytes")` and iterate +with `async for`. Sync full-text and hybrid queries do not support +`to_batches()`. + +The default `"lazy"` mode returns descriptors and keeps automatic `_rowid` +columns visible for manual fetches. `"descriptions"` returns the five-field +descriptors. Both `"bytes"` and `"descriptions"` hide automatic row IDs; +explicitly requested row IDs remain visible. + +Bytes mode has the same restrictions as `to_arrow(blob_mode="bytes")`, +including no legacy blobs, blobs inside lists, remote tables, or namespace +`QueryTable` pushdown. 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} -These examples require pylance. The query example expects `tbl_async`, +The lazy-handle calls below require pylance. The query example expects `tbl_async`, `query_vector`, and a `video` column, which the `photos` table does not define. diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index 46462553c..010c92009 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -280,7 +280,11 @@ async def replace_v2_blob_columns_with_bytes( row_ids = _hit_row_ids(tbl, present) for output_name, source_name in present.items(): tbl = _set_blob_column( - tbl, output_name, await fetch_blobs(source_name, row_ids) + tbl, + output_name, + await fetch_blobs(source_name, row_ids) + if row_ids + else pa.array([], type=pa.large_binary()), ) return tbl @@ -295,7 +299,13 @@ def replace_v2_blob_columns_with_bytes_sync( return tbl row_ids = _hit_row_ids(tbl, present) for output_name, source_name in present.items(): - tbl = _set_blob_column(tbl, output_name, fetch_blobs(source_name, row_ids)) + tbl = _set_blob_column( + tbl, + output_name, + fetch_blobs(source_name, row_ids) + if row_ids + else pa.array([], type=pa.large_binary()), + ) return tbl diff --git a/python/python/lancedb/arrow.py b/python/python/lancedb/arrow.py index 404efafea..c5a278ead 100644 --- a/python/python/lancedb/arrow.py +++ b/python/python/lancedb/arrow.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors from functools import singledispatch -from typing import List, Optional, Tuple, Union +from typing import AsyncIterator, List, Optional, Tuple, Union from lancedb.pydantic import LanceModel, model_to_dict import pyarrow as pa @@ -30,6 +30,7 @@ class AsyncRecordBatchReader: The schema of the batches produced by the stream. Accessing the schema does not consume any data from the stream """ + self._inner: AsyncIterator[pa.RecordBatch] if isinstance(inner, pa.Table): self._inner = self._async_iter_from_table(inner, max_batch_length) self.schema: pa.Schema = inner.schema @@ -39,6 +40,16 @@ class AsyncRecordBatchReader: else: raise TypeError("inner must be a RecordBatchStream or a Table") + @classmethod + def _from_batches( + cls, schema: pa.Schema, batches: AsyncIterator[pa.RecordBatch] + ) -> "AsyncRecordBatchReader": + """Wrap a transformed stream without consuming its first batch.""" + reader = cls.__new__(cls) + reader.schema = schema + reader._inner = batches + return reader + async def read_all(self) -> List[pa.RecordBatch]: """ Read all the record batches from the stream diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index b2093712e..03560ec31 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -104,11 +104,13 @@ class _LanceScanner(Protocol): def _unsupported_blob_pandas_error(reason: str) -> RuntimeError: + reason = reason.rstrip(".") + reason = reason[0].upper() + reason[1:] return RuntimeError( - "blob columns require Lance native scanner conversion for query " - f"to_pandas(), but {reason}. Use a plain scan query, pass " - "blob_mode='bytes' or blob_mode='descriptions', or remove blob columns " - "from the projection." + "Blob columns require Lance native scanner conversion for this query's " + f"to_pandas(). {reason}. For supported blob v2 queries, pass " + "blob_mode='bytes' or blob_mode='descriptions'. Otherwise, install " + "pylance for scanner conversion or remove blob columns from the projection." ) @@ -1139,9 +1141,10 @@ class LanceQueryBuilder(ABC): The maximum time to wait for the query to complete. If None, wait indefinitely. blob_mode: str, default "lazy" - Controls how blob columns are returned. Eligible local plain scans - use Lance's scanner, which returns lazy blob handles in ``"lazy"`` - mode and also supports legacy blob columns. Other local blob v2 + Controls how blob columns are returned. Local blob v2 queries support + ``"bytes"`` and ``"descriptions"`` without pylance. Eligible local + plain scans use pylance's scanner for lazy blob handles in ``"lazy"`` + mode and for legacy blob columns. Other local blob v2 queries return descriptors in ``"lazy"`` and ``"descriptions"`` modes, or payloads in ``"bytes"`` mode, through :meth:`to_arrow`. Remote blob queries support only ``"descriptions"`` until a stable @@ -1219,15 +1222,16 @@ class LanceQueryBuilder(ABC): batch_size: Optional[int] = None, *, timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", ) -> pa.RecordBatchReader: """ Execute the query and return the results as a pyarrow [RecordBatchReader](https://arrow.apache.org/docs/python/generated/pyarrow.RecordBatchReader.html) - For v2 blob projections, ``to_batches`` keeps the auto ``_rowid`` - column visible so batch consumers can call ``fetch_blobs``. Use - ``to_arrow``, ``to_list``, or ``to_pandas`` if you want LanceDB to hide - auto row ids in the final collected result. + For v2 blob projections, the default ``blob_mode="lazy"`` keeps the + auto ``_rowid`` column visible so batch consumers can call + ``fetch_blobs``. The ``"bytes"`` and ``"descriptions"`` modes hide + automatic row ids. Parameters ---------- @@ -1236,6 +1240,16 @@ class LanceQueryBuilder(ABC): timeout: Optional[timedelta] The maximum time to wait for the query to complete. If None, wait indefinitely. + blob_mode: str, default "lazy" + ``"lazy"`` keeps blob descriptors and auto row ids for manual fetches. + ``"descriptions"`` omits auto row ids. ``"bytes"`` fetches blob v2 + payloads one batch at a time, without requiring pylance. Bytes mode + has the same restrictions as :meth:`to_arrow`. + + Examples + -------- + Stream blob payloads with ``query.to_batches(blob_mode="bytes")`` and + convert each batch separately with ``batch.to_pandas()``. """ raise NotImplementedError @@ -1499,6 +1513,28 @@ class LanceQueryBuilder(ABC): finalized = finalized.drop_columns(["_rowid"]) return finalized + def _batches_in_mode( + self, + reader: pa.RecordBatchReader, + blob_mode: BlobMode, + blob_sources: dict[str, str], + ) -> pa.RecordBatchReader: + if blob_mode == "lazy": + return reader + + def convert(tbl: pa.Table) -> pa.Table: + return self._finalize_blob_query_table( + tbl, blob_mode=blob_mode, blob_sources=blob_sources + ) + + schema = convert(pa.Table.from_batches([], schema=reader.schema)).schema + + def batches(): + for batch in reader: + yield from convert(pa.Table.from_batches([batch])).to_batches() + + return pa.RecordBatchReader.from_batches(schema, batches()) + def with_row_address(self, with_row_address: bool = True) -> Self: """Set whether to return row addresses. @@ -1717,7 +1753,18 @@ class LanceQueryBuilder(ABC): get_uri_scheme(self._table._dataset_path) == "memory" ): return None - dataset = self._table.to_lance() + # Blob v2 payloads and descriptions use LanceDB's native fetch path. + # Only lazy handles, legacy blobs, and scanner-only options need pylance. + if ( + blob_mode != "lazy" + and not _query_needs_lance_scanner(self) + and blob_v2_projection_sources(self._table.schema, query.columns) + ): + return None + try: + dataset = self._table.to_lance() + except ImportError as err: + raise _unsupported_blob_pandas_error(str(err)) from err blob_auto_row_id = self._blob_auto_row_id_enabled() blob_sources = ( blob_v2_projection_sources(self._table.schema, query.columns) @@ -2076,6 +2123,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): batch_size: Optional[int] = None, *, timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", ) -> pa.RecordBatchReader: """ Execute the query and return the result as a RecordBatchReader object. @@ -2087,11 +2135,15 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): timeout: timedelta, default None The maximum time to wait for the query to complete. If None, wait indefinitely. + blob_mode: str, default "lazy" + Controls blob conversion as described in + :meth:`LanceQueryBuilder.to_batches`. Returns ------- pa.RecordBatchReader """ + blob_sources = self._validated_blob_sources(blob_mode) vector = self._query if isinstance(self._query, list) else self._query.tolist() if isinstance(vector[0], np.ndarray): vector = [v.tolist() for v in vector] @@ -2108,7 +2160,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): result_set.schema, result_set.to_batches() ) - return result_set + return self._batches_in_mode(result_set, blob_mode, blob_sources) def where( self, where: Union[str, Expr], prefilter: bool = None @@ -2323,7 +2375,12 @@ class LanceFtsQueryBuilder(LanceQueryBuilder): ) def to_batches( - self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None + self, + /, + batch_size: Optional[int] = None, + timeout: Optional[timedelta] = None, + *, + blob_mode: BlobMode = "lazy", ): raise NotImplementedError("to_batches on an FTS query") @@ -2379,10 +2436,19 @@ class LanceEmptyQueryBuilder(LanceQueryBuilder): return self._table._output_schema(query) def to_batches( - self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None + self, + /, + batch_size: Optional[int] = None, + timeout: Optional[timedelta] = None, + *, + blob_mode: BlobMode = "lazy", ) -> pa.RecordBatchReader: + blob_sources = self._validated_blob_sources(blob_mode) query = self._query_for_scan() - return self._table._execute_query(query, batch_size=batch_size, timeout=timeout) + reader = self._table._execute_query( + query, batch_size=batch_size, timeout=timeout + ) + return self._batches_in_mode(reader, blob_mode, blob_sources) def rerank(self, reranker: Reranker) -> LanceEmptyQueryBuilder: """Rerank the results using the specified reranker. @@ -2676,7 +2742,12 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): return results def to_batches( - self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None + self, + /, + batch_size: Optional[int] = None, + timeout: Optional[timedelta] = None, + *, + blob_mode: BlobMode = "lazy", ): raise NotImplementedError("to_batches not yet supported on a hybrid query") @@ -3279,6 +3350,7 @@ class AsyncQueryBase(object): *, max_batch_length: Optional[int] = None, timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", ) -> AsyncRecordBatchReader: """ Execute the query and return the results as an Apache Arrow RecordBatchReader. @@ -3295,7 +3367,47 @@ class AsyncQueryBase(object): The maximum time to wait for the query to complete. If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. + blob_mode: str, default "lazy" + ``"lazy"`` keeps blob descriptors and auto row ids for manual fetches. + ``"descriptions"`` omits auto row ids. ``"bytes"`` fetches blob v2 + payloads one batch at a time, without requiring pylance. Bytes mode + has the same restrictions as :meth:`to_arrow`. + + Examples + -------- + Use ``async for batch in await query.to_batches(blob_mode="bytes")`` + to stream blob payloads and convert each batch with ``batch.to_pandas()``. """ + blob_sources = await self._validated_blob_sources(blob_mode) + reader = await self._execute_batches( + max_batch_length=max_batch_length, timeout=timeout + ) + if blob_mode == "lazy": + return reader + + async def convert(tbl: pa.Table) -> pa.Table: + tbl = self._finalize_blob_query_table(tbl) + tbl = await self._apply_arrow_blob_mode(tbl, blob_mode, blob_sources) + if not self._user_requested_row_id() and "_rowid" in tbl.column_names: + tbl = tbl.drop_columns(["_rowid"]) + return tbl + + schema = (await convert(pa.Table.from_batches([], schema=reader.schema))).schema + + async def batches(): + async for batch in reader: + tbl = await convert(pa.Table.from_batches([batch])) + for converted in tbl.to_batches(): + yield converted + + return AsyncRecordBatchReader._from_batches(schema, batches()) + + async def _execute_batches( + self, + *, + max_batch_length: Optional[int] = None, + timeout: Optional[timedelta] = None, + ) -> AsyncRecordBatchReader: await self._maybe_add_blob_row_id() return AsyncRecordBatchReader( await self._inner.execute( @@ -3444,9 +3556,10 @@ class AsyncQueryBase(object): If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. blob_mode: str, default "lazy" - Controls how blob columns are returned. Eligible local plain scans - use Lance's scanner, which returns lazy blob handles in ``"lazy"`` - mode and also supports legacy blob columns. Other local blob v2 + Controls how blob columns are returned. Local blob v2 queries support + ``"bytes"`` and ``"descriptions"`` without pylance. Eligible local + plain scans use pylance's scanner for lazy blob handles in ``"lazy"`` + mode and for legacy blob columns. Other local blob v2 queries return descriptors in ``"lazy"`` and ``"descriptions"`` modes, or payloads in ``"bytes"`` mode, through :meth:`to_arrow`. Remote blob queries support only ``"descriptions"`` until a stable @@ -3536,7 +3649,16 @@ class AsyncQueryBase(object): get_uri_scheme(await self._table.uri()) == "memory" ): return None - dataset = await self._table.to_lance() + if ( + blob_mode != "lazy" + and not _query_needs_lance_scanner(self) + and blob_v2_projection_sources(schema, query.columns) + ): + return None + try: + dataset = await self._table.to_lance() + except ImportError as err: + raise _unsupported_blob_pandas_error(str(err)) from err scanner = dataset.scanner( **_scanner_kwargs_for_query( query, @@ -4036,13 +4158,13 @@ class AsyncFTSQuery(AsyncStandardQuery): self._table, ) - async def to_batches( + async def _execute_batches( self, *, max_batch_length: Optional[int] = None, timeout: Optional[timedelta] = None, ) -> AsyncRecordBatchReader: - reader = await super().to_batches(timeout=timeout) + reader = await super()._execute_batches(timeout=timeout) results = pa.Table.from_batches(await reader.read_all(), reader.schema) if self._reranker: results = self._reranker.rerank_fts(self.get_query(), results) @@ -4327,13 +4449,13 @@ class AsyncVectorQuery(AsyncStandardQuery, AsyncVectorQueryBase): self._inner.nearest_to_text({"query": query}), self._table ) - async def to_batches( + async def _execute_batches( self, *, max_batch_length: Optional[int] = None, timeout: Optional[timedelta] = None, ) -> AsyncRecordBatchReader: - reader = await super().to_batches(timeout=timeout) + reader = await super()._execute_batches(timeout=timeout) results = pa.Table.from_batches(await reader.read_all(), reader.schema) if self._reranker: results = self._reranker.rerank_vector(self._query_string, results) @@ -4456,7 +4578,7 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): return fts_query, vec_query, limit, offset - async def to_batches( + async def _execute_batches( self, *, max_batch_length: Optional[int] = None, @@ -4726,6 +4848,7 @@ class BaseQueryBuilder(object): *, max_batch_length: Optional[int] = None, timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", ) -> pa.RecordBatchReader: """ Execute the query and return the results as an Apache Arrow RecordBatchReader. @@ -4742,9 +4865,14 @@ class BaseQueryBuilder(object): The maximum time to wait for the query to complete. If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. + blob_mode: str, default "lazy" + Controls blob conversion as described in + :meth:`AsyncQueryBase.to_batches`. """ async_reader = LOOP.run( - self._inner.to_batches(max_batch_length=max_batch_length, timeout=timeout) + self._inner.to_batches( + max_batch_length=max_batch_length, timeout=timeout, blob_mode=blob_mode + ) ) def iter_sync(): diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index f4c6a9e24..96bf88c24 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -843,8 +843,7 @@ def test_to_pandas_does_not_drop_fragment_selection_on_fallback(): table.search().with_row_address(True).to_pandas(blob_mode="descriptions") -@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes"]) -def test_to_pandas_reports_lance_dataset_failures(tmp_path, blob_mode): +def test_to_pandas_reports_lance_dataset_failures(tmp_path): table = lancedb.connect(tmp_path).create_table( "dataset_failure", schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]), @@ -856,7 +855,59 @@ def test_to_pandas_reports_lance_dataset_failures(tmp_path, blob_mode): table.to_lance = reject_credentials with pytest.raises(PermissionError, match="storage credentials rejected"): - table.search().to_pandas(blob_mode=blob_mode) + table.search().to_pandas(blob_mode="lazy") + + +@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes"]) +@pytest.mark.parametrize("with_row_id", [None, True, False]) +def test_to_batches_blob_modes(tmp_path, monkeypatch, blob_mode, with_row_id): + table = _nested_blob_table(lancedb.connect(tmp_path), "batch_blob_modes") + query = table.search().select({"renamed": "info"}) + if with_row_id is not None: + query.with_row_id(with_row_id) + fetched = [] + fetch_blobs = table.fetch_blobs + + def fetch(column, row_ids): + fetched.append(row_ids) + return fetch_blobs(column, row_ids) + + monkeypatch.setattr(table, "fetch_blobs", fetch) + reader = query.to_batches(batch_size=1, blob_mode=blob_mode) + assert fetched == [] + first = reader.read_next_batch() + assert first.schema == reader.schema + assert len(fetched) == (1 if blob_mode == "bytes" else 0) + batches = [first, *reader] + assert all(batch.schema == reader.schema for batch in batches) + assert all(batch.num_rows == 1 for batch in batches) + assert len(fetched) == (3 if blob_mode == "bytes" else 0) + if blob_mode == "bytes": + assert [row["renamed"] for batch in batches for row in batch.to_pylist()] == [ + {"name": "first", "blob": b"hello"}, + {"name": "second", "blob": None}, + None, + ] + if blob_mode != "lazy": + assert ("_rowid" in reader.schema.names) == (with_row_id is True) + assert "_lance_row_id" not in str(reader.schema) + else: + assert "_rowid" in reader.schema.names + + +@pytest.mark.parametrize("blob_mode", ["descriptions", "bytes"]) +def test_empty_to_batches_blob_schema(blob_mode): + table = _blob_table("empty_blob_batches", [{"id": 1, "image": b"one"}]) + query = table.search().where("id < 0") + reader = query.to_batches(blob_mode=blob_mode) + assert reader.schema == query.to_arrow(blob_mode=blob_mode).schema + assert reader.read_all().num_rows == 0 + + +def test_to_batches_validates_blob_mode(): + table = _blob_table("invalid_batch_blob_mode", [{"id": 1, "image": b"one"}]) + with pytest.raises(ValueError, match="blob_mode must be one of"): + table.search().to_batches(blob_mode="handles") def test_to_pandas_keeps_to_arrow_overrides_without_blob_mode(): @@ -894,6 +945,12 @@ def test_take_queries_honor_blob_mode(): ] description = table.take_row_ids([0]).to_list(blob_mode="descriptions")[0]["image"] assert "_lance_row_id" not in description + batches = table.take_row_ids([0]).to_batches(blob_mode="bytes") + assert batches.read_all().to_pylist() == [{"id": 1, "image": b"one"}] + batches = table.take_row_ids([0]).to_batches(blob_mode="descriptions") + assert "_rowid" not in batches.schema.names + assert "_lance_row_id" not in str(batches.schema) + assert batches.read_all().num_rows == 1 def test_bytes_mode_follows_swapped_blob_aliases(): @@ -939,6 +996,8 @@ def test_bytes_mode_refused_under_namespace_query_pushdown(tmp_path): table.search().to_list(blob_mode="bytes") with pytest.raises(NotImplementedError, match="stable table snapshot"): table.take_row_ids([0]).to_list(blob_mode="bytes") + with pytest.raises(NotImplementedError, match="stable table snapshot"): + table.search().to_batches(blob_mode="bytes") descriptor = table.to_arrow(blob_mode="descriptions")["image"].to_pylist()[0] assert "_lance_row_id" not in descriptor @@ -960,6 +1019,8 @@ async def test_async_bytes_mode_refused_under_namespace_query_pushdown(tmp_path) await table.to_arrow(blob_mode="bytes") with pytest.raises(NotImplementedError, match="stable table snapshot"): await table.query().to_list(blob_mode="bytes") + with pytest.raises(NotImplementedError, match="stable table snapshot"): + await table.query().to_batches(blob_mode="bytes") def test_fetch_blobs_round_trip(): @@ -2076,6 +2137,10 @@ def test_blob_mode_for_every_query_shape(shape, blob_mode, with_row_id): assert all("_lance_row_id" not in image for image in images.values()) df = query.to_pandas(blob_mode=blob_mode) assert dict(zip(df["id"], df["image"])) == images + if shape in ("scan", "vector"): + reader = query.to_batches(batch_size=1, blob_mode=blob_mode) + assert reader.schema == hits.schema + assert reader.read_all().equals(hits) @pytest.mark.asyncio @@ -2102,6 +2167,48 @@ async def test_async_blob_bytes_for_every_query_shape(shape): assert dict(zip(df["id"], df["image"])) == {1: b"alpha", 2: b"beta"} +@pytest.mark.asyncio +@pytest.mark.parametrize("shape", ["scan", "vector", "fts", "hybrid"]) +@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes"]) +async def test_async_to_batches_blob_modes(shape, blob_mode, monkeypatch): + db = await lancedb.connect_async("memory:///") + table = await _hybrid_blob_table_async(db) + query = table.query() + if shape in ("vector", "hybrid"): + query = query.nearest_to([1.0, 0.0]) + if shape in ("fts", "hybrid"): + query = query.nearest_to_text("hello") + if shape == "scan": + query = query.where("id < 3") + query = query.select({"id": "id", "payload": "image"}).limit(2) + fetched = [] + fetch_blobs = table.fetch_blobs + + async def fetch(column, row_ids): + fetched.append(row_ids) + return await fetch_blobs(column, row_ids) + + monkeypatch.setattr(table, "fetch_blobs", fetch) + reader = await query.to_batches(max_batch_length=1, blob_mode=blob_mode) + assert fetched == [] + first = await reader.__anext__() + assert len(fetched) == (1 if blob_mode == "bytes" else 0) + batches = [first, *[batch async for batch in reader]] + assert all(batch.num_rows == 1 for batch in batches) + assert all(batch.schema == reader.schema for batch in batches) + if blob_mode == "bytes": + assert len(fetched) == 2 + assert { + row["id"]: row["payload"] for b in batches for row in b.to_pylist() + } == { + 1: b"alpha", + 2: b"beta", + } + if blob_mode != "lazy": + assert "_rowid" not in reader.schema.names + assert "_lance_row_id" not in str(reader.schema) + + class _KeepAllScoresReranker(Reranker): """Rebuilds the merged table like model rerankers with return_score="all".""" diff --git a/python/python/tests/test_blob_query_without_pylance.py b/python/python/tests/test_blob_query_without_pylance.py new file mode 100644 index 000000000..cd9fd73ee --- /dev/null +++ b/python/python/tests/test_blob_query_without_pylance.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import subprocess +import sys +import textwrap + +import pytest + + +@pytest.mark.parametrize("api", ["sync", "async"]) +@pytest.mark.parametrize("result", ["pandas", "batches", "lazy_error"]) +def test_blob_query_without_pylance(tmp_path, api, result): + script = """\ +import importlib.abc +import sys + +class MissingLance(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "lance" or fullname.startswith("lance."): + raise ModuleNotFoundError(fullname, name="lance") + +sys.meta_path.insert(0, MissingLance()) +""" + textwrap.dedent( + """\ + import asyncio + import lancedb + import pyarrow as pa + + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("video")]) + rows = [{"id": 1, "video": b"one"}, {"id": 2, "video": None}, + {"id": 3, "video": b"three"}] + + def check(data): + assert data == [{"filename": 2, "payload": None}, + {"filename": 3, "payload": b"three"}], data + + async def main(): + if sys.argv[2] == "async": + db = await lancedb.connect_async(sys.argv[1]) + table = await db.create_table("videos", schema=schema) + await table.add(rows) + query = table.query() + else: + db = lancedb.connect(sys.argv[1]) + table = db.create_table("videos", schema=schema) + table.add(rows) + query = table.search() + query = query.select({"filename": "id", "payload": "video"}).offset(1) + if sys.argv[3] == "pandas": + frame = query.to_pandas(blob_mode="bytes", split_blocks=True) + if sys.argv[2] == "async": + frame = await frame + check(frame.to_dict("records")) + frame = table.to_pandas(blob_mode="bytes") + if sys.argv[2] == "async": + frame = await frame + assert frame.to_dict("records") == rows + elif sys.argv[3] == "batches": + if sys.argv[2] == "async": + reader = await query.to_batches( + max_batch_length=1, blob_mode="bytes") + batches = [batch async for batch in reader] + else: + reader = query.to_batches(batch_size=1, blob_mode="bytes") + batches = list(reader) + assert reader.schema.field("payload").type == pa.large_binary() + assert all(batch.num_rows <= 1 for batch in batches) + assert all(batch.schema == reader.schema for batch in batches) + check([row for batch in batches for row in batch.to_pylist()]) + else: + try: + frame = query.to_pandas() + if sys.argv[2] == "async": + await frame + except RuntimeError as error: + message = str(error) + assert "pip install pylance" in message, message + assert ".." not in message, message + assert "Use a plain scan query" not in message, message + assert "but The" not in message, message + else: + raise AssertionError("lazy handles should require pylance") + assert "lance" not in sys.modules + + asyncio.run(main()) + """ + ) + completed = subprocess.run( + [sys.executable, "-c", script, str(tmp_path), api, result], + capture_output=True, + text=True, + ) + assert completed.returncode == 0, completed.stderr diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index 3d7c01fc5..b898a8888 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -675,10 +675,10 @@ def test_blob_v2_vector_projection_and_plain_scan_options(tmp_db, monkeypatch): query = table.search().where("id = 1") - def fail_to_arrow(*args, **kwargs): - raise AssertionError("the native scanner should handle bytes with flatten") + def fail_to_lance(*args, **kwargs): + raise AssertionError("blob v2 bytes should not open a pylance dataset") - monkeypatch.setattr(query, "to_arrow", fail_to_arrow) + monkeypatch.setattr(table, "to_lance", fail_to_lance) assert query.to_pandas(blob_mode="bytes", flatten=True)["blob"].tolist() == [b"one"] flattened_lazy = table.search().where("id = 1").to_pandas(flatten=True)