diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 35f50c5a8..0e7164c1a 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -793,6 +793,15 @@ with `search().with_row_id(True).to_arrow()` and pass the result to the fetch AP One missing or deleted row ID rejects the entire blob-read batch; the APIs do not return partial results. +Local and remote `BlobFile` handles prefetch 512 KiB for sequential reads by +default. Pass `buffer_size=0` to `fetch_blob_files`, `open_blob_file`, or +`open_blob_files` for sparse reads, or choose another window size in bytes. +Call `file.set_buffer_size(size)` to change it without moving the cursor; +changing the size discards unused prefetched bytes. Explicit `read_range` and +`read_ranges` calls always read exactly the requested intervals and leave the +cursor unchanged. Remote prefetch consumes the HTTP response completely before +returning, so idle handles do not hold request concurrency slots or connections. + To construct inline Blob V2 arrays for ingestion in record batches, use [blob_array][lancedb.schema.blob_array]. It accepts bytes and `None`, preserving the distinction between an empty blob (`b""`) and a null blob: diff --git a/python/benchmarks/bench_fetch_blobs_rss.py b/python/benchmarks/bench_fetch_blobs_rss.py new file mode 100644 index 000000000..8281490f0 --- /dev/null +++ b/python/benchmarks/bench_fetch_blobs_rss.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Measure fetch_blobs RSS in a fresh client process against real HTTP. + +Run with the uv environment's Python. The server lives in a separate process so +payload generation and Arrow serialization are excluded from the client's RSS. +Each row contains 1 MiB; requests over the server's 64 MiB limit are split by the +SDK. Example: python benchmarks/bench_fetch_blobs_rss.py --rows 48 144 +""" + +import argparse +import json +import resource +import subprocess +import sys +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + + +def serve(batch_rows, row_mib): + import pyarrow as pa + + payload = b"x" * (row_mib << 20) + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *args): + pass + + def respond(self, status, data, content_type="application/json"): + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(data))) + self.send_header("phalanx-version", "0.5.0") + self.send_header("x-lancedb-version", "1") + self.end_headers() + for offset in range(0, len(data), 8192): + self.wfile.write(data[offset : offset + 8192]) + + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + request_body = self.rfile.read(length) + if self.path.endswith("/describe/"): + self.do_GET() + return + body = json.loads(request_body) + count = len(body["row_ids"]) + if count * row_mib > 64: + self.respond( + 400, + b"fetch_blobs accepts at most 67108864 total blob bytes", + ) + return + schema = pa.schema([pa.field("image", pa.large_binary())]) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, schema) as writer: + for start in range(0, count, batch_rows): + array = pa.array( + [payload] * min(batch_rows, count - start), + type=pa.large_binary(), + ) + writer.write_batch(pa.record_batch([array], schema=schema)) + self.respond(200, sink.getvalue(), "application/vnd.apache.arrow.stream") + + def do_GET(self): + if "/blob/image/" in self.path: + start, end = map(int, self.headers["Range"][6:].split("-")) + body = memoryview(payload)[start : end + 1] + self.send_response(206) + self.send_header("Content-Length", str(len(body))) + self.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}") + self.send_header("x-lancedb-version", "1") + self.end_headers() + for offset in range(0, len(body), 8192): + self.wfile.write(body[offset : offset + 8192]) + return + self.respond( + 200, + json.dumps( + { + "table": "test", + "version": 1, + "schema": {"fields": []}, + } + ).encode(), + ) + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + print(server.server_port, flush=True) + server.serve_forever() + + +def client(port, rows, row_mib): + import lancedb + + db = lancedb.connect( + "db://rss-test", api_key="test", host_override=f"http://127.0.0.1:{port}" + ) + table = db.open_table("test") + # Reset the Linux high-water mark after imports and connection setup. + with open("/proc/self/clear_refs", "w") as reset: + reset.write("5") + + def memory_kib(field): + with open("/proc/self/status") as status: + return next( + int(line.split()[1]) for line in status if line.startswith(field + ":") + ) + + before = memory_kib("VmRSS") + values = table.fetch_blobs("image", list(range(rows))) + peak = memory_kib("VmHWM") + assert len(values) == rows + assert values.buffers()[2].size == rows * (row_mib << 20) + # Inspect slices without copying the result to Python bytes. + for i in (0, rows - 1): + assert values.slice(i, 1).buffers()[2][i * (row_mib << 20)] == ord("x") + print( + json.dumps( + { + "rows": rows, + "payload_mib": rows * row_mib, + "baseline_rss_kib": before, + "peak_rss_kib": peak, + "peak_delta_kib": peak - before, + "peak_delta_per_payload": (peak - before) / (rows * row_mib * 1024), + "rusage_maxrss_kib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, + } + ), + flush=True, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", nargs="+", type=int, default=[48, 144]) + parser.add_argument("--batch-rows", type=int, default=48) + parser.add_argument("--row-mib", type=int, default=1) + parser.add_argument("--serve", action="store_true") + parser.add_argument("--client", type=int) + args = parser.parse_args() + if args.serve: + serve(args.batch_rows, args.row_mib) + elif args.client is not None: + client(args.client, args.rows[0], args.row_mib) + else: + server = subprocess.Popen( + [ + sys.executable, + __file__, + "--serve", + "--batch-rows", + str(args.batch_rows), + "--row-mib", + str(args.row_mib), + ], + stdout=subprocess.PIPE, + text=True, + ) + try: + port = server.stdout.readline().strip() + for rows in args.rows: + subprocess.run( + [ + sys.executable, + __file__, + "--client", + port, + "--rows", + str(rows), + "--row-mib", + str(args.row_mib), + ], + check=True, + ) + finally: + server.terminate() + server.wait() + + +if __name__ == "__main__": + main() diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index 2c545a17e..fc1925ca6 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -35,6 +35,15 @@ FetchBlobsSync = Callable[[str, list[int]], pa.Array | pa.ChunkedArray] FetchBlobsAsync = Callable[[str, list[int]], Awaitable[pa.Array | pa.ChunkedArray]] +def _validate_blob_buffer_size(size: Optional[int]) -> Optional[int]: + if size is None: + return None + size = operator.index(size) + if size < 0: + raise ValueError("buffer_size must be non-negative") + return size + + def _normalize_blob_descriptor(value): if isinstance(value, pa.Scalar): value = value.as_py() @@ -58,6 +67,16 @@ class BlobFile(io.RawIOBase): def __init__(self, inner) -> None: self._inner = inner + def set_buffer_size(self, size: int) -> None: + """Set sequential read-ahead bytes (default 512 KiB); 0 disables it. + + Changing the size discards unused prefetched bytes without moving the + cursor. Explicit ``read_range`` and ``read_ranges`` bypass read-ahead. + """ + if self.closed: + raise ValueError("set_buffer_size of closed file") + self._inner.set_buffer_size(_validate_blob_buffer_size(size)) + async def aread(self) -> bytes: if self.closed: raise ValueError("read of closed file") diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index a64014cb7..4f1da9ba7 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -397,6 +397,7 @@ class BlobFile: def seek(self, position: int) -> None: ... def tell(self) -> int: ... def size(self) -> int: ... + def set_buffer_size(self, size: int) -> None: ... def read_range(self, offset: int, length: int) -> bytes: ... def read_ranges(self, ranges: list[tuple[int, int]]) -> list[bytes]: ... def read_up_to(self, length: int) -> bytes: ... @@ -755,9 +756,11 @@ class Table: requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: ... async def fetch_blob_files( - self, column: str, row_ids: list[int] + self, column: str, row_ids: list[int], buffer_size: Optional[int] = None ) -> list[Optional[BlobFile]]: ... - async def open_blob_file(self, descriptor: str) -> BlobFile: ... + async def open_blob_file( + self, descriptor: str, buffer_size: Optional[int] = None + ) -> BlobFile: ... def vector_search(self) -> VectorQuery: ... class Tags: diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 3f6bff45e..21749a0a7 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1420,13 +1420,20 @@ class RemoteTable(Table): self, column: str, row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], + *, + buffer_size: Optional[int] = None, ) -> "list[Optional[BlobFile]]": - return LOOP.run(self._table.fetch_blob_files(column, row_ids)) + return LOOP.run( + self._table.fetch_blob_files(column, row_ids, buffer_size=buffer_size) + ) def open_blob_file( - self, descriptor: Union[dict, pa.StructScalar, None] + self, + descriptor: Union[dict, pa.StructScalar, None], + *, + buffer_size: Optional[int] = None, ) -> Optional[BlobFile]: - return LOOP.run(self._table.open_blob_file(descriptor)) + return LOOP.run(self._table.open_blob_file(descriptor, buffer_size=buffer_size)) def head(self, n=5) -> pa.Table: """ diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 762f826c2..424b4b64c 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -40,6 +40,7 @@ from ._blob import ( BlobFile, _normalize_blob_descriptor, _normalize_blob_row_ids, + _validate_blob_buffer_size, _wrap_blob_files, strip_auto_row_ids, validate_blob_mode, @@ -2664,6 +2665,8 @@ class Table(ABC): self, column: str, row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], + *, + buffer_size: Optional[int] = None, ) -> "list[Optional[BlobFile]]": """Open lazy, seekable :class:`~lancedb._blob.BlobFile` handles. @@ -2673,6 +2676,8 @@ class Table(ABC): ``pyarrow.Table`` carrying row identity via ``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null blob values return ``None``. LanceDB Enterprise requires server version 0.5.0 or newer. + ``buffer_size`` controls sequential read-ahead in bytes (default 512 KiB). + Set it to 0 for sparse reads. Explicit range reads bypass this buffer. ``_rowid`` values stay valid after compaction when the table has stable row ids. @@ -2682,7 +2687,10 @@ class Table(ABC): @abstractmethod def open_blob_file( - self, descriptor: Union[dict, pa.StructScalar, None] + self, + descriptor: Union[dict, pa.StructScalar, None], + *, + buffer_size: Optional[int] = None, ) -> Optional[BlobFile]: """Open a seekable file from a SQL ``blob_descriptor`` result. @@ -2691,6 +2699,7 @@ class Table(ABC): not change this table's checkout. Null descriptors return ``None``. Payload bytes load only when the returned handle is read. An external blob with an unspecified size may require a metadata request on opening. + ``buffer_size`` controls sequential read-ahead (default 512 KiB; 0 disables). ```python rows = db.execute_query( @@ -2701,7 +2710,9 @@ class Table(ABC): ``` """ - def open_blob_files(self, descriptors: Iterable) -> Iterator[Optional[BlobFile]]: + def open_blob_files( + self, descriptors: Iterable, *, buffer_size: Optional[int] = None + ) -> Iterator[Optional[BlobFile]]: """Iterate over lazy files from descriptor dictionaries or Arrow scalars. Accepts an Arrow array/chunked array or any iterable. Opens each handle @@ -2716,7 +2727,10 @@ class Table(ABC): ``` """ for descriptor in descriptors: - yield self.open_blob_file(descriptor) + if buffer_size is None: + yield self.open_blob_file(descriptor) + else: + yield self.open_blob_file(descriptor, buffer_size=buffer_size) @abstractmethod def _execute_query( @@ -3998,13 +4012,20 @@ class LanceTable(Table): self, column: str, row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], + *, + buffer_size: Optional[int] = None, ) -> "list[Optional[BlobFile]]": - return LOOP.run(self._table.fetch_blob_files(column, row_ids)) + return LOOP.run( + self._table.fetch_blob_files(column, row_ids, buffer_size=buffer_size) + ) def open_blob_file( - self, descriptor: Union[dict, pa.StructScalar, None] + self, + descriptor: Union[dict, pa.StructScalar, None], + *, + buffer_size: Optional[int] = None, ) -> Optional[BlobFile]: - return LOOP.run(self._table.open_blob_file(descriptor)) + return LOOP.run(self._table.open_blob_file(descriptor, buffer_size=buffer_size)) @property def tags(self) -> Tags: @@ -8620,26 +8641,39 @@ class AsyncTable: self, column: str, row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], + *, + buffer_size: Optional[int] = None, ) -> "list[Optional[BlobFile]]": + """Open handles with configurable read-ahead. + + See [Table.fetch_blob_files][lancedb.table.Table.fetch_blob_files]. + """ + buffer_size = _validate_blob_buffer_size(buffer_size) handles = await self._inner.fetch_blob_files( - column, _normalize_blob_row_ids(row_ids, column) + column, _normalize_blob_row_ids(row_ids, column), buffer_size ) return _wrap_blob_files(handles) async def open_blob_file( - self, descriptor: Union[dict, pa.StructScalar, None] + self, + descriptor: Union[dict, pa.StructScalar, None], + *, + buffer_size: Optional[int] = None, ) -> Optional[BlobFile]: """Open a snapshot-bound descriptor as a seekable file. See [Table.open_blob_file][lancedb.table.Table.open_blob_file]. """ + buffer_size = _validate_blob_buffer_size(buffer_size) descriptor = _normalize_blob_descriptor(descriptor) if descriptor is None: return None - return BlobFile(await self._inner.open_blob_file(json.dumps(descriptor))) + return BlobFile( + await self._inner.open_blob_file(json.dumps(descriptor), buffer_size) + ) async def open_blob_files( - self, descriptors: Iterable + self, descriptors: Iterable, *, buffer_size: Optional[int] = None ) -> AsyncIterator[Optional[BlobFile]]: """Asynchronously iterate over files in descriptor order, preserving nulls. @@ -8654,7 +8688,7 @@ class AsyncTable: ``` """ for descriptor in descriptors: - yield await self.open_blob_file(descriptor) + yield await self.open_blob_file(descriptor, buffer_size=buffer_size) @property def tags(self) -> AsyncTags: diff --git a/python/python/tests/test_blob_buffering.py b/python/python/tests/test_blob_buffering.py new file mode 100644 index 000000000..09fbfb862 --- /dev/null +++ b/python/python/tests/test_blob_buffering.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import json +import threading +from contextlib import contextmanager +from urllib.parse import urlparse + +import pyarrow as pa +import pytest + +import lancedb +from test_remote_db import ( + BLOB_DESCRIBE_RESPONSE, + blob_descriptor_take_table, + mock_lancedb_connection, +) + + +@contextmanager +def remote_buffer_table(*, http_requests=None, truncate_first=False): + payload = bytes(range(256)) * 4096 + requests = [] + + def handler(request): + path = urlparse(request.path).path + if http_requests is not None: + http_requests.append(path) + if path.endswith("/describe/"): + body = json.dumps(BLOB_DESCRIBE_RESPONSE).encode() + content_type = "application/json" + elif path.endswith("/query/"): + data = json.loads( + request.rfile.read(int(request.headers["Content-Length"])) + ) + assert data["columns"] == ["image"] + table = blob_descriptor_take_table().slice(0, 1) + images = table["image"].combine_chunks() + images = pa.StructArray.from_arrays( + [ + images.field("kind"), + images.field("position"), + pa.array([len(payload)], pa.uint64()), + images.field("blob_id"), + images.field("blob_uri"), + ], + fields=list(images.type), + ) + table = table.set_column(0, table.schema.field(0), images) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + body = sink.getvalue().to_pybytes() + content_type = "application/vnd.apache.arrow.stream" + else: + assert path == "/v1/table/test/blob/image/10/bytes" + assert "version=1" in request.path + start, end = map(int, request.headers["Range"][6:].split("-")) + requests.append((start, end + 1)) + body = payload[start : end + 1] + if truncate_first and len(requests) == 1: + body = body[:5] + request.send_response(206) + request.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}") + request.send_header("Content-Length", str(len(body))) + request.end_headers() + request.wfile.write(body) + return + request.send_response(200) + request.send_header("Content-Type", content_type) + request.send_header("Content-Length", str(len(body))) + request.send_header("phalanx-version", "0.5.0") + request.send_header("x-lancedb-version", "1") + request.end_headers() + request.wfile.write(body) + + with mock_lancedb_connection(handler, blob_request_concurrency=1) as db: + yield db.open_table("test"), payload, requests + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("buffer_size", [None, 0, 8, 512 * 1024]) +async def test_remote_read_ahead_over_http(use_async, buffer_size): + with remote_buffer_table() as (table, payload, requests): + if use_async: + files = await table._table.fetch_blob_files( + "image", [10], buffer_size=buffer_size + ) + else: + files = table.fetch_blob_files("image", [10], buffer_size=buffer_size) + file = files[0] + assert requests == [] + assert file.read(2) == payload[:2] + window = 512 * 1024 if buffer_size is None else max(buffer_size, 2) + assert requests == [(0, window)] + assert file.read(2) == payload[2:4] + assert len(requests) == (2 if buffer_size == 0 else 1) + file.seek(1) + assert file.read(1) == payload[1:2] + assert len(requests) == (3 if buffer_size == 0 else 1) + before = len(requests) + assert file.read_range(0, 1) == payload[:1] + assert len(requests) == before + 1 + assert requests[-1] == (0, 1) + assert file.tell() == 2 + file.set_buffer_size(0) + assert file.read(1) == payload[2:3] + assert requests[-1] == (2, 3) + file.set_buffer_size(16) + assert file.read(1) == payload[3:4] + assert requests[-1] == (3, 19) + file.seek(len(payload) - 1) + assert file.read(4) == payload[-1:] + assert requests[-1] == (len(payload) - 1, len(payload)) + before = len(requests) + assert file.read(1) == b"" + assert len(requests) == before + file.close() + with pytest.raises(ValueError, match="closed"): + file.set_buffer_size(0) + + +def test_short_prefetch_over_http_preserves_cursor(): + with remote_buffer_table(truncate_first=True) as (table, payload, requests): + file = table.fetch_blob_files("image", [10], buffer_size=8)[0] + with pytest.raises(RuntimeError, match="returned 5 bytes, expected 8"): + file.read(6) + assert file.tell() == 0 + assert file.read(3) == payload[:3] + assert file.tell() == 3 + assert requests == [(0, 8), (0, 8)] + file.close() + + +def test_idle_prefetch_releases_http_concurrency_slot(): + with remote_buffer_table() as (table, payload, requests): + first, second = table.fetch_blob_files("image", [10, 10]) + assert first.read(1) == payload[:1] + results = [] + errors = [] + + def read_other(): + try: + results.append(second.read(1)) + except Exception as error: + errors.append(error) + + reader = threading.Thread(target=read_other, daemon=True) + reader.start() + reader.join(timeout=2) + completed_while_first_idle = not reader.is_alive() + first.close() + reader.join(timeout=2) + second.close() + assert completed_while_first_idle, ( + "idle prefetch retained the HTTP request slot" + ) + assert not errors + assert results == [payload[:1]] + assert requests == [(0, 512 * 1024), (0, 512 * 1024)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_local_buffer_configuration_through_python(tmp_path, use_async): + db = lancedb.connect(tmp_path) + table = db.create_table( + "blobs", [{"image": b"0123456789"}], schema=pa.schema([lancedb.blob("image")]) + ) + row_ids = table.search().with_row_id(True).to_arrow()["_rowid"] + if use_async: + files = await table._table.fetch_blob_files("image", row_ids, buffer_size=0) + else: + files = table.fetch_blob_files("image", row_ids, buffer_size=0) + file = files[0] + assert file.read(2) == b"01" + file.set_buffer_size(4) + assert file.read(2) == b"23" + file.seek(1) + assert file.read(3) == b"123" + assert file.read_range(8, 2) == b"89" + with pytest.raises(ValueError, match="non-negative"): + file.set_buffer_size(-1) + with pytest.raises(TypeError): + file.set_buffer_size(1.5) + file.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_open_descriptor_files_forward_buffer_size(use_async): + with remote_buffer_table() as (table, payload, requests): + descriptor = { + "kind": 1, + "position": 0, + "size": len(payload), + "blob_id": 0, + "blob_uri": "", + "source": { + "uri": "s3://bucket/test.lance", + "version": 1, + "committed_at_micros": 1, + "column": "image", + "row_reference": 10, + "row_is_address": False, + }, + } + if use_async: + files = [ + file + async for file in table._table.open_blob_files( + [None, descriptor], buffer_size=8 + ) + ] + else: + files = list(table.open_blob_files([None, descriptor], buffer_size=8)) + assert files[0] is None + assert files[1].read(1) == payload[:1] + assert requests == [(0, 8)] + files[1].close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_invalid_buffer_size_rejects_before_http(use_async): + http_requests = [] + with remote_buffer_table(http_requests=http_requests) as (table, _, requests): + http_requests.clear() + with pytest.raises(ValueError, match="non-negative"): + if use_async: + await table._table.fetch_blob_files("image", [10], buffer_size=-1) + else: + table.fetch_blob_files("image", [10], buffer_size=-1) + assert requests == [] + assert http_requests == [] diff --git a/python/src/table.rs b/python/src/table.rs index c4211e44d..62c2ee765 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1243,6 +1243,15 @@ impl PyBlobFile { .map_err(|e| blob_file_error("tell", e)) } + fn set_buffer_size(self_: PyRef<'_, Self>, size: usize) -> PyResult<()> { + self_.ensure_open()?; + let inner = self_.inner.clone(); + self_ + .py() + .detach(move || block_on(async move { inner.set_buffer_size(size).await })) + .map_err(|e| blob_file_error("set_buffer_size", e)) + } + fn size(self_: PyRef<'_, Self>) -> u64 { self_.inner.size() } @@ -2019,11 +2028,12 @@ impl Table { } /// Open lazy blob handles for `row_ids` from blob v2 column `column`. - #[pyo3(signature = (column, row_ids))] + #[pyo3(signature = (column, row_ids, buffer_size=None))] pub fn fetch_blob_files( self_: PyRef<'_, Self>, column: String, row_ids: Vec, + buffer_size: Option, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { @@ -2031,6 +2041,13 @@ impl Table { .fetch_blob_files(column.as_str(), &row_ids) .await .infer_error()?; + if let Some(size) = buffer_size { + for file in handles.iter().flatten() { + file.set_buffer_size(size) + .await + .map_err(|e| blob_file_error("set_buffer_size", e))?; + } + } Ok(handles .into_iter() .map(|handle| { @@ -2043,16 +2060,22 @@ impl Table { }) } - #[pyo3(signature = (descriptor))] + #[pyo3(signature = (descriptor, buffer_size=None))] pub fn open_blob_file<'py>( self_: PyRef<'py, Self>, descriptor: &str, + buffer_size: Option, ) -> PyResult> { let descriptor: BlobDescriptor = serde_json::from_str(descriptor) .map_err(|error| PyValueError::new_err(format!("invalid blob descriptor: {error}")))?; let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { let file = inner.open_blob_file(&descriptor).await.infer_error()?; + if let Some(size) = buffer_size { + file.set_buffer_size(size) + .await + .map_err(|e| blob_file_error("set_buffer_size", e))?; + } Ok(PyBlobFile { inner: Arc::new(file), closed: AtomicBool::new(false), diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index a8a816fdd..ead1f2410 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -204,6 +204,30 @@ impl BlobFile { lance::dataset::BlobFile::new_external(object_store, path, uri, position, size).into() } + /// Default sequential read-ahead window, in bytes. + pub const DEFAULT_READ_BUFFER_SIZE: usize = 512 * 1024; + + /// Set sequential read-ahead bytes. `0` disables read-ahead. + /// + /// Changing the size discards unused prefetched bytes without moving the + /// cursor. [`Self::read_range`] and [`Self::read_ranges`] bypass read-ahead. + /// The default is [`Self::DEFAULT_READ_BUFFER_SIZE`]. + /// + /// ``` + /// # async fn sparse_read(file: &lancedb::blob::BlobFile) -> lance_core::Result<()> { + /// file.set_buffer_size(0).await?; + /// let header = file.read_up_to(16).await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn set_buffer_size(&self, size: usize) -> lance_core::Result<()> { + match &self.inner { + BlobFileInner::Native(file) => file.set_buffer_size(size).await, + #[cfg(feature = "remote")] + BlobFileInner::Remote(file) => file.set_buffer_size(size).await, + } + } + /// Close the handle. pub async fn close(&self) -> lance_core::Result<()> { match &self.inner { diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index 65ab13935..345b731c5 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -3,6 +3,11 @@ //! Cloud blob column listing, whole-byte fetch, and seekable HTTP byte-range handles. +mod aggregate; + +use aggregate::BlobAccumulator; +use arrow_buffer::{Buffer, OffsetBuffer}; + use std::collections::HashMap; use std::ops::Range; use std::sync::Arc; @@ -159,9 +164,25 @@ impl BlobRangeRequester for TableBlobRangeRequester { } } -#[derive(Debug, Default)] +#[derive(Debug)] struct RemoteBlobState { cursor: u64, + buffer_size: usize, +} + +impl Default for RemoteBlobState { + fn default() -> Self { + Self { + cursor: 0, + buffer_size: BlobFile::DEFAULT_READ_BUFFER_SIZE, + } + } +} + +#[derive(Debug, Default)] +struct RemoteBlobBuffer { + start: u64, + bytes: Bytes, } /// Seekable Cloud blob handle over HTTP Range. @@ -171,6 +192,9 @@ struct RemoteBlobState { pub struct RemoteBlobFile { requester: Arc, state: Mutex, + // Never held across I/O: close can release cached bytes even while a read + // holds the cursor lock. A reader may retain its own prefix until it ends. + buffer: std::sync::Mutex, closed: AtomicBool, size: u64, version: Option, @@ -181,6 +205,7 @@ impl RemoteBlobFile { Self { requester, state: Mutex::new(RemoteBlobState::default()), + buffer: std::sync::Mutex::new(RemoteBlobBuffer::default()), closed: AtomicBool::new(false), size, version, @@ -190,6 +215,18 @@ impl RemoteBlobFile { /// 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); + self.buffer.lock().unwrap().bytes = Bytes::new(); + Ok(()) + } + + pub(crate) async fn set_buffer_size(&self, size: usize) -> lance_core::Result<()> { + self.ensure_open()?; + let mut state = self.state.lock().await; + self.ensure_open()?; + if state.buffer_size != size { + state.buffer_size = size; + self.buffer.lock().unwrap().bytes = Bytes::new(); + } Ok(()) } @@ -228,7 +265,7 @@ impl RemoteBlobFile { let BlobRangeResponse { request_id, response, - _permit, + _permit: permit, } = self .requester .request_range(&range_header, RangeRequestMode::DataRead, self.version) @@ -236,19 +273,9 @@ impl RemoteBlobFile { .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)?; + let bytes = read_blob_body(response, range.end - range.start, request_id).await?; + drop(permit); 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) } @@ -275,7 +302,16 @@ impl RemoteBlobFile { 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 + // Joining a small cached prefix to a large tail would copy the entire + // blob while both allocations are live. Match the native reader's + // bounded reuse rule for whole-blob reads. + let held_len = cached_prefix(&self.buffer.lock().unwrap(), state.cursor, remaining).len(); + if remaining > held_len.saturating_mul(2) { + self.buffer.lock().unwrap().bytes = Bytes::new(); + } + let bytes = self.read_up_to_locked(&mut state, remaining).await?; + self.buffer.lock().unwrap().bytes = Bytes::new(); + Ok(bytes) } pub(crate) async fn read_up_to(&self, len: usize) -> lance_core::Result { @@ -296,14 +332,51 @@ impl RemoteBlobFile { return Ok(Bytes::new()); } - // A partial read must release its HTTP response before returning. Keeping - // it open across calls would hold a connection slot indefinitely and - // prevent other handles from making progress. - let end = state.cursor + target_len as u64; - let bytes = self.read_range(state.cursor..end).await?; + let prefix = cached_prefix(&self.buffer.lock().unwrap(), state.cursor, target_len); + if prefix.len() == target_len { + state.cursor += target_len as u64; + return Ok(prefix); + } + let fetch_start = state.cursor + prefix.len() as u64; + let still_need = target_len - prefix.len(); + let fetch_len = if state.buffer_size == 0 { + still_need as u64 + } else { + (state.buffer_size.max(still_need) as u64).min(self.size - fetch_start) + }; + // Fully consume the response before retaining only its Bytes. read_range + // releases both the connection and the concurrency permit before return. + let data = self + .read_range(fetch_start..fetch_start + fetch_len) + .await?; self.ensure_open()?; - state.cursor = end; - Ok(bytes) + let tail = data.slice(..still_need); + let result = if prefix.is_empty() { + tail + } else { + let mut joined = Vec::new(); + joined.try_reserve_exact(target_len).map_err(|error| { + remote_blob_error(format!( + "cannot allocate blob read of {target_len} bytes: {error}" + )) + })?; + joined.extend_from_slice(&prefix); + joined.extend_from_slice(&tail); + Bytes::from(joined) + }; + { + let mut buffer = self.buffer.lock().unwrap(); + // Serialize publishing a new window with close clearing the cache. + self.ensure_open()?; + if state.buffer_size != 0 { + buffer.start = fetch_start; + buffer.bytes = data; + } else { + buffer.bytes = Bytes::new(); + } + } + state.cursor += target_len as u64; + Ok(result) } pub(crate) async fn seek(&self, new_cursor: u64) -> lance_core::Result<()> { @@ -326,6 +399,83 @@ impl RemoteBlobFile { } } +/// Collect a range without retaining a second payload's worth of network chunks. +async fn read_blob_body( + response: Response, + expected: u64, + request_id: String, +) -> lance_core::Result { + let expected = usize::try_from(expected).map_err(|_| { + lance_core::Error::invalid_input("blob range length exceeds addressable memory") + })?; + let length_error = |received| { + remote_blob_error(format!( + "byte range returned {received} bytes, expected {expected}" + )) + }; + let mut chunks = response.bytes_stream(); + let first = chunks + .try_next() + .await + .err_to_http(request_id.clone()) + .map_err(remote_blob_error)? + .ok_or_else(|| length_error(0))?; + if first.len() > expected { + return Err(length_error(first.len())); + } + if first.len() == expected { + while let Some(chunk) = chunks + .try_next() + .await + .err_to_http(request_id.clone()) + .map_err(remote_blob_error)? + { + if !chunk.is_empty() { + return Err(length_error(expected.saturating_add(chunk.len()))); + } + } + return Ok(first); + } + let mut bytes = Vec::new(); + bytes.try_reserve_exact(expected).map_err(|error| { + remote_blob_error(format!( + "cannot allocate blob range of {expected} bytes: {error}" + )) + })?; + bytes.extend_from_slice(&first); + drop(first); + while let Some(chunk) = chunks + .try_next() + .await + .err_to_http(request_id.clone()) + .map_err(remote_blob_error)? + { + let received = bytes.len().saturating_add(chunk.len()); + if received > expected { + return Err(length_error(received)); + } + bytes.extend_from_slice(&chunk); + } + if bytes.len() != expected { + return Err(length_error(bytes.len())); + } + Ok(Bytes::from(bytes)) +} + +fn cached_prefix(buffer: &RemoteBlobBuffer, cursor: u64, length: usize) -> Bytes { + let Some(offset) = cursor.checked_sub(buffer.start) else { + return Bytes::new(); + }; + let Ok(offset) = usize::try_from(offset) else { + return Bytes::new(); + }; + if offset >= buffer.bytes.len() { + return Bytes::new(); + } + let end = offset.saturating_add(length).min(buffer.bytes.len()); + buffer.bytes.slice(offset..end) +} + fn remote_blob_error(error: impl std::fmt::Display) -> lance_core::Error { lance_core::Error::io(format!("remote blob read failed: {error}")) } @@ -434,10 +584,13 @@ impl RemoteTable { read_snapshot.version = Some(self.describe_read_snapshot(read_snapshot).await?.version); } let mut pending: Vec<&[u64]> = row_ids.chunks(MAX_FETCH_BLOBS_ROW_IDS).rev().collect(); - let mut chunks = Vec::new(); + let mut blobs = BlobAccumulator::new(row_ids.len()); while let Some(ids) = pending.pop() { - match self.fetch_blobs_chunk(column, ids, read_snapshot).await { - Ok(blobs) => chunks.push(blobs), + match self + .fetch_blobs_chunk(column, ids, read_snapshot, &mut blobs) + .await + { + Ok(()) => {} Err(error) if is_fetch_blobs_byte_limit_error(&error) => { // A single-request call may turn into several requests after a // byte-cap error. No bytes from the failed request were used. @@ -455,7 +608,15 @@ impl RemoteTable { Some(file) => Some(file.read().await?), None => None, }; - chunks.push(LargeBinaryArray::from(vec![blob.as_deref()])); + let blob = match blob { + Some(bytes) => LargeBinaryArray::new( + OffsetBuffer::new(vec![0, bytes.len() as i64].into()), + Buffer::from(bytes), + None, + ), + None => LargeBinaryArray::new_null(1), + }; + blobs.append(&blob)?; } else { let mid = ids.len() / 2; pending.push(&ids[mid..]); @@ -466,15 +627,7 @@ impl RemoteTable { } } - if chunks.len() == 1 { - return Ok(chunks.pop().unwrap()); - } - let chunk_refs: Vec<&dyn Array> = chunks.iter().map(|chunk| chunk as &dyn Array).collect(); - Ok(arrow::compute::concat(&chunk_refs)? - .as_any() - .downcast_ref::() - .expect("concatenating LargeBinary arrays returns LargeBinary") - .clone()) + Ok(blobs.finish()) } async fn fetch_blobs_chunk( @@ -482,7 +635,8 @@ impl RemoteTable { column: &str, row_ids: &[u64], read_snapshot: ReadSnapshot, - ) -> Result { + blobs: &mut BlobAccumulator, + ) -> Result<()> { let mut body = serde_json::json!({ "version": read_snapshot.version, "column": column, @@ -502,7 +656,7 @@ impl RemoteTable { Err(error) => return Err(self.explain_list_blob_bad_request(column, error).await), }; - let mut blob_chunks: Vec> = Vec::new(); + let mut rows = 0; 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(), @@ -525,28 +679,19 @@ impl RemoteTable { 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)? + let array = arrow::compute::cast(blob_column, &DataType::LargeBinary)?; + let array = array .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() { + .expect("casting to LargeBinary returns LargeBinary"); + rows += array.len(); + blobs.append(array)?; + } + if rows != row_ids.len() { return Err(Error::Http { source: format!( "fetch_blobs returned {} rows for {} row ids", - blobs.len(), + rows, row_ids.len() ) .into(), @@ -554,7 +699,7 @@ impl RemoteTable { status_code: None, }); } - Ok(blobs) + Ok(()) } /// Materialize ranges through version-pinned handles, preserving nulls and order. @@ -1893,9 +2038,8 @@ mod tests { [ ("bytes=0-0".to_string(), None), ("bytes=5-8".to_string(), Some("5".to_string())), - ("bytes=0-3".to_string(), Some("5".to_string())), ( - format!("bytes=10-{}", PAYLOAD.len() - 1), + format!("bytes=0-{}", PAYLOAD.len() - 1), Some("5".to_string()) ), ] @@ -1943,7 +2087,7 @@ mod tests { } #[tokio::test] - async fn remote_blob_file_releases_response_after_each_partial_read() { + async fn remote_blob_file_reuses_prefetched_bytes_across_reads_and_seeks() { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); let file = table @@ -1961,7 +2105,41 @@ mod tests { assert_eq!(file.tell().await.unwrap(), 24); assert_eq!( requests.lock().unwrap().as_slice(), - ["bytes=0-0", "bytes=0-3", "bytes=4-6", "bytes=20-23"] + ["bytes=0-0", &format!("bytes=0-{}", PAYLOAD.len() - 1)] + ); + } + + #[tokio::test] + async fn remote_buffer_boundaries_and_disabled_read_ahead() { + let requests = Arc::new(StdMutex::new(Vec::new())); + let table = mock_remote_blob_table(requests.clone()); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + file.set_buffer_size(4).await.unwrap(); + assert_eq!(file.read_up_to(3).await.unwrap(), &PAYLOAD[..3]); + assert_eq!(file.read_up_to(3).await.unwrap(), &PAYLOAD[3..6]); + file.seek(5).await.unwrap(); + assert_eq!(file.read_up_to(2).await.unwrap(), &PAYLOAD[5..7]); + assert_eq!(file.read_range(5..6).await.unwrap(), &PAYLOAD[5..6]); + assert_eq!(file.tell().await.unwrap(), 7); + file.set_buffer_size(0).await.unwrap(); + assert_eq!(file.read_up_to(1).await.unwrap(), &PAYLOAD[7..8]); + assert_eq!(file.read_up_to(1).await.unwrap(), &PAYLOAD[8..9]); + assert_eq!( + requests.lock().unwrap().as_slice(), + [ + "bytes=0-0", + "bytes=0-3", + "bytes=4-7", + "bytes=5-5", + "bytes=7-7", + "bytes=8-8" + ] ); } @@ -2198,6 +2376,7 @@ mod tests { .flatten() .unwrap(); + file.set_buffer_size(0).await.unwrap(); let error = file.read_up_to(10).await.unwrap_err(); assert!( error.to_string().contains("returned 5 bytes, expected 10"), @@ -2324,6 +2503,7 @@ mod tests { file.read_up_to(1).await.unwrap_err(), file.seek(0).await.unwrap_err(), file.tell().await.unwrap_err(), + file.set_buffer_size(0).await.unwrap_err(), ] { assert!(error.to_string().contains("already closed"), "got: {error}"); } @@ -2530,10 +2710,114 @@ mod tests { assert!(error.to_string().contains("206"), "got: {error}"); } + #[tokio::test] + async fn impossible_range_capacity_returns_an_error() { + let size = isize::MAX as u64 + 1; + let table = RemoteTable::new_mock( + "videos".to_string(), + move |request| { + let range = request + .headers() + .get(header::RANGE) + .unwrap() + .to_str() + .unwrap(); + let end = if range == "bytes=0-0" { 0 } else { size - 1 }; + http::Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header(header::CONTENT_RANGE, format!("bytes 0-{end}/{size}")) + .header(VERSION_HEADER, "5") + .body(vec![1]) + .unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + let file = table + .probe_files("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + let error = file.read().await.unwrap_err(); + assert!( + error.to_string().contains("cannot allocate blob range"), + "{error}" + ); + assert_eq!(file.tell().await.unwrap(), 0); + } + + #[derive(Debug)] + struct TrackedRangeBody { + bytes: Vec, + dropped: Arc, + } + + impl AsRef<[u8]> for TrackedRangeBody { + fn as_ref(&self) -> &[u8] { + &self.bytes + } + } + + impl Drop for TrackedRangeBody { + fn drop(&mut self) { + self.dropped.store(true, Ordering::Release); + } + } + + #[tokio::test] + async fn close_releases_existing_cache_after_in_flight_read_ends() { + for cancel in [false, true] { + let release = Arc::new(tokio::sync::Barrier::new(2)); + let started = Arc::new(tokio::sync::Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let file = Arc::new(RemoteBlobFile::new( + Arc::new(BlockingRangeRequester { + release: release.clone(), + started: started.clone(), + block_start: 4, + tracked: Some(dropped.clone()), + }), + PAYLOAD.len() as u64, + Some(5), + )); + file.set_buffer_size(4).await.unwrap(); + assert_eq!(file.read_up_to(1).await.unwrap().as_ref(), &PAYLOAD[..1]); + file.seek(3).await.unwrap(); + assert!( + !dropped.load(Ordering::Acquire), + "warm cache was not retained" + ); + let reader = { + let file = file.clone(); + tokio::spawn(async move { file.read_up_to(2).await }) + }; + started.notified().await; + tokio::time::timeout(Duration::from_millis(50), file.close()) + .await + .expect("close waited for I/O") + .unwrap(); + if cancel { + reader.abort(); + assert!(reader.await.unwrap_err().is_cancelled()); + } else { + release.wait().await; + let error = reader.await.unwrap().unwrap_err(); + assert!(error.to_string().contains("already closed"), "{error}"); + } + assert!( + dropped.load(Ordering::Acquire), + "closed handle retained its old window (cancel={cancel})" + ); + } + } + #[derive(Debug)] struct BlockingRangeRequester { release: Arc, started: Arc, + block_start: usize, + tracked: Option>, } #[async_trait::async_trait] @@ -2544,17 +2828,31 @@ mod tests { _mode: RangeRequestMode, _version: Option, ) -> Result { - if range_header == "bytes=0-3" { + let (start, end) = range_header + .strip_prefix("bytes=") + .unwrap() + .split_once('-') + .unwrap(); + let start: usize = start.parse().unwrap(); + let end: usize = end.parse().unwrap(); + if start == self.block_start { self.started.notify_one(); self.release.wait().await; } + let bytes = match &self.tracked { + Some(dropped) if start == 0 => Bytes::from_owner(TrackedRangeBody { + bytes: PAYLOAD[start..=end].to_vec(), + dropped: dropped.clone(), + }), + _ => Bytes::copy_from_slice(&PAYLOAD[start..=end]), + }; let response = http::Response::builder() .status(StatusCode::PARTIAL_CONTENT) .header( header::CONTENT_RANGE, - format!("bytes 0-{}/{}", PAYLOAD.len() - 1, PAYLOAD.len()), + format!("bytes {start}-{end}/{}", PAYLOAD.len()), ) - .body(PAYLOAD.to_vec()) + .body(bytes) .unwrap(); Ok(("hung".to_string(), Response::from(response)).into()) } @@ -2567,6 +2865,8 @@ mod tests { let requester: Arc = Arc::new(BlockingRangeRequester { release: release.clone(), started: started.clone(), + block_start: 0, + tracked: None, }); let file = Arc::new(RemoteBlobFile::new( requester, diff --git a/rust/lancedb/src/remote/table/blobs/aggregate.rs b/rust/lancedb/src/remote/table/blobs/aggregate.rs new file mode 100644 index 000000000..e8e794fa0 --- /dev/null +++ b/rust/lancedb/src/remote/table/blobs/aggregate.rs @@ -0,0 +1,121 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Flatten response batches once, releasing each source buffer as it is copied. + +use std::collections::VecDeque; + +use arrow_array::{Array, LargeBinaryArray}; +use arrow_buffer::{Buffer, MutableBuffer, NullBufferBuilder, OffsetBuffer}; +use arrow_schema::ArrowError; + +pub(super) struct BlobAccumulator { + offsets: Vec, + nulls: NullBufferBuilder, + values: VecDeque, + bytes: i64, +} + +impl BlobAccumulator { + pub(super) fn new(rows: usize) -> Self { + let mut offsets = Vec::with_capacity(rows + 1); + offsets.push(0); + Self { + offsets, + nulls: NullBufferBuilder::new(rows), + values: VecDeque::new(), + bytes: 0, + } + } + + pub(super) fn append(&mut self, array: &LargeBinaryArray) -> Result<(), ArrowError> { + let offsets = array.value_offsets(); + let start = offsets[0]; + let length = offsets[array.len()] - start; + let total = self.bytes.checked_add(length).ok_or_else(|| { + ArrowError::ComputeError("fetch_blobs result exceeds i64 byte offsets".into()) + })?; + self.offsets.extend( + offsets[1..] + .iter() + .map(|offset| self.bytes + (*offset - start)), + ); + if let Some(nulls) = array.nulls() { + self.nulls.append_buffer(nulls); + } else { + self.nulls.append_n_non_nulls(array.len()); + } + if length != 0 { + self.values.push_back( + array + .values() + .slice_with_length(start as usize, length as usize), + ); + } + self.bytes = total; + Ok(()) + } + + pub(super) fn finish(mut self) -> LargeBinaryArray { + let values = match self.values.len() { + 0 => Buffer::default(), + 1 => self.values.pop_front().unwrap(), + _ => { + let mut values = MutableBuffer::new(self.bytes as usize); + // Do not retain all decoded batches alongside the final allocation. + // Drop each source allocation immediately after copying its bytes. + while let Some(chunk) = self.values.pop_front() { + values.extend_from_slice(&chunk); + } + values.into() + } + }; + LargeBinaryArray::new( + OffsetBuffer::new(self.offsets.into()), + values, + self.nulls.finish(), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn single_buffer_is_not_copied() { + let array = LargeBinaryArray::from(vec![Some(b"alpha".as_slice()), None, Some(b"")]); + let mut result = BlobAccumulator::new(3); + result.append(&array).unwrap(); + let output = result.finish(); + assert_eq!(output.values().as_ptr(), array.values().as_ptr()); + assert_eq!(output, array); + } + + #[test] + fn slices_batches_nulls_and_empty_values_keep_order() { + let array = LargeBinaryArray::from(vec![ + Some(b"unused".as_slice()), + Some(b"alpha"), + None, + Some(b""), + Some(b"beta"), + ]); + let mut result = BlobAccumulator::new(6); + result.append(&array.slice(1, 3)).unwrap(); + result.append(&array.slice(4, 1)).unwrap(); + result.append(&array.slice(2, 2)).unwrap(); + assert_eq!( + result.finish(), + LargeBinaryArray::from(vec![ + Some(b"alpha".as_slice()), + None, + Some(b""), + Some(b"beta"), + None, + Some(b""), + ]) + ); + assert_eq!(BlobAccumulator::new(0).finish().len(), 0); + } +} diff --git a/rust/lancedb/src/remote/table/ipc_stream.rs b/rust/lancedb/src/remote/table/ipc_stream.rs index af31a9b28..08b3b1dea 100644 --- a/rust/lancedb/src/remote/table/ipc_stream.rs +++ b/rust/lancedb/src/remote/table/ipc_stream.rs @@ -3,6 +3,10 @@ //! Incremental decoding of response bodies that use Arrow IPC stream framing. +mod framing; + +use framing::{FramingError, IpcFraming}; + use arrow_array::RecordBatch; use arrow_buffer::Buffer; use arrow_ipc::reader::StreamDecoder; @@ -10,7 +14,6 @@ use bytes::Bytes; use datafusion_common::DataFusionError; use datafusion_physical_plan::SendableRecordBatchStream; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; -use futures::stream::BoxStream; use futures::{Stream, StreamExt, TryStreamExt}; use crate::error::{BoxError, Error, Result}; @@ -35,7 +38,7 @@ where E: Into + 'static, { let mut body = IpcStreamBody { - chunks: chunks.map_err(Into::into).boxed(), + framing: IpcFraming::new(chunks.map_err(Into::into).boxed()), chunk: Buffer::default(), decoder: StreamDecoder::new(), finished: false, @@ -62,7 +65,7 @@ where } struct IpcStreamBody { - chunks: BoxStream<'static, std::result::Result>, + framing: IpcFraming, /// The part of the current chunk that has not been decoded yet. chunk: Buffer, decoder: StreamDecoder, @@ -80,11 +83,12 @@ impl IpcStreamBody { async fn next_batch(&mut self) -> Result> { while !self.finished { if self.chunk.is_empty() { - match self.chunks.next().await { - Some(Ok(chunk)) => self.chunk = chunk.into(), - Some(Err(error)) => return Err(self.interrupted(error)), - None if read_end_of_stream(&mut self.decoder) => self.finished = true, - None => { + match self.framing.next_buffer().await { + Ok(Some(chunk)) => self.chunk = chunk, + Err(FramingError::Interrupted(error)) => return Err(self.interrupted(error)), + Err(FramingError::Arrow(error)) => return Err(error.into()), + Ok(None) if read_end_of_stream(&mut self.decoder) => self.finished = true, + Ok(None) => { return Err(self.interrupted( "the response ended before the Arrow end-of-stream marker".into(), )); @@ -208,6 +212,37 @@ mod tests { .collect() } + #[tokio::test] + async fn impossible_ipc_body_lengths_return_errors() { + let batches = batches(); + let messages = messages(&batches[0].schema(), &batches, IpcWriteOptions::default()); + for length in [-1, i64::MAX - 31, i64::MAX] { + let mut message = messages[1].clone(); + let metadata_len = u32::from_le_bytes(message[4..8].try_into().unwrap()) as usize; + let metadata = arrow_ipc::root_as_message(&message[8..8 + metadata_len]).unwrap(); + let field = metadata + ._tab + .vtable() + .get(arrow_ipc::Message::VT_BODYLENGTH); + assert_ne!(field, 0, "fixture must contain a body length field"); + let offset = 8 + metadata._tab.loc() + field as usize; + message[offset..offset + 8].copy_from_slice(&length.to_le_bytes()); + // A valid schema and message metadata followed by just one body byte. + message.truncate(8 + metadata_len + 1); + let (_, returned, error) = read(vec![ + Ok(Bytes::copy_from_slice(&messages[0])), + Ok(Bytes::from(message)), + ]) + .await; + assert!(returned.is_empty()); + let error = error.unwrap().to_string(); + assert!( + error.contains("body length") || error.contains("cannot allocate IPC message"), + "unexpected error for length {length}: {error}" + ); + } + } + /// The schema, the batches yielded, and the error that ended the stream. async fn read( chunks: Vec>, @@ -278,6 +313,78 @@ mod tests { assert_eq!(output, input); } + #[tokio::test] + async fn fragmented_body_keeps_only_its_declared_allocation() { + use arrow_array::LargeBinaryArray; + + let bytes = vec![42; 3 * 1024 * 1024 + 13]; + let array = LargeBinaryArray::from(vec![bytes.as_slice()]); + let batch = RecordBatch::try_from_iter([("blob", Arc::new(array) as ArrayRef)]).unwrap(); + let body = messages( + &batch.schema(), + std::slice::from_ref(&batch), + IpcWriteOptions::default(), + ) + .concat(); + let (_, output, error) = read(chunks(&body, 8191)).await; + assert!(error.is_none(), "{error:?}"); + assert_eq!(output, [batch]); + let blobs = output[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!( + blobs.values().capacity() <= bytes.len() + 256, + "IPC body retained {} bytes for {} value bytes", + blobs.values().capacity(), + bytes.len() + ); + } + + #[rstest] + #[case(false)] + #[case(true)] + #[tokio::test] + async fn zero_body_batches_and_legacy_prefixes(#[case] legacy: bool) { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let input = [vec![], vec![1, 2], vec![]] + .into_iter() + .map(|values| { + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(values))]) + .unwrap() + }) + .collect::>(); + let options = IpcWriteOptions::try_new(8, legacy, arrow_ipc::MetadataVersion::V4).unwrap(); + let body = messages(&schema, &input, options).concat(); + let (_, output, error) = read(chunks(&body, 1)).await; + assert!(error.is_none(), "{error:?}"); + assert_eq!(output, input); + } + + #[tokio::test] + async fn partial_next_prefix_still_delivers_a_zero_body_batch() { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let empty = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(Vec::::new()))], + ) + .unwrap(); + let nonempty = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))]) + .unwrap(); + let messages = messages( + &schema, + &[empty.clone(), nonempty], + IpcWriteOptions::default(), + ); + let cutoff = messages[..2].concat().len() + 1; + let body = messages.concat(); + let (_, output, error) = read(chunks(&body[..cutoff], 1)).await; + assert_eq!(output, [empty]); + assert_interrupted(error, 1, 0); + } + /// The server answers an empty result with the schema and the end-of-stream marker. #[tokio::test] async fn empty_result_keeps_its_schema() { diff --git a/rust/lancedb/src/remote/table/ipc_stream/framing.rs b/rust/lancedb/src/remote/table/ipc_stream/framing.rs new file mode 100644 index 000000000..da20c8e99 --- /dev/null +++ b/rust/lancedb/src/remote/table/ipc_stream/framing.rs @@ -0,0 +1,172 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Frame IPC bodies before decoding, avoiding StreamDecoder's growing scratch buffer. + +use arrow_buffer::Buffer; +use arrow_schema::ArrowError; +use bytes::Bytes; +use futures::StreamExt; +use futures::stream::BoxStream; + +use crate::error::BoxError; + +pub(super) enum FramingError { + Interrupted(BoxError), + Arrow(ArrowError), +} + +#[derive(Clone, Copy)] +enum Frame { + Header { + bytes: [u8; 4], + read: usize, + continuation: bool, + }, + Metadata(usize), + Body(usize), + End, +} + +impl Default for Frame { + fn default() -> Self { + Self::Header { + bytes: [0; 4], + read: 0, + continuation: false, + } + } +} + +pub(super) struct IpcFraming { + chunks: BoxStream<'static, Result>, + chunk: Buffer, + frame: Frame, +} + +impl IpcFraming { + pub(super) fn new(chunks: BoxStream<'static, Result>) -> Self { + Self { + chunks, + chunk: Buffer::default(), + frame: Frame::default(), + } + } + + pub(super) async fn next_buffer(&mut self) -> Result, FramingError> { + let buffer = match self.frame { + Frame::Header { + mut bytes, + mut read, + continuation, + } => { + // Forward even a partial prefix immediately. StreamDecoder may + // need that byte to yield the preceding zero-body record batch. + self.fill_chunk().await?; + let take = self.chunk.len().min(4 - read); + let buffer = self.chunk.slice_with_length(0, take); + self.chunk.advance(take); + bytes[read..read + take].copy_from_slice(&buffer); + read += take; + self.frame = if read == 4 { + match u32::from_le_bytes(bytes) { + u32::MAX if !continuation => Frame::Header { + bytes: [0; 4], + read: 0, + continuation: true, + }, + 0 => Frame::End, + length => Frame::Metadata(length as usize), + } + } else { + Frame::Header { + bytes, + read, + continuation, + } + }; + buffer + } + Frame::Metadata(length) => { + let metadata = self.read_exact(length).await?; + let message = arrow_ipc::root_as_message(&metadata).map_err(|error| { + FramingError::Arrow(ArrowError::IpcError(format!( + "Invalid IPC message: {error}" + ))) + })?; + let length = usize::try_from(message.bodyLength()).map_err(|_| { + FramingError::Arrow(ArrowError::IpcError( + "IPC message body length is negative or exceeds addressable memory".into(), + )) + })?; + self.frame = if length == 0 { + Frame::default() + } else { + Frame::Body(length) + }; + metadata + } + Frame::Body(length) => { + let body = self.read_exact(length).await?; + self.frame = Frame::default(); + body + } + Frame::End => { + // Forward trailing bytes to the decoder so it rejects data after EOS. + if !self.chunk.is_empty() { + return Ok(Some(std::mem::take(&mut self.chunk))); + } + return self + .chunks + .next() + .await + .transpose() + .map(|chunk| chunk.map(Buffer::from)) + .map_err(FramingError::Interrupted); + } + }; + Ok(Some(buffer)) + } + + async fn read_exact(&mut self, length: usize) -> Result { + self.fill_chunk().await?; + if self.chunk.len() >= length { + let buffer = self.chunk.slice_with_length(0, length); + self.chunk.advance(length); + return Ok(buffer); + } + // The declared length is untrusted. Arrow's infallible aligned allocator + // can panic on its layout or abort on OOM; Vec exposes both as errors. + // StreamDecoder handles any array-specific alignment when decoding. + let mut buffer = Vec::new(); + buffer.try_reserve_exact(length).map_err(|error| { + FramingError::Arrow(ArrowError::IpcError(format!( + "cannot allocate IPC message of {length} bytes: {error}" + ))) + })?; + while buffer.len() < length { + self.fill_chunk().await?; + let take = self.chunk.len().min(length - buffer.len()); + buffer.extend_from_slice(&self.chunk[..take]); + self.chunk.advance(take); + } + Ok(Buffer::from_vec(buffer)) + } + + async fn fill_chunk(&mut self) -> Result<(), FramingError> { + while self.chunk.is_empty() { + self.chunk = self + .chunks + .next() + .await + .ok_or_else(|| { + FramingError::Interrupted( + "the response ended before the Arrow end-of-stream marker".into(), + ) + })? + .map_err(FramingError::Interrupted)? + .into(); + } + Ok(()) + } +}