mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
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.
This commit is contained in:
@@ -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()
|
||||
@@ -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")
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 == []
|
||||
+25
-2
@@ -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<u64>,
|
||||
buffer_size: Option<usize>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<usize>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
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),
|
||||
|
||||
Reference in New Issue
Block a user