From 3dd9c598e9f0cd6004bd0bd30a67f9b279cef8c1 Mon Sep 17 00:00:00 2001 From: Drew Gallardo Date: Mon, 3 Aug 2026 08:38:08 -0700 Subject: [PATCH] feat(remote): add seekable blob range reads (#3750) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Implements Cloud `fetch_blob_files`: returns real seekable `BlobFile` handles over HTTP Range instead of `NotSupported`. - Completes the second Cloud blob read verb after #3684 (`fetch_blobs` = eager whole bytes; this = lazy / partial / sequential reads). - Same public handle API as local (`read_range`, `read_up_to`, `seek`, `tell`, `close`), so one code path works for local and Cloud. Large blobs (video, audio, PDFs) should not require downloading the whole object to inspect a header or stream a slice. After search, callers open a handle and read only what they need: ```python hits = table.search(vec).select(["id", "video"]).limit(5).to_arrow() with table.fetch_blob_files("video", hits)[0] as f: header = f.read_range(0, 256) f.seek(keyframe_offset) chunk = f.read_up_to(1 << 20) ``` ### Behavior - Handle creation probes size with `bytes=0-0` (bounded concurrency, input order preserved). - `204` → null (`None`); `416` with `bytes */0` → valid empty blob; other `416` → error. - `read_range` validates `Content-Range` and body length; OOB ranges fail with `invalid_input` before the request (aligned with Lance). - `read_up_to` reuses one open-ended Range response across sequential reads; `seek` drops it. - Servers older than 0.5.0 get a clear `NotSupported` (does not suggest `fetch_blobs`, which they also lack). ## Testing - `cargo test --features remote -p lancedb remote_blob` - `cargo test --features remote -p lancedb test_blob` - `cargo clippy --features remote --tests --examples` (no new warnings from this change) --------- Co-authored-by: Claude Fable 5 Co-authored-by: Cursor --- python/python/lancedb/table.py | 4 +- python/python/tests/test_remote_db.py | 35 +- python/src/table.rs | 36 +- rust/lancedb/src/blob.rs | 200 +++- rust/lancedb/src/remote/table.rs | 27 +- rust/lancedb/src/remote/table/blobs.rs | 1229 +++++++++++++++++++++++- rust/lancedb/src/table.rs | 2 +- 7 files changed, 1473 insertions(+), 60 deletions(-) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index ea0f97621..31a70c298 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1595,8 +1595,8 @@ class Table(ABC): 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``. Unsupported on LanceDB Cloud, where - :meth:`fetch_blobs` returns full bytes instead. + rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or + newer. """ @abstractmethod diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 757d8a758..d5d3569d3 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2055,6 +2055,24 @@ def blob_remote_table(*, server_version=Version("0.5.0")): request.send_header("phalanx-version", str(server_version)) request.end_headers() request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode()) + elif request.path.startswith("/v1/table/test/blob/image/"): + path = request.path.partition("?")[0] + row_id = int(path.split("/")[-2]) + payload = {10: b"alpha", 20: None, 30: b"gamma"}[row_id] + if payload is None: + request.send_response(204) + request.end_headers() + return + byte_range = request.headers["Range"].removeprefix("bytes=") + start_text, end_text = byte_range.split("-", maxsplit=1) + start = int(start_text) + end = int(end_text) if end_text else len(payload) - 1 + chunk = payload[start : end + 1] + request.send_response(206) + request.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}") + request.send_header("Content-Length", str(len(chunk))) + request.end_headers() + request.wfile.write(chunk) elif request.path == "/v1/table/test/query/": content_len = int(request.headers.get("Content-Length", 0)) body = json.loads(request.rfile.read(content_len)) @@ -2092,8 +2110,21 @@ def test_remote_blob_columns_and_fetch(): assert table.blob_columns() == ["image"] blobs = table.fetch_blobs("image", [10, 20, 30]) assert blobs.to_pylist() == [b"alpha", None, b"gamma"] - with pytest.raises(NotImplementedError, match="Use fetch_blobs for full bytes"): - table.fetch_blob_files("image", [10, 20, 30]) + + +def test_remote_blob_files_are_lazy_seekable_handles(): + with blob_remote_table() as table: + files = table.fetch_blob_files("image", [10, 20, 30]) + + assert len(files) == 3 + alpha, null_row, gamma = files + assert null_row is None + assert alpha is not None + assert gamma is not None + assert alpha.size() == 5 + assert alpha.read_range(1, 3) == b"lph" + gamma.seek(2) + assert gamma.read() == b"mma" def test_remote_blob_fetch_accepts_query_table(): diff --git a/python/src/table.rs b/python/src/table.rs index a4e50c183..d69609093 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -426,9 +426,11 @@ pub struct PyBlobFile { impl PyBlobFile { fn read_bytes(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner.clone(); - let bytes = block_on(async move { inner.read().await }) + let py = self_.py(); + let bytes = py + .detach(move || block_on(async move { inner.read().await })) .map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?; - Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind()) + Ok(PyBytes::new(py, bytes.as_ref()).unbind()) } pub fn read(self_: PyRef<'_, Self>) -> PyResult> { @@ -444,24 +446,32 @@ impl PyBlobFile { fn close(self_: PyRef<'_, Self>) -> PyResult<()> { let inner = self_.inner.clone(); - block_on(async move { inner.close().await }) + self_ + .py() + .detach(move || block_on(async move { inner.close().await })) .map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}"))) } fn is_closed(self_: PyRef<'_, Self>) -> bool { let inner = self_.inner.clone(); - block_on(async move { inner.is_closed().await }) + self_ + .py() + .detach(move || block_on(async move { inner.is_closed().await })) } fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> { let inner = self_.inner.clone(); - block_on(async move { inner.seek(position).await }) + self_ + .py() + .detach(move || block_on(async move { inner.seek(position).await })) .map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}"))) } fn tell(self_: PyRef<'_, Self>) -> PyResult { let inner = self_.inner.clone(); - block_on(async move { inner.tell().await }) + self_ + .py() + .detach(move || block_on(async move { inner.tell().await })) .map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}"))) } @@ -475,16 +485,20 @@ impl PyBlobFile { .checked_add(length as u64) .ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?; let inner = self_.inner.clone(); - let bytes = block_on(async move { inner.read_range(offset..end).await }) + let py = self_.py(); + let bytes = py + .detach(move || block_on(async move { inner.read_range(offset..end).await })) .map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?; - Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind()) + Ok(PyBytes::new(py, bytes.as_ref()).unbind()) } fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult> { let inner = self_.inner.clone(); - let bytes = block_on(async move { inner.read_up_to(length).await }) - .map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?; - Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind()) + let py = self_.py(); + let bytes = py + .detach(move || block_on(async move { inner.read_up_to(length).await })) + .map_err(|e| PyRuntimeError::new_err(format!("blob read_up_to failed: {e}")))?; + Ok(PyBytes::new(py, bytes.as_ref()).unbind()) } } diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index a21fe3c1b..3448257ab 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -9,6 +9,7 @@ //! //! Blob tables require Lance file format >= 2.2 and stable row ids at create. +use std::ops::Range; use std::sync::Arc; use arrow_array::LargeBinaryArray; @@ -17,10 +18,202 @@ use arrow_schema::{DataType, Field, Schema}; use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams}; use lance_arrow::FieldExt; use lance_encoding::version::LanceFileVersion; +use lance_io::object_store::ObjectStore; +use object_store::path::Path; use crate::error::{Error, Result}; -pub use lance::dataset::BlobFile; +/// Seekable handle for one blob value, backed by local storage or a remote +/// HTTP byte-range endpoint. +#[derive(Debug)] +pub struct BlobFile { + inner: BlobFileInner, +} + +#[derive(Debug)] +enum BlobFileInner { + Native(lance::dataset::BlobFile), + #[cfg(feature = "remote")] + Remote(Box), +} + +impl From for BlobFile { + fn from(value: lance::dataset::BlobFile) -> Self { + Self { + inner: BlobFileInner::Native(value), + } + } +} + +#[cfg(feature = "remote")] +impl From for BlobFile { + fn from(value: crate::remote::table::blobs::RemoteBlobFile) -> Self { + Self { + inner: BlobFileInner::Remote(Box::new(value)), + } + } +} + +impl BlobFile { + /// Inline reader over a data-file slice. + pub fn new_inline( + object_store: Arc, + path: Path, + position: u64, + size: u64, + ) -> Self { + lance::dataset::BlobFile::new_inline(object_store, path, position, size).into() + } + + /// Dedicated sidecar-file reader. + pub fn new_dedicated(object_store: Arc, path: Path, size: u64) -> Self { + lance::dataset::BlobFile::new_dedicated(object_store, path, size).into() + } + + /// Packed reader for a slice in a shared sidecar. + pub fn new_packed( + object_store: Arc, + path: Path, + position: u64, + size: u64, + ) -> Self { + lance::dataset::BlobFile::new_packed(object_store, path, position, size).into() + } + + /// External reader at a resolved object location. + pub fn new_external( + object_store: Arc, + path: Path, + uri: String, + position: u64, + size: u64, + ) -> Self { + lance::dataset::BlobFile::new_external(object_store, path, uri, position, size).into() + } + + /// Close the handle. + pub async fn close(&self) -> lance_core::Result<()> { + match &self.inner { + BlobFileInner::Native(file) => file.close().await, + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.close().await, + } + } + + /// Whether the handle is closed. + pub async fn is_closed(&self) -> bool { + match &self.inner { + BlobFileInner::Native(file) => file.is_closed().await, + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.is_closed(), + } + } + + /// Read a range without moving the cursor. + pub async fn read_range(&self, range: Range) -> lance_core::Result { + match &self.inner { + BlobFileInner::Native(file) => file.read_range(range).await, + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.read_range(range).await, + } + } + + /// Read ranges without moving the cursor. + pub async fn read_ranges( + &self, + ranges: &[Range], + ) -> lance_core::Result> { + match &self.inner { + BlobFileInner::Native(file) => file.read_ranges(ranges).await, + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.read_ranges(ranges).await, + } + } + + /// Read from the cursor to the end. + pub async fn read(&self) -> lance_core::Result { + match &self.inner { + BlobFileInner::Native(file) => file.read().await, + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.read().await, + } + } + + /// Read up to `len` bytes and advance the cursor. + pub async fn read_up_to(&self, len: usize) -> lance_core::Result { + match &self.inner { + BlobFileInner::Native(file) => file.read_up_to(len).await, + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.read_up_to(len).await, + } + } + + /// Move the cursor to `new_cursor`. + pub async fn seek(&self, new_cursor: u64) -> lance_core::Result<()> { + match &self.inner { + BlobFileInner::Native(file) => file.seek(new_cursor).await, + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.seek(new_cursor).await, + } + } + + /// Current cursor position. + pub async fn tell(&self) -> lance_core::Result { + match &self.inner { + BlobFileInner::Native(file) => file.tell().await, + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.tell().await, + } + } + + /// Blob length in bytes. + pub fn size(&self) -> u64 { + match &self.inner { + BlobFileInner::Native(file) => file.size(), + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.size(), + } + } + + /// Physical byte offset in the data file. `None` on remote handles. The + /// Cloud byte-range route does not expose storage layout. + pub fn position(&self) -> Option { + match &self.inner { + BlobFileInner::Native(file) => Some(file.position()), + #[cfg(feature = "remote")] + BlobFileInner::Remote(_) => None, + } + } + + /// Path of the data file holding the blob. `None` on remote handles. The + /// Cloud byte-range route does not expose storage layout. + pub fn data_path(&self) -> Option<&Path> { + match &self.inner { + BlobFileInner::Native(file) => Some(file.data_path()), + #[cfg(feature = "remote")] + BlobFileInner::Remote(_) => None, + } + } + + /// Native storage layout. `None` on remote handles. The Cloud byte-range + /// route does not expose layout. + pub fn kind(&self) -> Option { + match &self.inner { + BlobFileInner::Native(file) => Some(file.kind()), + #[cfg(feature = "remote")] + BlobFileInner::Remote(_) => None, + } + } + + /// External URI for native handles. Remote handles do not expose storage URIs. + pub fn uri(&self) -> Option<&str> { + match &self.inner { + BlobFileInner::Native(file) => file.uri(), + #[cfg(feature = "remote")] + BlobFileInner::Remote(_) => None, + } + } +} /// One row-specific blob range read request. /// @@ -264,7 +457,10 @@ pub(crate) async fn take_blob_files_aligned( let handles = dataset.take_blobs(row_ids, column).await?; ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?; - Ok(handles) + Ok(handles + .into_iter() + .map(|handle| handle.map(Into::into)) + .collect()) } #[cfg(test)] diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index e9c200521..de3bbc94a 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -mod blobs; +pub mod blobs; pub mod insert; use self::insert::{RemoteWriteExec, WriteOp}; @@ -4300,32 +4300,9 @@ mod tests { "fetch_blobs", ); - let message = table - .fetch_blob_files("image", &[1]) - .await - .unwrap_err() - .to_string(); - assert!( - message.contains("fetch_blob_files is not supported on LanceDB Cloud"), - "got: {message}" - ); - assert!( - !message.contains("Use fetch_blobs"), - "old server must not be told to use fetch_blobs, got: {message}" - ); - } - - #[tokio::test] - async fn test_blob_files_point_at_fetch_blobs_on_a_blob_capable_server() { - let table = Table::new_with_handler_version( - "my_table", - semver::Version::new(0, 5, 0), - |_| -> http::Response { panic!("fetch_blob_files must not reach the server") }, - ); - assert_not_supported_error( table.fetch_blob_files("image", &[1]).await.unwrap_err(), - "Use fetch_blobs for full bytes", + "requires LanceDB Cloud server 0.5.0 or newer", ); } diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index 550f72a4b..bb8a5d103 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -1,21 +1,343 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! Cloud blob column listing and whole-byte fetch. +//! Cloud blob column listing, whole-byte fetch, and seekable HTTP byte-range handles. +use std::ops::Range; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use arrow_array::{Array, LargeBinaryArray}; use arrow_schema::DataType; -use futures::TryStreamExt; +use bytes::{Bytes, BytesMut}; +use futures::{StreamExt, TryStreamExt}; +use reqwest::{Response, StatusCode, header}; +use tokio::sync::Mutex; use crate::Error; use crate::blob::BlobFile; use crate::error::Result; -use crate::remote::client::HttpSend; +use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; use crate::table::BaseTable; -use super::RemoteTable; +use super::{FreshnessHeaders, RemoteTable}; + +#[derive(Debug, Clone, Copy)] +enum RangeRequestMode { + SizeProbe, + DataRead, +} + +#[async_trait::async_trait] +trait BlobRangeRequester: Send + Sync + std::fmt::Debug { + async fn request_range( + &self, + range_header: &str, + mode: RangeRequestMode, + ) -> Result<(String, Response)>; +} + +#[derive(Debug)] +struct TableBlobRangeRequester { + client: RestfulLanceDbClient, + path: String, + version: Option, + branch: Option, + freshness: FreshnessHeaders, +} + +#[async_trait::async_trait] +impl BlobRangeRequester for TableBlobRangeRequester { + async fn request_range( + &self, + range_header: &str, + mode: RangeRequestMode, + ) -> Result<(String, Response)> { + let mut request = self + .freshness + .apply(self.client.get(&self.path)) + .header(header::RANGE, range_header); + if let Some(version) = self.version { + request = request.query(&[("version", version)]); + } + if let Some(branch) = &self.branch { + request = request.query(&[("branch", branch)]); + } + let (request_id, response) = self.client.send_with_retry(request, None, true).await?; + // Preserve 416 size-probe responses so the caller can detect empty blobs. + if response.status() == StatusCode::RANGE_NOT_SATISFIABLE + && matches!(mode, RangeRequestMode::SizeProbe) + { + return Ok((request_id, response)); + } + let response = self.client.check_response(&request_id, response).await?; + Ok((request_id, response)) + } +} + +#[derive(Debug)] +struct SequentialResponse { + response: Response, + request_id: String, + buffered: Bytes, +} + +#[derive(Debug, Default)] +struct RemoteBlobState { + cursor: u64, + sequential_response: Option, +} + +/// Seekable Cloud blob handle over HTTP Range. +#[derive(Debug)] +pub(crate) struct RemoteBlobFile { + requester: Arc, + state: Mutex, + closed: AtomicBool, + size: u64, +} + +impl RemoteBlobFile { + fn new(requester: Arc, size: u64) -> Self { + Self { + requester, + state: Mutex::new(RemoteBlobState::default()), + closed: AtomicBool::new(false), + size, + } + } + + /// Close the handle without waiting for an in-flight read. + pub(crate) async fn close(&self) -> lance_core::Result<()> { + self.closed.store(true, Ordering::Release); + // Drop a retained response when the state lock is immediately available. + // A reader holding the lock drops it instead once it observes the flag. + if let Ok(mut state) = self.state.try_lock() { + state.sequential_response = None; + } + Ok(()) + } + + pub(crate) fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) + } + + fn ensure_open(&self) -> lance_core::Result<()> { + if self.closed.load(Ordering::Acquire) { + Err(lance_core::Error::invalid_input( + "blob file is already closed", + )) + } else { + Ok(()) + } + } + + pub(crate) async fn read_range(&self, range: Range) -> lance_core::Result { + self.ensure_open()?; + if range.start > range.end { + return Err(lance_core::Error::invalid_input(format!( + "blob range start {} exceeds end {}", + range.start, range.end + ))); + } + if range.end > self.size { + return Err(lance_core::Error::invalid_input(format!( + "blob range end {} exceeds blob size {}", + range.end, self.size + ))); + } + if range.is_empty() { + return Ok(Bytes::new()); + } + let range_header = format!("bytes={}-{}", range.start, range.end - 1); + let (request_id, response) = self + .requester + .request_range(&range_header, RangeRequestMode::DataRead) + .await + .map_err(remote_blob_error)?; + self.ensure_open()?; + validate_partial_response(&response, range.clone(), self.size)?; + let bytes = response + .bytes() + .await + .err_to_http(request_id) + .map_err(remote_blob_error)?; + self.ensure_open()?; + if bytes.len() as u64 != range.end - range.start { + return Err(remote_blob_error(format!( + "byte range returned {} bytes, expected {}", + bytes.len(), + range.end - range.start + ))); + } + Ok(bytes) + } + + /// Read ranges concurrently while preserving input order. + pub(crate) async fn read_ranges( + &self, + ranges: &[Range], + ) -> lance_core::Result> { + futures::stream::iter(ranges.iter().cloned().map(|range| self.read_range(range))) + .buffered(BLOB_REQUEST_CONCURRENCY) + .try_collect() + .await + } + + /// Read from the cursor to the end of the blob. + /// + /// Holds the state lock across cursor calculation and reading so a concurrent + /// seek cannot change the cursor between them. + pub(crate) async fn read(&self) -> lance_core::Result { + self.ensure_open()?; + let mut state = self.state.lock().await; + self.ensure_open()?; + let remaining = self.size.saturating_sub(state.cursor); + let remaining = usize::try_from(remaining).map_err(|_| { + lance_core::Error::invalid_input("remaining blob length exceeds addressable memory") + })?; + self.read_up_to_locked(&mut state, remaining).await + } + + pub(crate) async fn read_up_to(&self, len: usize) -> lance_core::Result { + self.ensure_open()?; + let mut state = self.state.lock().await; + self.ensure_open()?; + self.read_up_to_locked(&mut state, len).await + } + + /// Read up to `len` bytes using caller-validated, locked state. + async fn read_up_to_locked( + &self, + state: &mut RemoteBlobState, + len: usize, + ) -> lance_core::Result { + let target_len = self.size.saturating_sub(state.cursor).min(len as u64) as usize; + if target_len == 0 { + return Ok(Bytes::new()); + } + + // Remove the retained response from shared state before awaiting. Failed or + // cancelled reads leave the committed cursor unchanged and force the next + // read to open a fresh response. + let mut sequential_response = state.sequential_response.take(); + let mut cursor = state.cursor; + let mut output = BytesMut::with_capacity(target_len); + while output.len() < target_len { + if sequential_response.is_none() { + let range_header = format!("bytes={cursor}-"); + let (request_id, response) = self + .requester + .request_range(&range_header, RangeRequestMode::DataRead) + .await + .map_err(remote_blob_error)?; + self.ensure_open()?; + validate_partial_response(&response, cursor..self.size, self.size)?; + sequential_response = Some(SequentialResponse { + response, + request_id, + buffered: Bytes::new(), + }); + } + + let needed = target_len - output.len(); + let active = sequential_response.as_mut().unwrap(); + if !active.buffered.is_empty() { + let take = needed.min(active.buffered.len()); + output.extend_from_slice(&active.buffered.split_to(take)); + cursor += take as u64; + continue; + } + let chunk = active + .response + .chunk() + .await + .err_to_http(active.request_id.clone()) + .map_err(remote_blob_error)?; + self.ensure_open()?; + let chunk = chunk.ok_or_else(|| { + remote_blob_error("response ended before the requested blob range") + })?; + active.buffered = chunk; + } + self.ensure_open()?; + state.cursor = cursor; + if state.cursor < self.size { + state.sequential_response = sequential_response; + } + Ok(output.freeze()) + } + + pub(crate) async fn seek(&self, new_cursor: u64) -> lance_core::Result<()> { + self.ensure_open()?; + let mut state = self.state.lock().await; + self.ensure_open()?; + state.sequential_response = None; + state.cursor = new_cursor; + Ok(()) + } + + pub(crate) async fn tell(&self) -> lance_core::Result { + self.ensure_open()?; + let state = self.state.lock().await; + self.ensure_open()?; + Ok(state.cursor) + } + + pub(crate) fn size(&self) -> u64 { + self.size + } +} + +fn remote_blob_error(error: impl std::fmt::Display) -> lance_core::Error { + lance_core::Error::io(format!("remote blob read failed: {error}")) +} + +fn parse_content_range(value: &str) -> Option<(u64, u64, u64)> { + let value = value.strip_prefix("bytes ")?; + let (range, total) = value.split_once('/')?; + let (start, end) = range.split_once('-')?; + Some((start.parse().ok()?, end.parse().ok()?, total.parse().ok()?)) +} + +/// Parse the total from an unsatisfied-range header, `bytes */{total}`. +fn parse_unsatisfied_content_range(value: &str) -> Option { + value.strip_prefix("bytes */")?.parse().ok() +} + +/// Validate a partial response against the requested range and blob size. +/// +/// Reject `200 OK` because it may contain the entire object. +fn validate_partial_response( + response: &Response, + expected: Range, + size: u64, +) -> lance_core::Result<()> { + if response.status() != StatusCode::PARTIAL_CONTENT { + return Err(remote_blob_error(format!( + "expected HTTP 206 Partial Content, got {}", + response.status() + ))); + } + let content_range = response + .headers() + .get(header::CONTENT_RANGE) + .and_then(|value| value.to_str().ok()) + .and_then(parse_content_range) + .ok_or_else(|| remote_blob_error("response is missing a valid Content-Range header"))?; + if content_range != (expected.start, expected.end - 1, size) { + return Err(remote_blob_error(format!( + "expected Content-Range 'bytes {}-{}/{}', got 'bytes {}-{}/{}'", + expected.start, + expected.end - 1, + size, + content_range.0, + content_range.1, + content_range.2 + ))); + } + Ok(()) +} impl RemoteTable { /// Blob v2 columns are marked in field metadata, which `describe` returns. Reading @@ -30,8 +352,7 @@ impl RemoteTable { column: &str, row_ids: &[u64], ) -> Result { - // An empty selection already has its answer, so skip the round trip and the - // server requirement entirely. Local fetch_blobs returns early the same way. + // Empty requests do not require blob-route support. if row_ids.is_empty() { return Ok(LargeBinaryArray::from(Vec::>::new())); } @@ -109,19 +430,893 @@ impl RemoteTable { Ok(blobs) } + /// Open seekable handles for `row_ids`. + /// + /// Each non-null row is probed once to determine its size. pub(super) async fn fetch_blob_files_impl( &self, - _column: &str, - _row_ids: &[u64], + column: &str, + row_ids: &[u64], ) -> Result>> { - let message = if self.server_version.support_blobs() { - "fetch_blob_files is not supported on LanceDB Cloud yet. \ - Use fetch_blobs for full bytes" - } else { - "fetch_blob_files is not supported on LanceDB Cloud" - }; - Err(Error::NotSupported { - message: message.into(), - }) + // Empty requests do not require blob-route support. + if row_ids.is_empty() { + return Ok(Vec::new()); + } + if !self.server_version.support_blobs() { + return Err(Error::NotSupported { + message: "fetch_blob_files requires LanceDB Cloud server 0.5.0 or newer".into(), + }); + } + + let version = self.current_version().await; + let freshness = self.snapshot_freshness_headers(); + let encoded_column = urlencoding::encode(column); + let requesters = row_ids + .iter() + .map(|row_id| { + let path = format!( + "/v1/table/{}/blob/{encoded_column}/{row_id}/bytes", + self.identifier + ); + let requester: Arc = Arc::new(TableBlobRangeRequester { + client: self.client.clone(), + path, + version, + branch: self.branch.clone(), + freshness, + }); + requester + }) + .collect(); + probe_blob_files(requesters).await + } +} + +/// Probe blob sizes while preserving row order. +async fn probe_blob_files( + requesters: Vec>, +) -> Result>> { + // Collect before buffering to satisfy the async trait lifetime bounds. + let probe_futures: Vec<_> = requesters.into_iter().map(probe_blob_file).collect(); + futures::stream::iter(probe_futures) + .buffered(BLOB_REQUEST_CONCURRENCY) + .try_collect() + .await +} + +const BLOB_REQUEST_CONCURRENCY: usize = 8; + +/// Probe one blob's size. +/// +/// `204` represents null. `416` with `bytes */0` represents an empty blob. +async fn probe_blob_file(requester: Arc) -> Result> { + let (request_id, response) = requester + .request_range("bytes=0-0", RangeRequestMode::SizeProbe) + .await?; + match response.status() { + StatusCode::NO_CONTENT => { + response.bytes().await.err_to_http(request_id)?; + Ok(None) + } + StatusCode::RANGE_NOT_SATISFIABLE => { + let total = response + .headers() + .get(header::CONTENT_RANGE) + .and_then(|value| value.to_str().ok()) + .and_then(parse_unsatisfied_content_range); + match total { + Some(0) => {} + Some(total) => { + return Err(Error::Http { + source: format!( + "blob size probe returned HTTP 416 for a {total}-byte blob" + ) + .into(), + request_id, + status_code: Some(StatusCode::RANGE_NOT_SATISFIABLE), + }); + } + None => { + return Err(Error::Http { + source: "blob size probe returned an invalid Content-Range header".into(), + request_id, + status_code: Some(StatusCode::RANGE_NOT_SATISFIABLE), + }); + } + } + Ok(Some(RemoteBlobFile::new(requester, 0).into())) + } + StatusCode::PARTIAL_CONTENT => { + let size = response + .headers() + .get(header::CONTENT_RANGE) + .and_then(|value| value.to_str().ok()) + .and_then(parse_content_range) + .and_then(|(start, end, total)| { + (start == 0 && end == 0 && total > 0).then_some(total) + }) + .ok_or_else(|| Error::Http { + source: "blob size probe returned an invalid Content-Range header".into(), + request_id: request_id.clone(), + status_code: Some(StatusCode::PARTIAL_CONTENT), + })?; + let probe_body = response.bytes().await.err_to_http(request_id.clone())?; + if probe_body.len() != 1 { + return Err(Error::Http { + source: format!( + "blob size probe returned {} bytes, expected 1", + probe_body.len() + ) + .into(), + request_id, + status_code: Some(StatusCode::PARTIAL_CONTENT), + }); + } + Ok(Some(RemoteBlobFile::new(requester, size).into())) + } + status => Err(Error::Http { + source: format!("blob size probe expected HTTP 206 Partial Content, got {status}") + .into(), + request_id, + status_code: Some(status), + }), + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex as StdMutex}; + use std::time::Duration; + + use reqwest::Request; + use semver::Version; + + use super::*; + + const PAYLOAD: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyz"; + + fn null_blob_response() -> http::Response> { + http::Response::builder() + .status(StatusCode::NO_CONTENT) + .body(Vec::new()) + .unwrap() + } + + fn empty_blob_response() -> http::Response> { + http::Response::builder() + .status(StatusCode::RANGE_NOT_SATISFIABLE) + .header(header::CONTENT_RANGE, "bytes */0") + .body(Vec::new()) + .unwrap() + } + + fn range_response(request: &Request, payload: &[u8]) -> http::Response> { + let value = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap(); + let (start, end) = value + .strip_prefix("bytes=") + .unwrap() + .split_once('-') + .unwrap(); + let start = start.parse::().unwrap(); + let end = if end.is_empty() { + payload.len() - 1 + } else { + end.parse::().unwrap() + }; + http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes {start}-{end}/{}", payload.len()), + ) + .body(payload[start..=end].to_vec()) + .unwrap() + } + + fn mock_remote_blob_table( + requests: Arc>>, + ) -> RemoteTable { + RemoteTable::new_mock( + "my_table".to_string(), + move |request| { + assert_eq!(request.method(), reqwest::Method::GET); + let path = request.url().path(); + assert!(path.starts_with("/v1/table/my_table/blob/image/")); + requests.lock().unwrap().push( + request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(), + ); + if path.contains("/20/bytes") { + return null_blob_response(); + } + if path.contains("/30/bytes") { + return empty_blob_response(); + } + range_response(&request, PAYLOAD) + }, + Some(Version::new(0, 5, 0)), + ) + } + + #[tokio::test] + async fn remote_blob_files_probe_sizes_and_preserve_nulls() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table + .fetch_blob_files_impl("image", &[10, 20]) + .await + .unwrap(); + + assert_eq!(files.len(), 2); + assert!(files[1].is_none()); + let file = files[0].take().unwrap(); + assert_eq!(file.size(), PAYLOAD.len() as u64); + assert_eq!( + requests.lock().unwrap().as_slice(), + ["bytes=0-0", "bytes=0-0"] + ); + } + + #[tokio::test] + async fn remote_blob_file_reads_the_requested_range() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + let file = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + assert_eq!(file.read_range(5..12).await.unwrap(), &PAYLOAD[5..12]); + assert!(requests.lock().unwrap().contains(&"bytes=5-11".to_string())); + } + + #[tokio::test] + async fn remote_blob_file_reuses_sequential_response_until_seek() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + let file = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + assert_eq!(file.read_up_to(4).await.unwrap(), b"0123".as_slice()); + assert_eq!(file.read_up_to(3).await.unwrap(), b"456".as_slice()); + file.seek(20).await.unwrap(); + assert_eq!(file.read_up_to(4).await.unwrap(), &PAYLOAD[20..24]); + assert_eq!(file.tell().await.unwrap(), 24); + assert_eq!( + requests.lock().unwrap().as_slice(), + ["bytes=0-0", "bytes=0-", "bytes=20-"] + ); + } + + #[tokio::test] + async fn remote_blob_files_return_a_handle_for_an_empty_blob() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table + .fetch_blob_files_impl("image", &[10, 20, 30]) + .await + .unwrap(); + + assert_eq!(files.len(), 3); + assert!(files[1].is_none()); + let empty = files[2].take().unwrap(); + assert_eq!(empty.size(), 0); + assert!(empty.read_range(0..0).await.unwrap().is_empty()); + assert!(empty.read().await.unwrap().is_empty()); + let nonempty = files[0].take().unwrap(); + assert_eq!(nonempty.size(), PAYLOAD.len() as u64); + } + + #[tokio::test] + async fn remote_blob_files_reject_unsatisfied_probe_for_nonempty_blob() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |_| { + http::Response::builder() + .status(StatusCode::RANGE_NOT_SATISFIABLE) + .header(header::CONTENT_RANGE, "bytes */36") + .body(Vec::new()) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let error = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("blob size probe returned HTTP 416 for a 36-byte blob"), + "got: {error}" + ); + } + + #[tokio::test] + async fn remote_blob_files_reject_a_self_contradictory_probe_response() { + // A zero-length blob must use `416 Content-Range: bytes */0`. + let table = RemoteTable::new_mock( + "my_table".to_string(), + |_| { + http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header(header::CONTENT_RANGE, "bytes 0-0/0") + .body(vec![0u8]) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let error = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("blob size probe returned an invalid Content-Range header"), + "got: {error}" + ); + } + + #[tokio::test] + async fn remote_blob_files_reject_unsupported_server_version() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |_| -> http::Response { + panic!("old servers must be rejected before a range request") + }, + Some(Version::new(0, 4, 9)), + ); + let error = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap_err(); + assert!(matches!(error, Error::NotSupported { .. })); + assert!(error.to_string().contains("0.5.0")); + } + + #[tokio::test] + async fn remote_blob_file_rejects_mismatched_content_range() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = RemoteTable::new_mock( + "my_table".to_string(), + { + let requests = requests.clone(); + move |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + requests.lock().unwrap().push(range.clone()); + if range == "bytes=0-0" { + return range_response(&request, PAYLOAD); + } + http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 6-12/{}", PAYLOAD.len()), + ) + .body(PAYLOAD[6..=12].to_vec()) + .unwrap() + } + }, + Some(Version::new(0, 5, 0)), + ); + let file = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + let error = file.read_range(5..12).await.unwrap_err(); + assert!( + error.to_string().contains("expected Content-Range"), + "got: {error}" + ); + } + + #[tokio::test] + async fn remote_blob_file_rejects_short_response_body() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + move |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap(); + if range == "bytes=0-0" { + return range_response(&request, PAYLOAD); + } + http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 5-11/{}", PAYLOAD.len()), + ) + .body(PAYLOAD[5..=7].to_vec()) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + let file = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + let error = file.read_range(5..12).await.unwrap_err(); + assert!( + error.to_string().contains("returned 3 bytes, expected 7"), + "got: {error}" + ); + } + + #[tokio::test] + async fn remote_blob_file_failed_read_preserves_cursor_and_retries_fresh() { + let sequential_requests = Arc::new(AtomicUsize::new(0)); + let table = RemoteTable::new_mock( + "my_table".to_string(), + { + let sequential_requests = sequential_requests.clone(); + move |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + if range == "bytes=0-0" { + return range_response(&request, PAYLOAD); + } + let attempt = sequential_requests.fetch_add(1, Ordering::SeqCst); + if attempt == 0 { + // End the response five bytes early to simulate a truncated + // sequential read. + return http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-{}/{}", PAYLOAD.len() - 1, PAYLOAD.len()), + ) + .body(PAYLOAD[..5].to_vec()) + .unwrap(); + } + range_response(&request, PAYLOAD) + } + }, + Some(Version::new(0, 5, 0)), + ); + let file = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + let error = file.read_up_to(10).await.unwrap_err(); + assert!( + error + .to_string() + .contains("response ended before the requested blob range"), + "got: {error}" + ); + assert_eq!(file.tell().await.unwrap(), 0); + + // Retry from the last committed cursor with a fresh request. + let retried = file.read_up_to(4).await.unwrap(); + assert_eq!(retried, b"0123".as_slice()); + assert_eq!(sequential_requests.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn remote_blob_file_empty_range_sends_no_request() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + let file = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + assert!(file.read_range(3..3).await.unwrap().is_empty()); + assert_eq!(requests.lock().unwrap().as_slice(), ["bytes=0-0"]); + } + + #[derive(Debug)] + struct CountingProbeRequester { + index: usize, + in_flight: Arc, + max_in_flight: Arc, + } + + #[async_trait::async_trait] + impl BlobRangeRequester for CountingProbeRequester { + async fn request_range( + &self, + range_header: &str, + _mode: RangeRequestMode, + ) -> Result<(String, Response)> { + assert_eq!(range_header, "bytes=0-0"); + let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(now, Ordering::SeqCst); + // Earlier probes sleep longer, so later probes finish first and the + // ordered collection has to do real reordering work. + tokio::time::sleep(Duration::from_millis(20 - self.index as u64)).await; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + let response = http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-0/{}", 100 + self.index), + ) + .body(vec![0u8]) + .unwrap(); + Ok((format!("probe-{}", self.index), Response::from(response))) + } + } + + #[tokio::test] + async fn remote_blob_file_probes_are_bounded_and_preserve_order() { + let in_flight = Arc::new(AtomicUsize::new(0)); + let max_in_flight = Arc::new(AtomicUsize::new(0)); + let probes = (0..16) + .map(|index| { + let requester: Arc = Arc::new(CountingProbeRequester { + index, + in_flight: in_flight.clone(), + max_in_flight: max_in_flight.clone(), + }); + requester + }) + .collect(); + + let files = probe_blob_files(probes).await.unwrap(); + + let sizes: Vec = files.into_iter().map(|file| file.unwrap().size()).collect(); + let expected: Vec = (0..16).map(|index| 100 + index as u64).collect(); + assert_eq!(sizes, expected); + let max = max_in_flight.load(Ordering::SeqCst); + assert!(max > 1, "probes never overlapped"); + assert!( + max <= BLOB_REQUEST_CONCURRENCY, + "{max} probes in flight exceeds the bound" + ); + } + + #[tokio::test] + async fn remote_blob_file_metadata_reports_none() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + + assert_eq!(file.size(), PAYLOAD.len() as u64); + assert_eq!(file.position(), None); + assert_eq!(file.kind(), None); + assert_eq!(file.data_path(), None); + assert_eq!(file.uri(), None); + } + + #[tokio::test] + async fn closed_remote_blob_file_rejects_every_operation_without_requests() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + let probe_requests = requests.lock().unwrap().len(); + + file.close().await.unwrap(); + + assert!(file.is_closed().await); + for error in [ + file.read().await.unwrap_err(), + file.read_range(0..1).await.unwrap_err(), + file.read_ranges(&[0..1, 1..2]).await.unwrap_err(), + file.read_up_to(1).await.unwrap_err(), + file.seek(0).await.unwrap_err(), + file.tell().await.unwrap_err(), + ] { + assert!(error.to_string().contains("already closed"), "got: {error}"); + } + assert_eq!(requests.lock().unwrap().len(), probe_requests); + } + + #[tokio::test] + async fn out_of_range_read_fails_without_a_request() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + + let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + let probe_requests = requests.lock().unwrap().len(); + + let past_end = file.read_range(0..file.size() + 1).await.unwrap_err(); + assert!( + past_end.to_string().contains("exceeds blob size"), + "got: {past_end}" + ); + let inverted = file + .read_range(Range { start: 3, end: 1 }) + .await + .unwrap_err(); + assert!( + inverted.to_string().contains("exceeds end"), + "got: {inverted}" + ); + assert_eq!(requests.lock().unwrap().len(), probe_requests); + } + + #[tokio::test] + async fn data_read_rejects_416_response() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + if range == "bytes=0-0" { + return http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-0/{}", PAYLOAD.len()), + ) + .body(vec![PAYLOAD[0]]) + .unwrap(); + } + http::Response::builder() + .status(StatusCode::RANGE_NOT_SATISFIABLE) + .header(header::CONTENT_RANGE, "bytes */0") + .body(b"stale range".to_vec()) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + + let error = file.read_range(1..3).await.unwrap_err(); + assert!(error.to_string().contains("416"), "got: {error}"); + } + + #[tokio::test] + async fn empty_row_id_lists_bypass_server_version_gate() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |_| -> http::Response { panic!("an empty selection sends no request") }, + Some(Version::new(0, 4, 9)), + ); + + let files = table.fetch_blob_files_impl("image", &[]).await.unwrap(); + assert!(files.is_empty()); + + let blobs = table.fetch_blobs_impl("image", &[]).await.unwrap(); + assert_eq!(blobs.len(), 0); + } + + #[derive(Debug)] + struct CountingRangeRequester { + in_flight: Arc, + max_in_flight: Arc, + } + + #[async_trait::async_trait] + impl BlobRangeRequester for CountingRangeRequester { + async fn request_range( + &self, + range_header: &str, + _mode: RangeRequestMode, + ) -> Result<(String, Response)> { + let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(now, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(5)).await; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + let (start, end) = range_header + .strip_prefix("bytes=") + .unwrap() + .split_once('-') + .unwrap(); + let start = start.parse::().unwrap(); + let end = end.parse::().unwrap(); + let response = http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes {start}-{end}/{}", PAYLOAD.len()), + ) + .body(PAYLOAD[start..=end].to_vec()) + .unwrap(); + Ok(("range".to_string(), Response::from(response))) + } + } + + #[tokio::test] + async fn read_ranges_run_bounded_and_preserve_order() { + let in_flight = Arc::new(AtomicUsize::new(0)); + let max_in_flight = Arc::new(AtomicUsize::new(0)); + let requester: Arc = Arc::new(CountingRangeRequester { + in_flight, + max_in_flight: max_in_flight.clone(), + }); + let file = RemoteBlobFile::new(requester, PAYLOAD.len() as u64); + + let ranges: Vec<_> = (0..16u64).map(|start| start..start + 2).collect(); + let output = file.read_ranges(&ranges).await.unwrap(); + + for (range, bytes) in ranges.iter().zip(&output) { + assert_eq!( + bytes.as_ref(), + &PAYLOAD[range.start as usize..range.end as usize] + ); + } + let max = max_in_flight.load(Ordering::SeqCst); + assert!(max > 1, "range reads never overlapped"); + assert!( + max <= BLOB_REQUEST_CONCURRENCY, + "{max} range reads in flight exceeds the bound" + ); + } + + #[tokio::test] + async fn read_ranges_reject_out_of_bounds_range() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + let file = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + + let oob = file + .read_ranges(&[0..2, 0..PAYLOAD.len() as u64 + 1]) + .await + .unwrap_err(); + assert!(oob.to_string().contains("exceeds blob size"), "got: {oob}"); + } + + #[tokio::test] + async fn range_read_rejects_200_response() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap() + .to_string(); + if range == "bytes=0-0" { + return http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-0/{}", PAYLOAD.len()), + ) + .body(vec![PAYLOAD[0]]) + .unwrap(); + } + http::Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_LENGTH, PAYLOAD.len().to_string()) + .body(PAYLOAD.to_vec()) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let file = files.remove(0).unwrap(); + let error = file.read_range(1..3).await.unwrap_err(); + assert!(error.to_string().contains("206"), "got: {error}"); + } + + #[derive(Debug)] + struct BlockingRangeRequester { + release: Arc, + started: Arc, + } + + #[async_trait::async_trait] + impl BlobRangeRequester for BlockingRangeRequester { + async fn request_range( + &self, + range_header: &str, + _mode: RangeRequestMode, + ) -> Result<(String, Response)> { + if range_header.ends_with('-') { + self.started.notify_one(); + self.release.wait().await; + } + let response = http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header( + header::CONTENT_RANGE, + format!("bytes 0-{}/{}", PAYLOAD.len() - 1, PAYLOAD.len()), + ) + .body(PAYLOAD.to_vec()) + .unwrap(); + Ok(("hung".to_string(), Response::from(response))) + } + } + + #[tokio::test] + async fn close_returns_while_a_sequential_read_is_in_flight() { + let release = Arc::new(tokio::sync::Barrier::new(2)); + let started = Arc::new(tokio::sync::Notify::new()); + let requester: Arc = Arc::new(BlockingRangeRequester { + release: release.clone(), + started: started.clone(), + }); + let file = Arc::new(RemoteBlobFile::new(requester, PAYLOAD.len() as u64)); + + let reader = { + let file = file.clone(); + tokio::spawn(async move { file.read_up_to(4).await }) + }; + started.notified().await; + + tokio::time::timeout(Duration::from_millis(50), file.close()) + .await + .expect("close waited on the hung read") + .unwrap(); + assert!(file.is_closed()); + + release.wait().await; + let error = reader.await.unwrap().unwrap_err(); + assert!(error.to_string().contains("already closed"), "got: {error}"); + assert!( + file.read_range(0..1) + .await + .unwrap_err() + .to_string() + .contains("already closed") + ); } } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 1c9d68f4c..6f8ecaf14 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -3,6 +3,7 @@ //! LanceDB Table APIs +use crate::blob::BlobFile; use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader}; use arrow_schema::{Schema, SchemaRef}; use async_trait::async_trait; @@ -12,7 +13,6 @@ use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::display::DisplayableExecutionPlan; use futures::StreamExt; use futures::stream::FuturesUnordered; -use lance::dataset::BlobFile; pub use lance::dataset::ColumnAlteration; pub use lance::dataset::NewColumnTransform; pub use lance::dataset::ReadParams;