rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-14 14:46:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support batched blob range reads

Browse Source
This commit is contained in:
Xuanwo
2026-07-22 15:29:53 +08:00
parent 2ce88f8e02
commit 0e5695d313
9 changed files with 396 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/_lancedb.pyi
+5
View File
@@ -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]]: ...
python/python/lancedb/remote/table.py
+5
View File
@@ -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"
python/python/lancedb/table.py
+32
View File
@@ -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]]":
python/python/tests/test_blob.py
+52
View File
@@ -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(
python/src/table.rs
+22 -1
View File
@@ -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<Bound<'_, PyAny>> {
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::<Vec<_>>();
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(
rust/lancedb/src/blob.rs
+114 -1
View File
@@ -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<data, uri>` 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<u64> {
.collect()
}
/// Materialize blob-local ranges (same length and order as `requests`, nulls preserved).
pub(crate) async fn take_blob_ranges_aligned(
dataset: &Arc<Dataset>,
column: &str,
requests: &[BlobRangeRequest],
) -> Result<LargeBinaryArray> {
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::<Vec<_>>();
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::<Vec<_>>();
if !missing_request_indices.is_empty() {
let missing_row_ids = missing_request_indices
.iter()
.map(|request_index| requests[*request_index].row_id)
.collect::<Vec<_>>();
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<Dataset>,
rust/lancedb/src/lib.rs
+1 -1
View File
@@ -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;
rust/lancedb/src/table.rs
+61
View File
@@ -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<LargeBinaryArray> {
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<dyn std::error::Error>> {
/// 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<str>,
requests: impl IntoIterator<Item = BlobRangeRequest>,
) -> Result<LargeBinaryArray> {
let requests = requests.into_iter().collect::<Vec<_>>();
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<LargeBinaryArray> {
let dataset = self.dataset.get().await?;
crate::blob::take_blob_ranges_aligned(&dataset, column, requests).await
}
async fn fetch_blob_files(
&self,
column: &str,
rust/lancedb/tests/blob_integration.rs
+104 -1
View File
@@ -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();