From 10f89b3c32577bf054774b5c74346ba24ee79576 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 13 Nov 2026 02:05:09 +0800 Subject: [PATCH] perf: reduce Blob read buffering and expose read window controls (#4979) Remote Blob reads retain and concatenate intermediate IPC buffers across batches and HTTP requests, increasing peak memory. Decode fragmented IPC bodies using their declared length, assemble Blob values once across batches/requests, and collect large-object Range responses into known-size buffers. Impossible IPC/Range allocations return errors instead of panicking. Closing a remote Blob file releases cached bytes immediately, including while a cross-window read is waiting for the network. Rust and Python expose configurable Blob read buffers. Python sync/async `fetch_blob_files` and `open_blob_file(s)` accept `buffer_size`; the default is 512 KiB and zero disables prefetch. Explicit ranges bypass prefetch. New PE batch-endpoint adoption is separate. Malformed-length and retained-cache regressions fail on the previous implementation and pass after the correction. Real HTTP tests cover buffer propagation, close during blocked input, batch/request aggregation, and large-object fallback. The final extension was rebuilt before Python validation. Non-Blob Arrow arrays may still require alignment copies; a contiguous multi-batch result still requires its final allocation. ## Benchmark Peak client RSS increment in KiB, using `python/benchmarks/bench_fetch_blobs_rss.py`, independent clients and real loopback HTTP. Both extensions were rebuilt from clean, unchanged source with the same toolchain/configuration; loaded artifact hashes were checked. Baseline `42f11d13a2755e44b63876b5a5edd93901291150`, after `0a8a76de5286949d80368f95a5c529adf422c187`. | Equivalent fetch workload | Payload | Before | After | | --- | ---: | ---: | ---: | | Single response, one batch | 48 MiB | 63,812 | 63,896 | | Multiple responses, one batch per response | 144 MiB | 423,104 | 195,628 | | Single response, 12 rows per batch | 48 MiB | 130,520 | 71,348 | | Multiple responses, 12 rows per batch | 144 MiB | 363,144 | 172,748 | | One oversized row, Range fallback | 70 MiB | 182,816 | 85,548 | These are one observation per workload and build, not a statistical distribution or memory upper bound. The single-batch workload shows no improvement. Cloud/Sophon deployment, credentialed S3 and long-running stress have not been exercised. --- docs/src/python/python.md | 9 + python/benchmarks/bench_fetch_blobs_rss.py | 183 ++++++++ python/python/lancedb/_blob.py | 19 + python/python/lancedb/_lancedb.pyi | 7 +- python/python/lancedb/remote/table.py | 13 +- python/python/lancedb/table.py | 56 ++- python/python/tests/test_blob_buffering.py | 236 ++++++++++ python/src/table.rs | 27 +- rust/lancedb/src/blob.rs | 24 + rust/lancedb/src/remote/table/blobs.rs | 424 +++++++++++++++--- .../src/remote/table/blobs/aggregate.rs | 121 +++++ rust/lancedb/src/remote/table/ipc_stream.rs | 123 ++++- .../src/remote/table/ipc_stream/framing.rs | 172 +++++++ 13 files changed, 1326 insertions(+), 88 deletions(-) create mode 100644 python/benchmarks/bench_fetch_blobs_rss.py create mode 100644 python/python/tests/test_blob_buffering.py create mode 100644 rust/lancedb/src/remote/table/blobs/aggregate.rs create mode 100644 rust/lancedb/src/remote/table/ipc_stream/framing.rs 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(()) + } +}