diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index c922aa69b..229eeaafc 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -297,6 +297,11 @@ class Table: async def fetch_blobs( self, column: str, row_ids: list[int] ) -> pa.LargeBinaryArray: ... + async def fetch_blob_ranges( + self, + column: str, + requests: List[Tuple[int, int, int]], + ) -> pa.LargeBinaryArray: ... async def fetch_blob_files( self, column: str, row_ids: list[int] ) -> list[Optional[BlobFile]]: ... diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 8c702aace..d01224909 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1040,6 +1040,11 @@ class RemoteTable(Table): def fetch_blobs(self, column: str, row_ids) -> pa.LargeBinaryArray: raise NotImplementedError("fetch_blobs() is not supported on LanceDB Cloud") + def fetch_blob_ranges(self, column: str, requests) -> pa.LargeBinaryArray: + raise NotImplementedError( + "fetch_blob_ranges() is not supported on LanceDB Cloud" + ) + def fetch_blob_files(self, column: str, row_ids): raise NotImplementedError( "fetch_blob_files() is not supported on LanceDB Cloud" diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 0e181b457..b2c8682c1 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -20,6 +20,7 @@ from typing import ( List, Literal, Optional, + Sequence, Tuple, Union, overload, @@ -1537,6 +1538,23 @@ class Table(ABC): :meth:`fetch_blob_files`. """ + @abstractmethod + def fetch_blob_ranges( + self, + column: str, + requests: Sequence[Tuple[int, int, 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. + The result has the same length and order as ``requests``; null blobs + produce null slots and empty ranges on non-null blobs produce ``b""``. + + Row IDs can be obtained from a query with ``with_row_id(True)``. This + API is currently supported only by local tables. + """ + @abstractmethod def fetch_blob_files( self, column: str, row_ids: Union[list[int], pa.Table] @@ -2260,6 +2278,13 @@ class LanceTable(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]], + ) -> 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] ) -> "list[Optional[BlobFile]]": @@ -5823,6 +5848,13 @@ class AsyncTable: column, _normalize_blob_row_ids(row_ids, column) ) + async def fetch_blob_ranges( + self, + column: str, + requests: Sequence[Tuple[int, int, 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] ) -> "list[Optional[BlobFile]]": diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 05af14298..59c80a57c 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -198,6 +198,58 @@ def test_fetch_blobs_null_alignment(): assert blobs[2].as_py() == b"present" +def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls(): + table = _blob_table( + "range_alignment", + [{"id": 1, "image": b"abcdefghij"}, {"id": 2, "image": None}], + ) + by_id = _row_ids_by_id(table) + requests = [ + (by_id[1], 2, 3), + (by_id[2], 0, 0), + (by_id[1], 0, 2), + (by_id[1], 2, 3), + (by_id[1], 10, 0), + ] + + ranges = table.fetch_blob_ranges("image", requests) + + assert ranges.to_pylist() == [b"cde", None, b"ab", b"cde", b""] + + +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] + + with pytest.raises(ValueError, match="exceeds blob size"): + table.fetch_blob_ranges("image", [(row_id, 2, 2)]) + + with pytest.raises(ValueError, match="offset \\+ length overflowed"): + table.fetch_blob_ranges("image", [(row_id, 2**64 - 1, 1)]) + + with pytest.raises(ValueError, match="row ids"): + table.fetch_blob_ranges("image", [(2**64 - 1, 0, 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() == [] + + +@pytest.mark.asyncio +async def test_async_fetch_blob_ranges(): + db = await lancedb.connect_async("memory:///") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = await db.create_table("range_async", schema=schema) + await table.add([{"id": 1, "image": b"abcdefghij"}]) + hits = await table.query().with_row_id().to_arrow() + row_id = hits["_rowid"][0].as_py() + + ranges = await table.fetch_blob_ranges("image", [(row_id, 1, 3), (row_id, 6, 2)]) + + assert ranges.to_pylist() == [b"bcd", b"gh"] + + def test_fetch_blobs_nested_path(): db = lancedb.connect("memory:///") info = pa.StructArray.from_arrays( diff --git a/python/src/table.rs b/python/src/table.rs index 5d6910685..277a99b70 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -17,7 +17,7 @@ use arrow::{ ffi_stream::ArrowArrayStreamReader, pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; -use lancedb::blob::BlobFile; +use lancedb::blob::{BlobFile, BlobRangeRequest}; use lancedb::index::scalar::FtsIndexBuilder; use lancedb::table::{ AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken, @@ -1101,6 +1101,27 @@ impl Table { }) } + /// Read row-specific blob-local byte ranges in one planned operation. + #[pyo3(signature = (column, requests))] + pub fn fetch_blob_ranges( + self_: PyRef<'_, Self>, + column: String, + requests: Vec<(u64, u64, u64)>, + ) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + 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 + .infer_error()?; + Python::attach(|py| blobs.to_data().to_pyarrow(py).map(|obj| obj.unbind())) + }) + } + /// Open lazy blob handles for `row_ids` from blob v2 column `column`. #[pyo3(signature = (column, row_ids))] pub fn fetch_blob_files( diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index c8884f9b3..deaa22e84 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -14,7 +14,8 @@ use std::sync::Arc; use arrow_array::builder::LargeBinaryBuilder; use arrow_array::{Array, LargeBinaryArray, RecordBatch, StructArray, UInt8Array, UInt64Array}; use arrow_schema::{DataType, Field, Schema}; -use lance::dataset::{Dataset, WriteParams}; +use futures::TryStreamExt; +use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams}; use lance_arrow::FieldExt; use lance_core::datatypes::parse_field_path; use lance_encoding::version::LanceFileVersion; @@ -23,6 +24,40 @@ use crate::error::{Error, Result}; pub use lance::dataset::BlobFile; +/// One row-specific blob range read request. +/// +/// `row_id` is obtained from a query with row ids enabled. +/// `offset` and `length` are relative to the beginning of the logical blob. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BlobRangeRequest { + /// Row id of the blob value to read. + pub row_id: u64, + /// Byte offset from the beginning of the blob value. + pub offset: u64, + /// Number of bytes to read. + pub length: u64, +} + +impl BlobRangeRequest { + /// Create a row-specific blob range request. + pub const fn new(row_id: u64, offset: u64, length: u64) -> Self { + Self { + row_id, + offset, + length, + } + } +} + +fn map_blob_read_error(source: lance::Error) -> Error { + match source { + source @ lance::Error::InvalidInput { .. } => Error::InvalidInput { + message: source.to_string(), + }, + source => source.into(), + } +} + /// Creates an Arrow field for a Lance blob v2 column. /// /// `Struct` with the `lance.blob.v2` marker. Same layout Lance @@ -232,6 +267,84 @@ fn non_null_row_ids(row_ids: &[u64], null_mask: &[bool]) -> Vec { .collect() } +/// Materialize blob-local ranges (same length and order as `requests`, nulls preserved). +pub(crate) async fn take_blob_ranges_aligned( + dataset: &Arc, + column: &str, + requests: &[BlobRangeRequest], +) -> Result { + ensure_blob_v2_column(dataset.schema(), column)?; + if requests.is_empty() { + return Ok(LargeBinaryBuilder::new().finish()); + } + + let lance_requests = requests + .iter() + .map(|request| LanceBlobRangeRequest::new(request.row_id, request.offset, request.length)) + .collect::>(); + let reader = dataset + .read_blob_ranges(column) + .map_err(map_blob_read_error)? + .with_row_ids(lance_requests) + .preserve_order(false); + + let mut stream = reader + .try_into_stream() + .await + .map_err(map_blob_read_error)?; + let mut payloads = vec![None; requests.len()]; + while let Some(result) = stream.try_next().await.map_err(map_blob_read_error)? { + let slot = payloads + .get_mut(result.request_index) + .ok_or_else(|| Error::Runtime { + message: format!( + "blob range read for column '{column}' returned out-of-range request index {} for {} requests", + result.request_index, + requests.len() + ), + })?; + if slot.replace(result.data).is_some() { + return Err(Error::Runtime { + message: format!( + "blob range read for column '{column}' returned request index {} more than once", + result.request_index + ), + }); + } + } + + let missing_request_indices = payloads + .iter() + .enumerate() + .filter_map(|(request_index, payload)| payload.is_none().then_some(request_index)) + .collect::>(); + if !missing_request_indices.is_empty() { + let missing_row_ids = missing_request_indices + .iter() + .map(|request_index| requests[*request_index].row_id) + .collect::>(); + let null_mask = blob_null_mask(dataset, column, &missing_row_ids).await?; + for (request_index, is_null) in missing_request_indices.into_iter().zip(null_mask) { + if !is_null { + return Err(Error::Runtime { + message: format!( + "blob range read for column '{column}' did not return request index {request_index}" + ), + }); + } + } + } + + let mut builder = LargeBinaryBuilder::new(); + for payload in payloads { + match payload { + Some(data) => builder.append_value(data.as_ref()), + None => builder.append_null(), + } + } + Ok(builder.finish()) +} + /// Materialize blob bytes for `row_ids` (same length and order, nulls preserved). pub(crate) async fn take_blobs_aligned( dataset: &Arc, diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 0b2962682..057771bdc 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -196,7 +196,7 @@ use std::{fmt::Display, str::FromStr}; use serde::{Deserialize, Serialize}; -pub use blob::{blob, is_blob}; +pub use blob::{BlobRangeRequest, blob, is_blob}; pub use connection::{ConnectNamespaceBuilder, Connection}; pub use error::{Error, Result}; use lance_index::vector::ApproxMode as LanceApproxMode; diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 1f7b2f75b..fcf1997c0 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -46,6 +46,7 @@ use std::sync::Arc; use crate::connection::NamespaceClientPushdownOperation; use crate::DistanceType; +use crate::blob::BlobRangeRequest; use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions}; use crate::database::Database; use crate::database::read_freshness::TableFreshness; @@ -646,6 +647,16 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "fetch_blobs is not supported on this table type".into(), }) } + /// Materialize blob-local ranges. See [`Table::fetch_blob_ranges`]. + async fn fetch_blob_ranges( + &self, + _column: &str, + _requests: &[BlobRangeRequest], + ) -> Result { + Err(Error::NotSupported { + message: "fetch_blob_ranges is not supported on this table type".into(), + }) + } /// Open lazy blob handles for the given row ids. See [`Table::fetch_blob_files`]. async fn fetch_blob_files( &self, @@ -1055,6 +1066,47 @@ impl Table { self.inner.fetch_blobs(column.as_ref(), row_ids).await } + /// Materialize row-specific ranges from a blob v2 column. + /// + /// Each request contains a row id and a blob-local offset and length. + /// Requests may be duplicated or reordered, including multiple + /// ranges for the same blob. The output has the same length and order as + /// the requests. Null blobs produce null output slots; empty ranges on + /// non-null blobs produce empty byte strings. + /// + /// ``` + /// use lancedb::blob::BlobRangeRequest; + /// + /// # use lancedb::Table; + /// # async fn read_ranges(table: &Table, row_id: u64) -> Result<(), Box> { + /// let ranges = table + /// .fetch_blob_ranges( + /// "image", + /// [ + /// BlobRangeRequest::new(row_id, 0, 1024), + /// BlobRangeRequest::new(row_id, 4096, 1024), + /// ], + /// ) + /// .await?; + /// # let _ = ranges; + /// # Ok(()) + /// # } + /// ``` + /// + /// Returns [`Error::InvalidInput`] when a range is invalid, a requested + /// row id does not exist, or the column is not a blob v2 column. Returns + /// [`Error::NotSupported`] on table types without blob support. + pub async fn fetch_blob_ranges( + &self, + column: impl AsRef, + requests: impl IntoIterator, + ) -> Result { + let requests = requests.into_iter().collect::>(); + self.inner + .fetch_blob_ranges(column.as_ref(), &requests) + .await + } + /// Open lazy [`BlobFile`] handles for the given row ids. /// /// Same length and order as `row_ids`. Null rows are `None`. Bytes are not @@ -3070,6 +3122,15 @@ impl BaseTable for NativeTable { crate::blob::take_blobs_aligned(&dataset, column, row_ids).await } + async fn fetch_blob_ranges( + &self, + column: &str, + requests: &[BlobRangeRequest], + ) -> Result { + let dataset = self.dataset.get().await?; + crate::blob::take_blob_ranges_aligned(&dataset, column, requests).await + } + async fn fetch_blob_files( &self, column: &str, diff --git a/rust/lancedb/tests/blob_integration.rs b/rust/lancedb/tests/blob_integration.rs index ad0c8d9ff..d44eff105 100644 --- a/rust/lancedb/tests/blob_integration.rs +++ b/rust/lancedb/tests/blob_integration.rs @@ -12,7 +12,7 @@ use futures::TryStreamExt; use lance_encoding::version::LanceFileVersion; use lancedb::{ Connection, Error, Result, Table, - blob::blob, + blob::{BlobRangeRequest, blob}, connect, connect_namespace, database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS, query::{ExecutableQuery, QueryBase}, @@ -595,6 +595,75 @@ async fn fetch_blobs_aligns_with_reordered_and_duplicate_ids() -> Result<()> { Ok(()) } +#[tokio::test] +async fn fetch_blob_ranges_aligns_repeated_ranges_and_nulls() -> 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], &[Some(b"abcdefghij".as_slice()), None]) + .await?; + + let pairs = collect_id_rowid(&table).await?; + let by_id = |want: i64| pairs.iter().find(|(id, _)| *id == want).unwrap().1; + let requests = [ + BlobRangeRequest::new(by_id(1), 2, 3), + BlobRangeRequest::new(by_id(2), 0, 0), + BlobRangeRequest::new(by_id(1), 0, 2), + BlobRangeRequest::new(by_id(1), 2, 3), + BlobRangeRequest::new(by_id(1), 10, 0), + ]; + let bytes = table.fetch_blob_ranges("image", requests).await?; + + assert_eq!(bytes.len(), requests.len()); + assert_eq!(bytes.value(0), b"cde"); + assert!(bytes.is_null(1)); + assert_eq!(bytes.value(2), b"ab"); + assert_eq!(bytes.value(3), b"cde"); + assert_eq!(bytes.value(4), b""); + Ok(()) +} + +#[tokio::test] +async fn fetch_blob_ranges_validates_requests() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let table = create_inline_blob_table(&db, "t", &[1], &[Some(b"abc".as_slice())]).await?; + let row_id = collect_row_ids(&table).await?[0]; + + let err = table + .fetch_blob_ranges("image", [BlobRangeRequest::new(row_id, 2, 2)]) + .await + .unwrap_err(); + assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("exceeds blob size")); + + let err = table + .fetch_blob_ranges("image", [BlobRangeRequest::new(row_id, u64::MAX, 1)]) + .await + .unwrap_err(); + assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("offset + length overflowed")); + + let err = table + .fetch_blob_ranges("image", [BlobRangeRequest::new(u64::MAX, 0, 1)]) + .await + .unwrap_err(); + assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("row ids")); + Ok(()) +} + +#[tokio::test] +async fn fetch_blob_ranges_empty_requests_returns_empty_array() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let table = create_inline_blob_table(&db, "t", &[1], &[Some(b"x".as_slice())]).await?; + + let bytes = table.fetch_blob_ranges("image", std::iter::empty()).await?; + assert!(bytes.is_empty()); + Ok(()) +} + #[tokio::test] async fn fetch_blobs_empty_ids_returns_empty() -> Result<()> { let tmp = tempdir().unwrap(); @@ -843,6 +912,19 @@ async fn fetch_blobs_with_precompaction_row_ids_survives_compaction() -> Result< _ => unreachable!(), } } + + let ranges = ids_before + .iter() + .map(|row_id| BlobRangeRequest::new(*row_id, 5, 3)); + let ranges_after = table.fetch_blob_ranges("image", ranges).await?; + assert_eq!(ranges_after.len(), 2); + for (i, (id, _)) in pairs_before.iter().enumerate() { + match id { + 1 => assert_eq!(ranges_after.value(i), b"one"), + 2 => assert_eq!(ranges_after.value(i), b"two"), + _ => unreachable!(), + } + } Ok(()) } @@ -927,6 +1009,27 @@ async fn fetch_blobs_aligns_across_fragments_with_nulls_and_dups() -> Result<()> Ok(()) } +#[tokio::test] +async fn fetch_blob_ranges_aligns_across_fragments_with_nulls_and_dups() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let table = multi_fragment_dedicated_blob_table(&db).await?; + let row_ids = row_ids_for_logical(&table, &SCRAMBLED_LOGICAL_IDS).await?; + let requests = row_ids + .iter() + .map(|row_id| BlobRangeRequest::new(*row_id, 123, 8)); + + let bytes = table.fetch_blob_ranges("image", requests).await?; + assert_eq!(bytes.len(), SCRAMBLED_LOGICAL_IDS.len()); + for (slot, logical_id) in SCRAMBLED_LOGICAL_IDS.iter().enumerate() { + match logical_id { + 3 | 5 => assert!(bytes.is_null(slot)), + id => assert_eq!(bytes.value(slot), [*id as u8; 8]), + } + } + Ok(()) +} + #[tokio::test] async fn fetch_blob_files_aligns_across_fragments_with_nulls_and_dups() -> Result<()> { let tmp = tempdir().unwrap();