mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 20:18:37 +00:00
fix(python): bound scanner memory for wide-row bulk ingestion (#3625)
## Problem `table.add(dataset)` with a `pyarrow.dataset.Dataset` OOMs the client during bulk ingestion of wide rows (e.g. embedding columns), even against a remote table where the upload itself is streaming. The cause is in `to_scannable`: a `Dataset` is scanned with pyarrow's default scanner settings (`batch_size=131072` rows, `batch_readahead=16`, `fragment_readahead=4`). pyarrow's internal threads prefetch that read-ahead window independently of LanceDB's backpressure, so for wide rows a large fraction of the dataset is held in memory. On the remote path this is then multiplied across the multipart write partitions (one in-flight batch per partition, up to CPU-core count). Reproduced on a 10 GB / 1.55M-row dataset with two 768-dim float32 embeddings: peak client RSS ~11.7 GB for the scan alone (6.8 GB after consuming a *single* batch), ~15.4 GB for the full remote `add()`. ## Fix `to_scannable` now sizes the scanner from an estimate of bytes-per-row derived from the schema: - **Narrow datasets keep pyarrow's defaults** (empty scanner kwargs) — no throughput regression. The bound only engages above ~410 bytes/row. - **Wide rows** get a smaller `batch_size` (~16 MiB/batch) and reduced read-ahead (`batch_readahead=2`, `fragment_readahead=1`) so peak in-flight memory stays near a ~1 GiB budget. Read-ahead (not just batch size) has to drop, because pyarrow pins whole row-group buffers. On the 10 GB dataset this drops peak client RSS to ~1.4 GB, and it stays flat as the dataset grows. The `Dataset`/`LanceDataset` scannables remain rescannable (retry-safe). ## Also: expose `write_parallelism` on `add()` `AddDataBuilder::write_parallelism` already existed in Rust but was not exposed in Python. This PR forwards it through the async, sync, and remote `add()` methods, so users can cap the number of parallel write partitions (each buffers data in flight) to trade throughput for memory on large uploads. ## Tests - `test_scannable.py`: bytes-per-row estimation; narrow → defaults; wide → bounded; `Dataset` reader streams bounded batches and stays rescannable. - `test_table.py`: `write_parallelism` on sync and async `add()`, and that `write_parallelism=0` is rejected. Fixes ENT-1883 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -434,6 +434,29 @@ def test_add(mem_db: DBConnection):
|
||||
_add(table, schema)
|
||||
|
||||
|
||||
def test_add_write_parallelism(mem_db: DBConnection):
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
table = mem_db.create_table("test", schema=schema)
|
||||
|
||||
data = pa.table({"id": list(range(1000))}, schema=schema)
|
||||
table.add(data, write_parallelism=4)
|
||||
assert len(table) == 1000
|
||||
|
||||
# invalid parallelism is rejected
|
||||
with pytest.raises(ValueError, match="write_parallelism"):
|
||||
table.add(data, write_parallelism=0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_write_parallelism_async(mem_db_async: AsyncConnection):
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
table = await mem_db_async.create_table("test", schema=schema)
|
||||
|
||||
data = pa.table({"id": list(range(1000))}, schema=schema)
|
||||
await table.add(data, write_parallelism=4)
|
||||
assert await table.count_rows() == 1000
|
||||
|
||||
|
||||
def test_add_struct(mem_db: DBConnection):
|
||||
# https://github.com/lancedb/lancedb/issues/2114
|
||||
schema = pa.schema(
|
||||
|
||||
Reference in New Issue
Block a user