diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index 338c973be..926769f48 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -14,14 +14,10 @@ import pyarrow as pa from .expr import Expr from .schema import blob_v2_column_paths from .types import BlobMode, QueryProjection, QueryProjectionSpec -from .util import get_uri_scheme if TYPE_CHECKING: from _typeshed import WriteableBuffer - from .remote.table import RemoteTable - from .table import AsyncTable, Table - BLOB_MODE_TO_HANDLING = { "lazy": "blobs_descriptions", "bytes": "all_binary", @@ -104,22 +100,6 @@ def validate_blob_mode(blob_mode: BlobMode) -> None: raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}") -def supports_blob_auto_row_id(table: Table | AsyncTable | RemoteTable) -> bool: - """Blob auto row-id applies to native tables, not LanceDB Cloud.""" - from .remote.table import RemoteTable - - if isinstance(table, RemoteTable): - return False - - inner = getattr(table, "_inner", None) - if inner is not None: - uri = inner.database().uri - if isinstance(uri, str) and get_uri_scheme(uri) == "db": - return False - - return True - - def projection_includes_blob_column( projection: QueryProjection, blob_columns: Iterable[str], @@ -164,16 +144,14 @@ def v2_projection_needs_row_id( def blob_auto_row_id_for_scan( - table: Table | AsyncTable | RemoteTable, schema: pa.Schema, projection: QueryProjection, *, with_row_id: bool | None, ) -> bool: + """Auto row-id only applies when the caller said nothing about row ids.""" if with_row_id is not None: return False - if not supports_blob_auto_row_id(table): - return False return v2_projection_needs_row_id(schema, projection, with_row_id=False) @@ -186,6 +164,11 @@ def finalize_blob_query_table( ) -> pa.Table: if user_requested_row_id or not blob_auto_row_id: return tbl + if "_rowid" not in tbl.column_names: + # A backend that ignores the row-id request leaves nothing to stash. Hand + # back the projection as-is so fetch_blobs raises the error that names the + # ways to supply row ids, rather than failing here about a hidden column. + return tbl return stash_auto_row_ids(tbl, blob_paths) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index b85f50d7b..181eb9349 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -52,7 +52,6 @@ from ._blob import ( finalize_blob_query_table, replace_v2_blob_columns_with_bytes, replace_v2_blob_columns_with_bytes_sync, - supports_blob_auto_row_id, validate_blob_mode, ) from .types import BlobMode, QueryProjection @@ -1277,10 +1276,7 @@ class LanceQueryBuilder(ABC): return self._with_row_id is True def _blob_auto_row_id_enabled(self) -> bool: - if not supports_blob_auto_row_id(self._table): - return False return blob_auto_row_id_for_scan( - self._table, self._table.schema, self._columns, with_row_id=self._with_row_id, @@ -2771,7 +2767,7 @@ class AsyncQueryBase(object): ) async def _maybe_add_blob_row_id(self) -> None: - if self._table is None or not supports_blob_auto_row_id(self._table): + if self._table is None: self._blob_auto_row_id = False self._blob_paths = () return @@ -2779,7 +2775,6 @@ class AsyncQueryBase(object): req = self._inner.to_query_request() schema = await self._table.schema() self._blob_auto_row_id = blob_auto_row_id_for_scan( - self._table, schema, req.select, with_row_id=self._with_row_id, @@ -3031,7 +3026,6 @@ class AsyncQueryBase(object): schema = await self._table.schema() blob_auto_row_id = blob_auto_row_id_for_scan( - self._table, schema, query.columns, with_row_id=self._with_row_id, @@ -3875,10 +3869,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): req = fts_query._inner.to_query_request() blob_auto_row_id = False blob_paths: tuple[str, ...] = () - if self._table is not None and supports_blob_auto_row_id(self._table): + if self._table is not None: schema = await self._table.schema() blob_auto_row_id = blob_auto_row_id_for_scan( - self._table, schema, req.select, with_row_id=self._with_row_id, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 01b90c019..20a952c53 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -20,6 +20,7 @@ from typing import ( import warnings from lancedb import __version__ +from lancedb._blob import BlobFile from lancedb._lancedb import ( AddColumnsResult, @@ -1037,22 +1038,22 @@ class RemoteTable(Table): ) def blob_columns(self) -> list[str]: - raise NotImplementedError( - "blob_columns() is not yet supported on the LanceDB Cloud" - ) + return LOOP.run(self._table.blob_columns()) - def fetch_blobs(self, column: str, row_ids) -> pa.LargeBinaryArray: - raise NotImplementedError("fetch_blobs() is not supported on LanceDB Cloud") + def fetch_blobs( + self, column: str, row_ids: Union[list[int], pa.Table] + ) -> pa.LargeBinaryArray: + return LOOP.run(self._table.fetch_blobs(column, row_ids)) 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" - ) + def fetch_blob_files( + self, column: str, row_ids: Union[list[int], pa.Table] + ) -> "list[Optional[BlobFile]]": + return LOOP.run(self._table.fetch_blob_files(column, row_ids)) def head(self, n=5) -> pa.Table: """ diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 56764bfd3..26d78a427 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1574,8 +1574,10 @@ class Table(ABC): """Open lazy, seekable :class:`~lancedb._blob.BlobFile` handles. Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is - a ``list[int]`` or query ``pyarrow.Table`` with ``_rowid`` (or stashed - row-id metadata). Null rows are ``None``. Local tables only. + 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. """ @abstractmethod diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 55718f516..dd46628f1 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1020,7 +1020,7 @@ def query_test_table(query_handler, *, server_version=Version("0.1.0")): request.send_header("Content-Type", "application/json") request.send_header("phalanx-version", str(server_version)) request.end_headers() - request.wfile.write(b"{}") + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') elif request.path == "/v1/table/test/query/": content_len = int(request.headers.get("Content-Length")) body = request.rfile.read(content_len) @@ -1858,3 +1858,183 @@ def test_inherited_remote_table_reopens_after_fork(): finally: server.shutdown() server_thread.join() + + +BLOB_DESCRIBE_RESPONSE = { + "table": "test", + "version": 1, + "schema": { + "fields": [ + {"name": "id", "type": {"type": "int64"}, "nullable": False}, + { + "name": "image", + "type": { + "type": "struct", + "fields": [ + { + "name": "data", + "type": {"type": "large_binary"}, + "nullable": True, + }, + {"name": "uri", "type": {"type": "string"}, "nullable": True}, + ], + }, + "nullable": True, + "metadata": { + "ARROW:extension:name": "lance.blob.v2", + "ARROW:extension:metadata": "", + }, + }, + ] + }, +} + + +def blob_query_response_table(): + image_field = pa.field( + "image", + pa.struct( + [ + pa.field("kind", pa.uint8(), nullable=False), + pa.field("position", pa.uint64(), nullable=False), + pa.field("size", pa.uint64(), nullable=False), + pa.field("blob_id", pa.uint32(), nullable=False), + pa.field("blob_uri", pa.string(), nullable=False), + ] + ), + metadata={"lance-encoding:blob": "true"}, + ) + images = pa.StructArray.from_arrays( + [ + pa.array([1, 0, 0], type=pa.uint8()), + pa.array([0, 0, 0], type=pa.uint64()), + pa.array([5, 0, 5], type=pa.uint64()), + pa.array([1, 0, 2], type=pa.uint32()), + pa.array(["", "", ""], type=pa.string()), + ], + fields=image_field.type, + mask=pa.array([False, True, False]), + ) + return pa.Table.from_arrays( + [ + pa.array([1, 2, 3], type=pa.int64()), + images, + pa.array([10, 20, 30], type=pa.uint64()), + ], + schema=pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + image_field, + pa.field("_rowid", pa.uint64()), + ] + ), + ) + + +@contextlib.contextmanager +def blob_remote_table(*, server_version=Version("0.5.0")): + def handler(request): + if request.path == "/v1/table/test/describe/": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.send_header("phalanx-version", str(server_version)) + request.end_headers() + request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode()) + elif request.path == "/v1/table/test/query/": + content_len = int(request.headers.get("Content-Length", 0)) + body = json.loads(request.rfile.read(content_len)) + assert body["columns"] == ["id", "image"] + assert body["with_row_id"] is True + response_table = blob_query_response_table() + request.send_response(200) + request.send_header("Content-Type", "application/vnd.apache.arrow.file") + request.end_headers() + with pa.ipc.new_file(request.wfile, response_table.schema) as writer: + writer.write_table(response_table) + elif request.path == "/v1/table/test/fetch_blobs/": + content_len = int(request.headers.get("Content-Length", 0)) + body = json.loads(request.rfile.read(content_len)) + assert body["column"] == "image" + assert body["row_ids"] == [10, 20, 30] + response_table = pa.table( + {"image": pa.array([b"alpha", None, b"gamma"], type=pa.large_binary())} + ) + request.send_response(200) + request.send_header("Content-Type", "application/vnd.apache.arrow.stream") + request.end_headers() + with pa.ipc.new_stream(request.wfile, response_table.schema) as writer: + writer.write_table(response_table) + else: + request.send_response(404) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + yield db.open_table("test") + + +def test_remote_blob_columns_and_fetch(): + with blob_remote_table() as table: + 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_fetch_accepts_query_table(): + hits = pa.table({"_rowid": pa.array([10, 20, 30], type=pa.uint64())}) + + with blob_remote_table() as table: + blobs = table.fetch_blobs("image", hits) + + assert blobs.to_pylist() == [b"alpha", None, b"gamma"] + + +def test_remote_blob_query_stashes_row_ids_for_fetch(): + with blob_remote_table() as table: + hits = table.search().select(["id", "image"]).limit(3).to_arrow() + assert "_rowid" not in hits.column_names + assert "_lance_row_id" in hits.schema.field("image").type.names + blobs = table.fetch_blobs("image", hits) + + assert blobs.to_pylist() == [b"alpha", None, b"gamma"] + + +def test_remote_blob_query_survives_a_server_that_ignores_the_row_id_request(): + def handler(request): + if request.path == "/v1/table/test/describe/": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.send_header("phalanx-version", "0.5.0") + request.end_headers() + request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode()) + elif request.path == "/v1/table/test/query/": + content_len = int(request.headers.get("Content-Length", 0)) + assert json.loads(request.rfile.read(content_len))["with_row_id"] is True + response_table = blob_query_response_table().drop_columns(["_rowid"]) + request.send_response(200) + request.send_header("Content-Type", "application/vnd.apache.arrow.file") + request.end_headers() + with pa.ipc.new_file(request.wfile, response_table.schema) as writer: + writer.write_table(response_table) + else: + request.send_response(404) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + table = db.open_table("test") + hits = table.search().select(["id", "image"]).limit(3).to_arrow() + + assert hits.column_names == ["id", "image"] + assert "_lance_row_id" not in hits.schema.field("image").type.names + with pytest.raises(ValueError, match="pass a list of row ids"): + table.fetch_blobs("image", hits) + + +def test_remote_blob_byte_apis_not_supported_on_old_server(): + with blob_remote_table(server_version=Version("0.1.0")) as table: + assert table.blob_columns() == ["image"] + with pytest.raises(NotImplementedError, match="not supported"): + table.fetch_blobs("image", [1]) + with pytest.raises(NotImplementedError, match="not supported"): + table.fetch_blob_files("image", [1]) diff --git a/rust/lancedb/src/remote.rs b/rust/lancedb/src/remote.rs index 25e19c537..657ae1a47 100644 --- a/rust/lancedb/src/remote.rs +++ b/rust/lancedb/src/remote.rs @@ -14,7 +14,6 @@ pub(crate) mod table; pub(crate) mod util; const ARROW_STREAM_CONTENT_TYPE: &str = "application/vnd.apache.arrow.stream"; -#[cfg(test)] const ARROW_FILE_CONTENT_TYPE: &str = "application/vnd.apache.arrow.file"; #[cfg(test)] const JSON_CONTENT_TYPE: &str = "application/json"; diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 25fcd23b1..a67f9dbb3 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -79,6 +79,10 @@ impl ServerVersion { pub fn support_multipart_write(&self) -> bool { self.0 >= semver::Version::new(0, 4, 0) } + + pub fn support_blobs(&self) -> bool { + self.0 >= semver::Version::new(0, 5, 0) + } } pub const OPT_REMOTE_PREFIX: &str = "remote_database_"; @@ -661,6 +665,7 @@ impl Database for RemoteDatabase { RemoteTable::::handle_table_not_found(&request.name, rsp, &request_id).await?; let rsp = self.client.check_response(&request_id, rsp).await?; let version = parse_server_version(&request_id, &rsp)?; + let describe_body = rsp.text().await.ok(); let table_identifier = build_table_identifier( &request.name, &request.namespace_path, @@ -673,6 +678,12 @@ impl Database for RemoteDatabase { table_identifier, version, )); + // This describe already carries the schema, so hand it to the table + // instead of making the first schema read fetch it again. A version or + // branch pin applied after this invalidates the cache. + if let Some(body) = &describe_body { + table.seed_schema(body); + } let cache_key = build_cache_key(&request.name, &request.namespace_path); self.table_cache.insert(cache_key, table.clone()).await; Ok(table) @@ -923,6 +934,7 @@ impl From for RemoteOptions { mod tests { use super::{NamespaceHeaderProviderContext, build_cache_key}; use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, OnceLock}; use arrow_array::{Int32Array, RecordBatch}; @@ -1073,6 +1085,46 @@ mod tests { assert_eq!(table.name(), "table1"); } + #[tokio::test] + async fn test_open_table_seeds_the_schema_from_its_describe() { + let describe_calls = Arc::new(AtomicUsize::new(0)); + let counted = describe_calls.clone(); + let conn = Connection::new_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/table/table1/describe/"); + counted.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version": 1, "schema": {"fields": [ + {"name": "id", "type": {"type": "int64"}, "nullable": false} + ]}}"# + .to_string(), + ) + .unwrap() + }); + + let table = conn.open_table("table1").execute().await.unwrap(); + let schema = table.schema().await.unwrap(); + + assert_eq!(schema.field(0).name(), "id"); + assert_eq!(describe_calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_open_table_survives_a_describe_body_it_cannot_parse() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.url().path(), "/v1/table/table1/describe/"); + http::Response::builder() + .status(200) + .body(r#"{"table": "table1"}"#.to_string()) + .unwrap() + }); + + let table = conn.open_table("table1").execute().await.unwrap(); + + assert_eq!(table.name(), "table1"); + } + #[tokio::test] async fn test_open_table_branch_and_version() { let conn = Connection::new_with_handler(|request| { diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index b0d7dacc5..b668f6d67 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1,20 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +mod blobs; pub mod insert; use self::insert::{RemoteWriteExec, WriteOp}; -use crate::expr::expr_to_sql_string; -use crate::table::write_progress::FinishOnDrop; -// Used by the test module below (single-request write requests set these -// headers directly in test handlers); kept at module scope so both the -// library and its tests can name them. -#[cfg(test)] -use super::ARROW_STREAM_CONTENT_TYPE; use super::client::RequestResultExt; use super::client::{HttpSend, RestfulLanceDbClient, Sender}; use super::db::ServerVersion; +use super::{ARROW_FILE_CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE}; +use crate::blob::BlobFile; use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions}; +use crate::expr::expr_to_sql_string; use crate::index::Index; use crate::index::IndexStatistics; use crate::index::waiter::wait_for_index; @@ -31,6 +28,7 @@ use crate::table::Tags; use crate::table::UpdateResult; use crate::table::merge::MergeFilter; use crate::table::query::create_multi_vector_plan; +use crate::table::write_progress::FinishOnDrop; use crate::table::{AlterColumnsResult, FieldMetadataUpdate, UpdateFieldMetadataResult}; use crate::table::{AnyQuery, Filter, Predicate, PreprocessingOutput, TableStatistics}; use crate::utils::background_cache::BackgroundCache; @@ -47,16 +45,15 @@ use crate::{ merge::MergeInsertBuilder, }, }; -use arrow_array::RecordBatchReader; -use arrow_ipc::reader::FileReader; -use arrow_schema::{DataType, SchemaRef}; +use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader}; +use arrow_ipc::reader::{FileReader, StreamReader}; +use arrow_schema::{ArrowError, DataType, SchemaRef}; use async_trait::async_trait; use chrono::{DateTime, Utc}; use datafusion_common::DataFusionError; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ExecutionPlan, RecordBatchStream, SendableRecordBatchStream}; use futures::{StreamExt, TryStreamExt}; -#[cfg(test)] use http::header::CONTENT_TYPE; use http::{HeaderName, StatusCode}; use lance::arrow::json::{JsonDataType, JsonSchema}; @@ -314,6 +311,20 @@ impl RemoteTable { } } + /// Seed the schema cache from a `describe` body the caller already fetched. + /// + /// Best effort. `open_table` succeeds today without reading this body, so a + /// body we cannot parse leaves the cache empty and the next schema read + /// fetches it again through the path that reports a real error. + pub(crate) fn seed_schema(&self, describe_body: &str) { + let Ok(description) = serde_json::from_str::(describe_body) else { + return; + }; + if let Ok(schema) = arrow_schema::Schema::try_from(description.schema) { + self.schema_cache.seed(Arc::new(schema)); + } + } + /// Return a new handle scoped to `branch`, sharing the client but with fresh /// caches and version/freshness state (the branch tracks its own latest). /// Mirrors `NativeTable`'s handle-per-branch model. @@ -526,19 +537,39 @@ impl RemoteTable { result } - async fn read_arrow_stream( + async fn read_arrow_response( &self, request_id: &str, response: reqwest::Response, ) -> Result { let response = self.check_table_response(request_id, response).await?; - // There isn't a way to actually stream this data yet. I have an upstream issue: - // https://github.com/apache/arrow-rs/issues/6420 + // The header has to be read before the body, which consumes the response. + let content_type = response + .headers() + .get(CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + let framing = resolve_arrow_ipc_framing(content_type.as_deref(), request_id)?; + + // Buffer the whole body. File framing keeps its footer at the end, so /query + // cannot decode incrementally. Stream framing could, via + // arrow_ipc::reader::StreamDecoder, but fetch_blobs concatenates every batch + // before returning, so no caller would see data sooner. let body = response.bytes().await.err_to_http(request_id.into())?; - let reader = FileReader::try_new(Cursor::new(body), None)?; - let schema = reader.schema(); - let stream = futures::stream::iter(reader).map_err(DataFusionError::from); + type IpcBatchIterator = + Box> + Send>; + let (schema, batches): (SchemaRef, IpcBatchIterator) = match framing { + ArrowIpcFraming::Stream => { + let reader = StreamReader::try_new(Cursor::new(body), None)?; + (reader.schema(), Box::new(reader)) + } + ArrowIpcFraming::File => { + let reader = FileReader::try_new(Cursor::new(body), None)?; + (reader.schema(), Box::new(reader)) + } + }; + let stream = futures::stream::iter(batches).map_err(DataFusionError::from); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } @@ -910,7 +941,7 @@ impl RemoteTable { let futures = requests.into_iter().map(|req| async move { let (request_id, response) = self.send(req, true).await?; self.track_read_version_from_headers(response.headers()); - self.read_arrow_stream(&request_id, response).await + self.read_arrow_response(&request_id, response).await }); let streams = futures::future::try_join_all(futures); @@ -976,6 +1007,47 @@ struct TableDescription { location: Option, } +/// How a response body frames its Arrow IPC payload. `/query` answers with file framing +/// and `fetch_blobs` with stream framing, so the reader is picked per response. +enum ArrowIpcFraming { + File, + Stream, +} + +/// A response with no `Content-Type` uses file framing, preserving this helper's +/// behavior before `fetch_blobs` introduced stream responses. +fn resolve_arrow_ipc_framing( + content_type: Option<&str>, + request_id: &str, +) -> Result { + let Some(media_type) = content_type.map(base_media_type) else { + return Ok(ArrowIpcFraming::File); + }; + if media_type.eq_ignore_ascii_case(ARROW_STREAM_CONTENT_TYPE) { + return Ok(ArrowIpcFraming::Stream); + } + if media_type.eq_ignore_ascii_case(ARROW_FILE_CONTENT_TYPE) { + return Ok(ArrowIpcFraming::File); + } + Err(Error::Http { + source: format!( + "Expected an Arrow IPC response with Content-Type '{ARROW_STREAM_CONTENT_TYPE}' \ + or '{ARROW_FILE_CONTENT_TYPE}', got '{media_type}'" + ) + .into(), + request_id: request_id.into(), + status_code: None, + }) +} + +/// Strip media-type parameters before matching against the Arrow content types. +fn base_media_type(content_type: &str) -> &str { + match content_type.split_once(';') { + Some((media_type, _parameters)) => media_type.trim(), + None => content_type.trim(), + } +} + /// Extract an Error from Arc, reconstructing if the Arc is shared. /// This is needed because `Shared` futures cache results internally, so /// `Arc::try_unwrap` typically fails. @@ -2008,6 +2080,22 @@ impl BaseTable for RemoteTable { } } + async fn blob_columns(&self) -> Result> { + self.blob_columns_impl().await + } + + async fn fetch_blobs(&self, column: &str, row_ids: &[u64]) -> Result { + self.fetch_blobs_impl(column, row_ids).await + } + + async fn fetch_blob_files( + &self, + column: &str, + row_ids: &[u64], + ) -> Result>> { + self.fetch_blob_files_impl(column, row_ids).await + } + async fn explain_plan(&self, query: &AnyQuery, verbose: bool) -> Result { let base_request = self.post_read(&format!("/v1/table/{}/explain_plan/", self.identifier)); @@ -2911,7 +2999,9 @@ mod tests { use crate::table::{AddDataMode, FieldMetadataUpdate, FtsToken}; use arrow::{array::AsArray, compute::concat_batches, datatypes::Int32Type}; - use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, record_batch}; + use arrow_array::Array; + use arrow_array::builder::LargeBinaryBuilder; + use arrow_array::{BinaryArray, Int32Array, RecordBatch, RecordBatchIterator, record_batch}; use arrow_schema::{DataType, Field, Schema}; use chrono::{DateTime, Utc}; use futures::{StreamExt, TryFutureExt, future::BoxFuture}; @@ -2935,7 +3025,6 @@ mod tests { AnalyzePlanDistributedMetrics, ColumnOrdering, ExecutableQuery, QueryBase, QueryExecutionOptions, }, - remote::ARROW_FILE_CONTENT_TYPE, }; #[tokio::test] @@ -3117,6 +3206,17 @@ mod tests { body } + fn write_ipc_stream_uncompressed(data: &RecordBatch) -> Vec { + let mut body = Vec::new(); + { + let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, &data.schema()) + .expect("Failed to create writer"); + writer.write(data).expect("Failed to write data"); + writer.finish().expect("Failed to finish"); + } + body + } + fn write_ipc_file(data: &RecordBatch) -> Vec { let mut body = Vec::new(); { @@ -3788,6 +3888,503 @@ mod tests { assert_eq!(data[0].as_ref().unwrap(), &expected_data); } + fn blob_describe_response() -> http::Response { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + crate::blob("image", true), + Field::new("caption", DataType::Utf8, true), + crate::blob("thumbnail", true), + ]); + let json_schema = JsonSchema::try_from(&schema).unwrap(); + http::Response::builder() + .status(200) + .body(serde_json::json!({ "version": 1, "schema": json_schema }).to_string()) + .unwrap() + } + + #[rstest] + #[case(semver::Version::new(0, 1, 0))] + #[case(semver::Version::new(0, 5, 0))] + #[tokio::test] + async fn test_blob_columns_read_the_schema_on_any_server_version( + #[case] version: semver::Version, + ) { + let table = Table::new_with_handler_version("my_table", version, |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + blob_describe_response() + }); + + let columns = table.blob_columns().await.unwrap(); + assert_eq!(columns, vec!["image".to_string(), "thumbnail".to_string()]); + } + + #[tokio::test] + async fn test_fetch_blobs_decodes_null_aligned_bytes() { + let mut builder = LargeBinaryBuilder::new(); + builder.append_value(b"alpha"); + builder.append_null(); + builder.append_value(b"gamma"); + let blobs = builder.finish(); + let expected = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "image", + DataType::LargeBinary, + true, + )])), + vec![Arc::new(blobs)], + ) + .unwrap(); + let expected_ref = expected.clone(); + + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + move |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/fetch_blobs/"); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(body["column"], "image"); + assert_eq!(body["row_ids"], serde_json::json!([10, 20, 30])); + + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .body(write_ipc_stream_uncompressed(&expected_ref)) + .unwrap() + }, + ); + + let blobs = table.fetch_blobs("image", &[10, 20, 30]).await.unwrap(); + assert_eq!(blobs.len(), 3); + assert_eq!(blobs.value(0), b"alpha"); + assert!(blobs.is_null(1)); + assert_eq!(blobs.value(2), b"gamma"); + } + + #[tokio::test] + async fn test_fetch_blobs_concatenates_multiple_batches() { + let schema = Arc::new(Schema::new(vec![Field::new( + "image", + DataType::LargeBinary, + true, + )])); + + let mut first_builder = LargeBinaryBuilder::new(); + first_builder.append_value(b"alpha"); + first_builder.append_null(); + let first_batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(first_builder.finish())]).unwrap(); + + let mut second_builder = LargeBinaryBuilder::new(); + second_builder.append_value(b"gamma"); + let second_batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(second_builder.finish())]).unwrap(); + + let mut body = Vec::new(); + { + let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, &schema).unwrap(); + writer.write(&first_batch).unwrap(); + writer.write(&second_batch).unwrap(); + writer.finish().unwrap(); + } + + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/fetch_blobs/"); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .body(body.clone()) + .unwrap() + }, + ); + + let blobs = table.fetch_blobs("image", &[10, 20, 30]).await.unwrap(); + assert_eq!(blobs.len(), 3); + assert_eq!(blobs.value(0), b"alpha"); + assert!(blobs.is_null(1)); + assert_eq!(blobs.value(2), b"gamma"); + } + + fn table_with_fetch_blobs_response(body: Vec) -> Table { + table_with_fetch_blobs_content_type(Some(ARROW_STREAM_CONTENT_TYPE), body) + } + + fn table_with_fetch_blobs_content_type( + content_type: Option<&'static str>, + body: Vec, + ) -> Table { + Table::new_with_handler_version("my_table", semver::Version::new(0, 5, 0), move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/fetch_blobs/"); + let mut response = http::Response::builder().status(200); + if let Some(content_type) = content_type { + response = response.header(CONTENT_TYPE, content_type); + } + response.body(body.clone()).unwrap() + }) + } + + fn one_row_blob_batch(column: &str) -> RecordBatch { + let mut builder = LargeBinaryBuilder::new(); + builder.append_value(b"alpha"); + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + column, + DataType::LargeBinary, + true, + )])), + vec![Arc::new(builder.finish())], + ) + .unwrap() + } + + fn one_row_blob_ipc_stream(column: &str) -> Vec { + write_ipc_stream_uncompressed(&one_row_blob_batch(column)) + } + + #[tokio::test] + async fn test_fetch_blobs_sends_the_checked_out_version() { + let ipc = one_row_blob_ipc_stream("image"); + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + move |request| match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(br#"{"version": 42, "schema": {"fields": []}}"#.to_vec()) + .unwrap(), + "/v1/table/my_table/fetch_blobs/" => { + let body = request_body_json(&request); + assert_eq!( + body["version"], 42, + "blob reads must use the same snapshot as the query" + ); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .body(ipc.clone()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }, + ); + + table.checkout(42).await.unwrap(); + + assert_eq!(table.fetch_blobs("image", &[10]).await.unwrap().len(), 1); + } + + #[tokio::test] + async fn test_fetch_blobs_sends_the_checked_out_branch() { + use lance::dataset::refs::Ref; + let ipc = one_row_blob_ipc_stream("image"); + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body(b"{}".to_vec()) + .unwrap(), + "/v1/table/my_table/fetch_blobs/" => { + let body = request_body_json(&request); + assert_eq!(body["branch"], "exp", "blob reads must stay on the branch"); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .body(ipc.clone()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }, + ); + + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + + assert_eq!(branch.fetch_blobs("image", &[10]).await.unwrap().len(), 1); + } + + #[tokio::test] + async fn test_fetch_blobs_sends_a_nested_column_as_a_dotted_path() { + let ipc = one_row_blob_ipc_stream("info.blob"); + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + move |request| { + let body = request_body_json(&request); + assert_eq!(body["column"], "info.blob"); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .body(ipc.clone()) + .unwrap() + }, + ); + + let blobs = table.fetch_blobs("info.blob", &[10]).await.unwrap(); + + assert_eq!(blobs.value(0), b"alpha"); + } + + fn write_empty_ipc_stream(schema: &Schema) -> Vec { + let mut body = Vec::new(); + arrow_ipc::writer::StreamWriter::try_new(&mut body, schema) + .unwrap() + .finish() + .unwrap(); + body + } + + fn assert_fetch_blobs_http_error(error: Error, expected: &str) { + match error { + Error::Http { + source, request_id, .. + } => { + assert!(source.to_string().contains(expected)); + assert!(!request_id.is_empty()); + } + error => panic!("expected HTTP error, got {error}"), + } + } + + fn assert_not_supported_error(error: Error, expected: &str) { + match error { + Error::NotSupported { message } => assert!(message.contains(expected)), + error => panic!("expected not-supported error, got {error}"), + } + } + + #[tokio::test] + async fn test_fetch_blobs_rejects_missing_column() { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "other", + DataType::LargeBinary, + true, + )])), + vec![Arc::new(LargeBinaryArray::from(vec![Some( + b"value".as_slice(), + )]))], + ) + .unwrap(); + let table = table_with_fetch_blobs_response(write_ipc_stream_uncompressed(&batch)); + + let error = table.fetch_blobs("image", &[10]).await.unwrap_err(); + assert_fetch_blobs_http_error(error, "missing the 'image' column"); + } + + #[tokio::test] + async fn test_fetch_blobs_rejects_wrong_column_type() { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "image", + DataType::Int32, + false, + )])), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + let table = table_with_fetch_blobs_response(write_ipc_stream_uncompressed(&batch)); + + let error = table.fetch_blobs("image", &[10]).await.unwrap_err(); + assert_fetch_blobs_http_error( + error, + "type Int32, expected Binary, LargeBinary, or BinaryView", + ); + } + + #[rstest] + #[case(DataType::Binary)] + #[case(DataType::LargeBinary)] + #[case(DataType::BinaryView)] + #[tokio::test] + async fn test_fetch_blobs_accepts_binary_large_binary_and_binary_view( + #[case] data_type: DataType, + ) { + let binary_values = BinaryArray::from(vec![ + Some(b"alpha".as_slice()), + None, + Some(b"gamma".as_slice()), + ]); + let typed_column = arrow::compute::cast(&binary_values, &data_type).unwrap(); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("image", data_type, true)])), + vec![typed_column], + ) + .unwrap(); + let table = table_with_fetch_blobs_response(write_ipc_stream_uncompressed(&batch)); + + let blobs = table.fetch_blobs("image", &[10, 20, 30]).await.unwrap(); + assert_eq!(blobs.len(), 3); + assert_eq!(blobs.value(0), b"alpha"); + assert!(blobs.is_null(1)); + assert_eq!(blobs.value(2), b"gamma"); + } + + #[tokio::test] + async fn test_fetch_blobs_rejects_row_count_mismatch() { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "image", + DataType::LargeBinary, + true, + )])), + vec![Arc::new(LargeBinaryArray::from(vec![Some( + b"value".as_slice(), + )]))], + ) + .unwrap(); + let table = table_with_fetch_blobs_response(write_ipc_stream_uncompressed(&batch)); + + let error = table.fetch_blobs("image", &[10, 20]).await.unwrap_err(); + assert_fetch_blobs_http_error(error, "returned 1 rows for 2 row ids"); + } + + #[tokio::test] + async fn test_fetch_blobs_rejects_zero_batches_for_nonempty_row_ids() { + let schema = Schema::new(vec![Field::new("image", DataType::LargeBinary, true)]); + let table = table_with_fetch_blobs_response(write_empty_ipc_stream(&schema)); + + let error = table.fetch_blobs("image", &[10]).await.unwrap_err(); + assert_fetch_blobs_http_error(error, "returned 0 rows for 1 row ids"); + } + + #[tokio::test] + async fn test_fetch_blobs_skips_the_request_for_empty_row_ids() { + let table = Table::new_with_handler("my_table", |_| -> http::Response { + panic!("fetch_blobs must not call the server for an empty selection"); + }); + + let blobs = table.fetch_blobs("image", &[]).await.unwrap(); + assert!(blobs.is_empty()); + } + + #[tokio::test] + async fn test_blob_byte_apis_not_supported_on_old_server() { + let table = Table::new_with_handler("my_table", |_| -> http::Response { + panic!("blob request must not reach a server without blob support"); + }); + + assert_not_supported_error( + table.fetch_blobs("image", &[1]).await.unwrap_err(), + "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", + ); + } + + #[rstest] + #[case(ARROW_STREAM_CONTENT_TYPE)] + #[case("application/vnd.apache.arrow.stream; charset=utf-8")] + #[case("APPLICATION/VND.APACHE.ARROW.STREAM")] + #[tokio::test] + async fn test_fetch_blobs_accepts_stream_content_type_variants( + #[case] content_type: &'static str, + ) { + let table = table_with_fetch_blobs_content_type( + Some(content_type), + write_ipc_stream_uncompressed(&one_row_blob_batch("image")), + ); + + let blobs = table.fetch_blobs("image", &[10]).await.unwrap(); + + assert_eq!(blobs.value(0), b"alpha"); + } + + // Shared decoder accepts file framing because /query still returns it. + // fetch_blobs wire contract is stream. Enforced on the server. + #[rstest] + #[case(ARROW_FILE_CONTENT_TYPE)] + #[case("application/vnd.apache.arrow.file; charset=utf-8")] + #[case("APPLICATION/VND.APACHE.ARROW.FILE")] + #[tokio::test] + async fn test_fetch_blobs_accepts_file_content_type_variants( + #[case] content_type: &'static str, + ) { + let table = table_with_fetch_blobs_content_type( + Some(content_type), + write_ipc_file(&one_row_blob_batch("image")), + ); + + let blobs = table.fetch_blobs("image", &[10]).await.unwrap(); + + assert_eq!(blobs.value(0), b"alpha"); + } + + #[rstest] + #[case(ARROW_STREAM_CONTENT_TYPE, write_ipc_file(&one_row_blob_batch("image")))] + #[case( + ARROW_FILE_CONTENT_TYPE, + write_ipc_stream_uncompressed(&one_row_blob_batch("image")) + )] + #[tokio::test] + async fn test_fetch_blobs_fails_when_the_body_contradicts_the_content_type( + #[case] content_type: &'static str, + #[case] body: Vec, + ) { + let table = table_with_fetch_blobs_content_type(Some(content_type), body); + + let error = table.fetch_blobs("image", &[10]).await.unwrap_err(); + + assert!( + matches!(error, Error::Arrow { .. }), + "expected an Arrow decode failure, got {error}" + ); + } + + #[tokio::test] + async fn test_fetch_blobs_rejects_a_response_that_is_not_arrow_ipc() { + let table = table_with_fetch_blobs_content_type( + Some("application/json"), + br#"{"blobs": []}"#.to_vec(), + ); + + let error = table.fetch_blobs("image", &[10]).await.unwrap_err(); + assert_fetch_blobs_http_error(error, "got 'application/json'"); + } + + #[tokio::test] + async fn test_fetch_blobs_without_content_type_falls_back_to_file_framing() { + let table = + table_with_fetch_blobs_content_type(None, write_ipc_file(&one_row_blob_batch("image"))); + + let blobs = table.fetch_blobs("image", &[10]).await.unwrap(); + + assert_eq!(blobs.value(0), b"alpha"); + } + #[tokio::test] async fn test_query_vector_default_values() { let expected_data = RecordBatch::try_new( diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs new file mode 100644 index 000000000..550f72a4b --- /dev/null +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Cloud blob column listing and whole-byte fetch. + +use std::sync::Arc; + +use arrow_array::{Array, LargeBinaryArray}; +use arrow_schema::DataType; +use futures::TryStreamExt; + +use crate::Error; +use crate::blob::BlobFile; +use crate::error::Result; +use crate::remote::client::HttpSend; +use crate::table::BaseTable; + +use super::RemoteTable; + +impl RemoteTable { + /// Blob v2 columns are marked in field metadata, which `describe` returns. Reading + /// them from the cached schema needs no route of its own and no version gate. + pub(super) async fn blob_columns_impl(&self) -> Result> { + let schema = self.schema().await?; + Ok(crate::blob::blob_column_names(schema.as_ref())) + } + + pub(super) async fn fetch_blobs_impl( + &self, + 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. + if row_ids.is_empty() { + return Ok(LargeBinaryArray::from(Vec::>::new())); + } + if !self.server_version.support_blobs() { + return Err(Error::NotSupported { + message: "fetch_blobs is not supported on this LanceDB Cloud server".into(), + }); + } + let version = self.current_version().await; + let mut body = serde_json::json!({ + "version": version, + "column": column, + "row_ids": row_ids, + }); + self.apply_branch_body(&mut body); + + let request = self + .post_read(&format!("/v1/table/{}/fetch_blobs/", self.identifier)) + .json(&body); + let (request_id, response) = self.send(request, true).await?; + let mut stream = self.read_arrow_response(&request_id, response).await?; + + let mut blob_chunks: Vec> = Vec::new(); + while let Some(batch) = stream.try_next().await? { + let blob_column = batch.column_by_name(column).ok_or_else(|| Error::Http { + source: format!("fetch_blobs response is missing the '{column}' column").into(), + request_id: request_id.clone(), + status_code: None, + })?; + // The server returns LargeBinary today. Accept the other binary types so a + // server that switches encodings does not break older clients. + if !matches!( + blob_column.data_type(), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView + ) { + return Err(Error::Http { + source: format!( + "fetch_blobs response column has type {}, expected Binary, LargeBinary, or BinaryView", + blob_column.data_type() + ) + .into(), + request_id: request_id.clone(), + status_code: None, + }); + } + blob_chunks.push(arrow::compute::cast(blob_column, &DataType::LargeBinary)?); + } + let blobs = if blob_chunks.is_empty() { + LargeBinaryArray::from(Vec::>::new()) + } else { + let blob_chunk_refs: Vec<&dyn Array> = blob_chunks.iter().map(AsRef::as_ref).collect(); + arrow::compute::concat(&blob_chunk_refs)? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::Http { + source: "fetch_blobs could not read the concatenated response as LargeBinary" + .into(), + request_id: request_id.clone(), + status_code: None, + })? + .clone() + }; + if blobs.len() != row_ids.len() { + return Err(Error::Http { + source: format!( + "fetch_blobs returned {} rows for {} row ids", + blobs.len(), + row_ids.len() + ) + .into(), + request_id, + status_code: None, + }); + } + Ok(blobs) + } + + pub(super) async fn fetch_blob_files_impl( + &self, + _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(), + }) + } +}