diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index dc4ed37df..eb0f1a118 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -6,6 +6,7 @@ from __future__ import annotations import io +import operator from collections.abc import Awaitable, Callable, Iterable from typing import TYPE_CHECKING, Optional, Union @@ -391,16 +392,30 @@ def _leaf_struct_column(tbl: pa.Table, path: str) -> pa.StructArray: def _normalize_blob_row_ids( - row_ids: Union[list[int], pa.Table], blob_column: str + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], blob_column: str ) -> list[int]: if isinstance(row_ids, pa.Table): - return read_row_ids_from_hits(row_ids, blob_column) - if isinstance(row_ids, (pa.Array, pa.ChunkedArray)): - raise ValueError( - "pass a query table with _rowid, not a column array " - "(use fetch_blobs('image', hits), not fetch_blobs('image', hits['image']))" - ) - return list(row_ids) + values = read_row_ids_from_hits(row_ids, blob_column) + elif isinstance(row_ids, (pa.Array, pa.ChunkedArray)): + if not pa.types.is_integer(row_ids.type): + raise ValueError( + "row id arrays must have an integer type; pass a query table with " + "_rowid when using a blob descriptor column" + ) + values = row_ids.to_pylist() + else: + values = list(row_ids) + + normalized = [] + for value in values: + try: + row_id = operator.index(value) + except TypeError as exc: + raise ValueError("row ids must be integers without nulls") from exc + if row_id < 0: + raise ValueError("row ids must be non-negative") + normalized.append(row_id) + return normalized def _wrap_blob_files(handles: Iterable[object]) -> list[Optional[BlobFile]]: diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index ba4e3a0df..a2cf4c4cf 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -1,6 +1,6 @@ from datetime import date, datetime, timedelta from decimal import Decimal -from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal +from typing import Dict, List, Optional, Sequence, Tuple, Any, TypedDict, Union, Literal from uuid import UUID import pyarrow as pa @@ -529,7 +529,7 @@ class Table: async def fetch_blob_ranges( self, column: str, - requests: List[Tuple[int, int, int]], + requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: ... async def fetch_blob_files( self, column: str, row_ids: list[int] diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 2a6512fc5..c23d8cf58 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1140,7 +1140,9 @@ class RemoteTable(Table): return LOOP.run(self._table.blob_columns()) def fetch_blobs( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> pa.LargeBinaryArray: return LOOP.run(self._table.fetch_blobs(column, row_ids)) @@ -1150,7 +1152,9 @@ class RemoteTable(Table): ) def fetch_blob_files( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 297b4237f..895a73dcb 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1924,15 +1924,18 @@ class Table(ABC): @abstractmethod def fetch_blobs( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> pa.LargeBinaryArray: """Materialize full blob bytes for ``column`` at the given rows. The result has the same length and order as ``row_ids``. Null blobs produce null slots; valid empty blobs produce ``b""``. - ``_rowid`` values stay valid after compaction when the table has stable - row ids. + ``row_ids`` may be a list, an integer Arrow array or chunked array, + or a query table with row ids. ``_rowid`` values stay valid after + compaction when the table has stable row ids. Convenience for small payloads. For large values use :meth:`fetch_blob_files`. @@ -1942,12 +1945,13 @@ class Table(ABC): def fetch_blob_ranges( self, column: str, - requests: Sequence[Tuple[int, int, int]], + requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: """Materialize row-specific byte ranges from a blob v2 column. - Each request is a ``(row_id, offset, length)`` tuple. Requests may be - repeated or reordered, including multiple ranges for the same blob. + Each request is a three-item ``(row_id, offset, length)`` sequence. + Requests may be repeated or reordered, including multiple ranges for + the same blob. The result has the same length and order as ``requests``; null blobs produce null slots and empty ranges on non-null blobs produce ``b""``. @@ -1960,15 +1964,17 @@ class Table(ABC): @abstractmethod def fetch_blob_files( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> "list[Optional[BlobFile]]": """Open lazy, seekable :class:`~lancedb._blob.BlobFile` handles. Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is - a ``list[int]`` or a query ``pyarrow.Table`` carrying row identity via - ``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null - rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or - newer. + a ``list[int]``, an integer Arrow array or chunked array, or a query + ``pyarrow.Table`` carrying row identity via ``_rowid`` or a + ``_lance_row_id`` field on the blob descriptor. Null blob values return + ``None``. Remote tables require LanceDB Cloud server 0.5.0 or newer. ``_rowid`` values stay valid after compaction when the table has stable row ids. @@ -2854,19 +2860,23 @@ class LanceTable(Table): return LOOP.run(self._table.blob_columns()) def fetch_blobs( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> pa.LargeBinaryArray: return LOOP.run(self._table.fetch_blobs(column, row_ids)) def fetch_blob_ranges( self, column: str, - requests: Sequence[Tuple[int, int, int]], + requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: return LOOP.run(self._table.fetch_blob_ranges(column, list(requests))) def fetch_blob_files( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) @@ -6860,7 +6870,9 @@ class AsyncTable: return await self._inner.blob_columns() async def fetch_blobs( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> pa.LargeBinaryArray: return await self._inner.fetch_blobs( column, _normalize_blob_row_ids(row_ids, column) @@ -6869,12 +6881,14 @@ class AsyncTable: async def fetch_blob_ranges( self, column: str, - requests: Sequence[Tuple[int, int, int]], + requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: return await self._inner.fetch_blob_ranges(column, list(requests)) async def fetch_blob_files( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> "list[Optional[BlobFile]]": handles = await self._inner.fetch_blob_files( column, _normalize_blob_row_ids(row_ids, column) diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 0d3a89c21..6954b73e4 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -68,7 +68,7 @@ def _row_ids_by_id(table): def _assert_missing_blob_row_ids(exc_info): message = str(exc_info.value) - assert "row ids" in message + assert "row ids but some do not exist in the table" in message assert "rowaddr" not in message assert "fragment" not in message @@ -713,6 +713,38 @@ def test_fetch_blobs_accepts_query_result(): assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"} +def test_fetch_blob_apis_accept_integer_arrow_row_ids(): + table = _blob_table( + "arrow_row_ids", + [{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}], + ) + by_id = _row_ids_by_id(table) + ids = [by_id[2], by_id[1]] + + for row_ids in ( + pa.array(ids, type=pa.uint64()), + pa.chunked_array([ids[:1], ids[1:]], type=pa.int64()), + ): + assert table.fetch_blobs("image", row_ids).to_pylist() == [b"beta", b"alpha"] + files = table.fetch_blob_files("image", row_ids) + assert [file.read() for file in files] == [b"beta", b"alpha"] + + +def test_fetch_blob_apis_validate_arrow_and_negative_row_ids(): + table = _blob_table("invalid_arrow_row_ids", [{"id": 1, "image": b"alpha"}]) + hits = table.search().with_row_id(True).to_arrow() + + for fetch in (table.fetch_blobs, table.fetch_blob_files): + with pytest.raises(ValueError, match="non-negative"): + fetch("image", [-1]) + with pytest.raises(ValueError, match="non-negative"): + fetch("image", pa.array([-1], type=pa.int64())) + with pytest.raises(ValueError, match="without nulls"): + fetch("image", pa.array([None], type=pa.int64())) + with pytest.raises(ValueError, match="integer type"): + fetch("image", hits["image"]) + + def test_fetch_blobs_after_compact_with_stable_row_ids(tmp_path): db = lancedb.connect( tmp_path, storage_options={"new_table_enable_stable_row_ids": "true"} @@ -844,6 +876,15 @@ def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls(): assert ranges.to_pylist() == [b"cde", None, b"ab", b"cde", b""] +def test_fetch_blob_ranges_accepts_three_item_sequences(): + table = _blob_table("range_sequences", [{"id": 1, "image": b"abcdef"}]) + row_id = _row_ids_by_id(table)[1] + + ranges = table.fetch_blob_ranges("image", [[row_id, 1, 2], (row_id, 4, 2)]) + + assert ranges.to_pylist() == [b"bc", b"ef"] + + def test_fetch_blob_ranges_validates_requests(): table = _blob_table("range_validation", [{"id": 1, "image": b"abc"}]) row_id = _row_ids_by_id(table)[1] @@ -858,6 +899,30 @@ def test_fetch_blob_ranges_validates_requests(): table.fetch_blob_ranges("image", [(2**64 - 1, 0, 1)]) _assert_missing_blob_row_ids(exc_info) + for index, name in enumerate(("row id", "offset", "length")): + request = [row_id, 0, 1] + request[index] = -1 + with pytest.raises( + ValueError, match=f"blob range {name} must be a non-negative" + ): + table.fetch_blob_ranges("image", [request]) + + with pytest.raises(ValueError, match="three values"): + table.fetch_blob_ranges("image", [[row_id, 0]]) + + +def test_fetch_blob_apis_reject_row_past_fragment_end(): + table = _blob_table( + "past_fragment_end", + [ + {"id": 1, "image": b"a"}, + {"id": 2, "image": b"b"}, + {"id": 3, "image": b"c"}, + ], + ) + _assert_fetch_apis_reject_missing_row_ids(table, [3]) + _assert_fetch_apis_reject_missing_row_ids(table, [3, 0]) + def test_fetch_blob_apis_reject_missing_fragment_row_addr(): table = _blob_table("missing_frag", [{"id": 1, "image": b"x"}]) @@ -875,6 +940,19 @@ def test_fetch_blob_apis_reject_deleted_row_ids(): _assert_fetch_apis_reject_missing_row_ids(table, [by_id[2], by_id[1]]) +def test_fetch_blob_apis_reject_deleted_stable_row_ids(tmp_path): + db = lancedb.connect( + tmp_path, storage_options={"new_table_enable_stable_row_ids": "true"} + ) + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("stable_deleted", schema=schema) + table.add([{"id": 1, "image": b"one"}, {"id": 2, "image": b"two"}]) + by_id = _row_ids_by_id(table) + table.delete("id = 2") + + _assert_fetch_apis_reject_missing_row_ids(table, [by_id[2], by_id[1]]) + + def test_fetch_blob_ranges_empty_requests_returns_empty_array(): table = _blob_table("range_empty", [{"id": 1, "image": b"x"}]) assert table.fetch_blob_ranges("image", []).to_pylist() == [] diff --git a/python/src/table.rs b/python/src/table.rs index a3dad3d86..70533a572 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1421,14 +1421,30 @@ impl Table { pub fn fetch_blob_ranges( self_: PyRef<'_, Self>, column: String, - requests: Vec<(u64, u64, u64)>, + requests: Vec>, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); + let requests = requests + .into_iter() + .map(|values| { + let [row_id, offset, length] = <[i128; 3]>::try_from(values).map_err(|_| { + PyValueError::new_err("each blob range request must have three values") + })?; + let to_u64 = |value, name| { + u64::try_from(value).map_err(|_| { + PyValueError::new_err(format!( + "blob range {name} must be a non-negative 64-bit integer" + )) + }) + }; + Ok(BlobRangeRequest::new( + to_u64(row_id, "row id")?, + to_u64(offset, "offset")?, + to_u64(length, "length")?, + )) + }) + .collect::>>()?; future_into_py(self_.py(), async move { - let requests = requests - .into_iter() - .map(|(row_id, offset, length)| BlobRangeRequest::new(row_id, offset, length)) - .collect::>(); let blobs: LargeBinaryArray = inner .fetch_blob_ranges(column, requests) .await diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index cbc0724c6..382211285 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -368,17 +368,21 @@ pub(crate) fn ensure_blob_v2_column( } } +fn missing_blob_row_ids_error(column: &str, requested: usize) -> Error { + Error::InvalidInput { + message: format!( + "blob read for column '{column}' requested {requested} row ids but some \ + do not exist in the table; pass row ids collected from this table" + ), + } +} + fn ensure_all_row_ids_resolved(column: &str, requested: usize, resolved: usize) -> Result<()> { if requested == resolved { return Ok(()); } if resolved < requested { - Err(Error::InvalidInput { - message: format!( - "blob read for column '{column}' requested {requested} row ids but only {resolved} \ - exist in the table; pass row ids collected from this table" - ), - }) + Err(missing_blob_row_ids_error(column, requested)) } else { Err(Error::Runtime { message: format!( @@ -388,25 +392,17 @@ fn ensure_all_row_ids_resolved(column: &str, requested: usize, resolved: usize) } } -/// Lance take reports a missing physical row address as NotSupported or InvalidInput. +/// Lance reports missing physical and stable row ids through several error paths +/// without a shared typed error. fn map_blob_take_error(column: &str, requested: usize, err: lance::Error) -> Error { - let missing_row_addr = match &err { - lance::Error::NotSupported { source, .. } => { - source.to_string().contains("must not target deleted rows") - } - lance::Error::InvalidInput { source, .. } => source - .to_string() - .contains("belongs to non-existent fragment"), - _ => false, - }; + let message = err.to_string(); + let missing_row_id = message.contains("must not target deleted rows") + || message.contains("belongs to non-existent fragment") + || (message.contains("Invalid read params ") && message.contains("addressable rows")) + || message.contains("Could not resolve all requested row IDs"); - if missing_row_addr { - Error::InvalidInput { - message: format!( - "blob read for column '{column}' requested {requested} row ids but some \ - do not exist in the table; pass row ids collected from this table" - ), - } + if missing_row_id { + missing_blob_row_ids_error(column, requested) } else { err.into() } diff --git a/rust/lancedb/tests/blob_integration.rs b/rust/lancedb/tests/blob_integration.rs index dd4051657..4a14d2732 100644 --- a/rust/lancedb/tests/blob_integration.rs +++ b/rust/lancedb/tests/blob_integration.rs @@ -501,7 +501,10 @@ async fn collect_id_rowid(table: &Table) -> Result> { fn assert_missing_blob_row_ids(err: &Error) { assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); let message = err.to_string(); - assert!(message.contains("row ids"), "{message}"); + assert!( + message.contains("row ids but some do not exist in the table"), + "{message}" + ); assert!(!message.contains("rowaddr"), "{message}"); assert!(!message.contains("fragment"), "{message}"); } @@ -790,6 +793,18 @@ async fn fetch_blobs_out_of_range_id_errors_without_panic() -> Result<()> { Ok(()) } +#[tokio::test] +async fn fetch_blob_apis_reject_row_past_fragment_end() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let table = + create_inline_blob_table(&db, "t", &[1, 2, 3], &[Some(b"a"), Some(b"b"), Some(b"c")]) + .await?; + + assert_fetch_apis_reject_missing_row_ids(&table, &[3]).await?; + assert_fetch_apis_reject_missing_row_ids(&table, &[3, 0]).await +} + #[tokio::test] async fn fetch_blob_files_rejects_missing_fragment_row_addr() -> Result<()> { let tmp = tempdir().unwrap(); @@ -831,6 +846,23 @@ async fn fetch_blob_apis_reject_deleted_row_ids() -> Result<()> { assert_fetch_apis_reject_missing_row_ids(&table, &[deleted_row_addr, live_row_addr]).await } +#[tokio::test] +async fn fetch_blob_apis_reject_deleted_stable_row_ids() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()) + .storage_option(OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS, "true") + .execute() + .await?; + let table = create_inline_blob_table(&db, "t", &[1, 2], &[Some(b"one"), Some(b"two")]).await?; + let pairs = collect_id_rowid(&table).await?; + let deleted_row_id = pairs.iter().find(|(id, _)| *id == 2).unwrap().1; + let live_row_id = pairs.iter().find(|(id, _)| *id == 1).unwrap().1; + + table.delete("id = 2").await?; + + assert_fetch_apis_reject_missing_row_ids(&table, &[deleted_row_id, live_row_id]).await +} + #[tokio::test] async fn fetch_blobs_rejects_non_blob_column() -> Result<()> { let tmp = tempdir().unwrap();