diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 9c000a831..448947a29 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 80c250bf5..682b8533c 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1042,6 +1042,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 f61a1f1cf..6c3854734 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, @@ -1538,10 +1539,30 @@ class Table(ABC): ) -> 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""``. + Convenience for small payloads. For large values use :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] @@ -2265,6 +2286,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]]": @@ -5829,6 +5857,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..04a99ea16 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -184,18 +184,75 @@ def test_fetch_blobs_accepts_query_result(): assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"} -def test_fetch_blobs_null_alignment(): +def test_fetch_blobs_preserves_null_and_empty_values(): table = _blob_table( "nulls", - [{"id": 1, "image": b"present"}, {"id": 2, "image": None}], + [ + {"id": 1, "image": b"present"}, + {"id": 2, "image": None}, + {"id": 3, "image": b""}, + ], ) by_id = _row_ids_by_id(table) - request = [by_id[1], by_id[2], by_id[1]] + request = [by_id[1], by_id[2], by_id[3], by_id[1]] blobs = table.fetch_blobs("image", request) assert len(blobs) == len(request) assert blobs[0].as_py() == b"present" assert blobs[1].as_py() is None - assert blobs[2].as_py() == b"present" + assert blobs[2].as_py() == b"" + assert blobs[3].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(RuntimeError, match="exceeds blob size"): + table.fetch_blob_ranges("image", [(row_id, 2, 2)]) + + with pytest.raises(RuntimeError, 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(): diff --git a/python/src/table.rs b/python/src/table.rs index b62aa7377..bb44fc023 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 de666aa10..a21fe3c1b 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -11,18 +11,42 @@ use std::sync::Arc; +use arrow_array::LargeBinaryArray; 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 lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams}; use lance_arrow::FieldExt; -use lance_core::datatypes::parse_field_path; use lance_encoding::version::LanceFileVersion; 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, + } + } +} + /// Creates an Arrow field for a Lance blob v2 column. /// /// `Struct` with the `lance.blob.v2` marker. Same layout Lance @@ -145,91 +169,57 @@ pub(crate) fn ensure_blob_v2_column( } } -/// Returns the leaf descriptor `StructArray` for `column` in a descriptor batch. -fn leaf_descriptor_struct<'a>(batch: &'a RecordBatch, column: &str) -> Result<&'a StructArray> { - let path = parse_field_path(column).map_err(|e| Error::InvalidInput { - message: format!("invalid blob column path '{column}': {e}"), - })?; - let not_struct = || Error::Runtime { - message: format!("blob column '{column}' did not read back as a descriptor struct"), - }; - let mut current = batch - .column_by_name(&path[0]) - .and_then(|c| c.as_any().downcast_ref::()) - .ok_or_else(not_struct)?; - for segment in &path[1..] { - current = current - .column_by_name(segment) - .and_then(|c| c.as_any().downcast_ref::()) - .ok_or_else(not_struct)?; +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" + ), + }) + } else { + Err(Error::Runtime { + message: format!( + "blob read for column '{column}' returned {resolved} results for {requested} row ids" + ), + }) } - Ok(current) } -/// Null rows in `row_ids`, from a descriptor take. -/// -/// Lance `read_blobs` / `take_blobs` skip null rows (`kind == 0 && position == 0 && size == 0`). -/// TODO(lance): aligned read API would drop this pass. -async fn blob_null_mask( +/// Materialize blob-local ranges (same length and order as `requests`, nulls preserved). +pub(crate) async fn take_blob_ranges_aligned( dataset: &Arc, column: &str, - row_ids: &[u64], -) -> Result> { - let projection = dataset.schema().project(&[column])?; - let descriptors = dataset.take_builder(row_ids, projection)?.execute().await?; - if descriptors.num_rows() != row_ids.len() { - return Err(Error::InvalidInput { - message: format!( - "blob take for column '{column}' requested {} row ids but only {} exist in the \ - table; pass row ids collected from this table", - row_ids.len(), - descriptors.num_rows() - ), - }); + requests: &[BlobRangeRequest], +) -> Result { + ensure_blob_v2_column(dataset.schema(), column)?; + if requests.is_empty() { + return Ok(LargeBinaryBuilder::new().finish()); } - let descriptor_struct = leaf_descriptor_struct(&descriptors, column)?; - let child = |name: &str| { - descriptor_struct - .column_by_name(name) - .ok_or_else(|| Error::Runtime { - message: format!("blob descriptor for '{column}' is missing the '{name}' field"), - }) - }; - let kinds = child("kind")? - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::Runtime { - message: format!("blob descriptor 'kind' for '{column}' is not a UInt8 array"), - })?; - let positions = child("position")? - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::Runtime { - message: format!("blob descriptor 'position' for '{column}' is not a UInt64 array"), - })?; - let sizes = child("size")? - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::Runtime { - message: format!("blob descriptor 'size' for '{column}' is not a UInt64 array"), - })?; - // Match Lance `collect_blob_entries_v2` skip condition (`BlobKind::Inline` == 0). - Ok((0..descriptor_struct.len()) - .map(|i| { - descriptor_struct.is_null(i) - || kinds.is_null(i) - || (kinds.value(i) == 0 && positions.value(i) == 0 && sizes.value(i) == 0) - }) - .collect()) -} - -fn non_null_row_ids(row_ids: &[u64], null_mask: &[bool]) -> Vec { - row_ids + let lance_requests = requests .iter() - .zip(null_mask) - .filter_map(|(row_id, is_null)| (!is_null).then_some(*row_id)) - .collect() + .map(|request| LanceBlobRangeRequest::new(request.row_id, request.offset, request.length)) + .collect::>(); + let payloads = dataset + .read_blob_ranges(column)? + .with_row_ids(lance_requests) + .preserve_order(true) + .execute() + .await?; + ensure_all_row_ids_resolved(column, requests.len(), payloads.len())?; + + let mut builder = LargeBinaryBuilder::new(); + for payload in payloads { + match payload.data { + Some(data) => builder.append_value(data), + None => builder.append_null(), + } + } + Ok(builder.finish()) } /// Materialize blob bytes for `row_ids` (same length and order, nulls preserved). @@ -243,42 +233,19 @@ pub(crate) async fn take_blobs_aligned( return Ok(LargeBinaryBuilder::new().finish()); } - let null_mask = blob_null_mask(dataset, column, row_ids).await?; - let non_null_row_ids = non_null_row_ids(row_ids, &null_mask); - let non_null_count = non_null_row_ids.len(); - let payloads = if non_null_count == 0 { - Vec::new() - } else { - dataset - .read_blobs(column)? - .with_row_ids(non_null_row_ids) - .preserve_order(true) - .execute() - .await? - }; - - if payloads.len() != non_null_count { - return Err(Error::Runtime { - message: format!( - "blob read for column '{column}' returned {} payloads for {} non-null rows", - payloads.len(), - non_null_count - ), - }); - } + let payloads = dataset + .read_blobs(column)? + .with_row_ids(row_ids.to_vec()) + .preserve_order(true) + .execute() + .await?; + ensure_all_row_ids_resolved(column, row_ids.len(), payloads.len())?; let mut builder = LargeBinaryBuilder::new(); - let mut payload_idx = 0; - for is_null in &null_mask { - if *is_null { - builder.append_null(); - } else { - if let Some(data) = &payloads[payload_idx].data { - builder.append_value(data); - } else { - builder.append_null(); - } - payload_idx += 1; + for payload in payloads { + match payload.data { + Some(data) => builder.append_value(data), + None => builder.append_null(), } } Ok(builder.finish()) @@ -295,34 +262,9 @@ pub(crate) async fn take_blob_files_aligned( return Ok(Vec::new()); } - let null_mask = blob_null_mask(dataset, column, row_ids).await?; - let non_null_row_ids = non_null_row_ids(row_ids, &null_mask); - let handles = if non_null_row_ids.is_empty() { - Vec::new() - } else { - dataset.take_blobs(&non_null_row_ids, column).await? - }; - if handles.len() != non_null_row_ids.len() { - return Err(Error::Runtime { - message: format!( - "blob take for column '{column}' returned {} handles for {} non-null rows", - handles.len(), - non_null_row_ids.len() - ), - }); - } - - let mut handles = handles.into_iter(); - Ok(null_mask - .iter() - .map(|is_null| { - if *is_null { - None - } else { - handles.next().flatten() - } - }) - .collect()) + let handles = dataset.take_blobs(row_ids, column).await?; + ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?; + Ok(handles) } #[cfg(test)] diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 35a564d41..3cdf33615 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -201,7 +201,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..0e9354ffc 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, @@ -1019,8 +1030,9 @@ impl Table { /// Materialize blob bytes for the given row ids. /// - /// Output matches `row_ids` in length and order. Null and zero-length rows - /// are null. Prefer [`Self::fetch_blob_files`] for large selections. + /// Output matches `row_ids` in length and order. Null blobs are null; + /// valid empty blobs contain empty byte strings. Prefer + /// [`Self::fetch_blob_files`] for large selections. /// /// ``` /// use arrow_array::UInt64Array; @@ -1055,6 +1067,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 an error 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 +3123,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..47b80b433 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,73 @@ 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!(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!(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(); @@ -617,6 +684,32 @@ async fn fetch_blobs_out_of_range_id_errors_without_panic() -> Result<()> { Ok(()) } +#[tokio::test] +async fn fetch_blob_apis_reject_mixed_valid_and_missing_row_ids() -> 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 row_id = collect_row_ids(&table).await?[0]; + let row_ids = [u64::MAX, row_id]; + + let err = table.fetch_blobs("image", &row_ids).await.unwrap_err(); + assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("row ids")); + + let err = table.fetch_blob_files("image", &row_ids).await.unwrap_err(); + assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("row ids")); + + let requests = row_ids.map(|row_id| BlobRangeRequest::new(row_id, 0, 1)); + let err = table + .fetch_blob_ranges("image", requests) + .await + .unwrap_err(); + assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("row ids")); + Ok(()) +} + #[tokio::test] async fn fetch_blobs_rejects_non_blob_column() -> Result<()> { let tmp = tempdir().unwrap(); @@ -843,11 +936,24 @@ 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(()) } #[tokio::test] -async fn zero_length_blob_reads_back_as_null() -> Result<()> { +async fn empty_blob_reads_back_as_empty_bytes() -> 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"".as_slice())]).await?; @@ -855,7 +961,8 @@ async fn zero_length_blob_reads_back_as_null() -> Result<()> { let ids = collect_row_ids(&table).await?; let bytes = table.fetch_blobs("image", &ids).await?; assert_eq!(bytes.len(), 1); - assert!(bytes.is_null(0)); + assert!(!bytes.is_null(0)); + assert!(bytes.value(0).is_empty()); Ok(()) } @@ -927,6 +1034,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();