fix(python): bound scanner memory for wide-row bulk ingestion#3625
Conversation
Previously, `table.add(dataset)` with a `pyarrow.dataset.Dataset` (or `lance.LanceDataset`) used pyarrow's default scanner settings (`batch_size=131072`, `batch_readahead=16`, `fragment_readahead=4`). For wide rows (e.g. embedding columns) this buffers a large read-ahead window in host memory and can OOM the client during bulk ingestion, even though the add path is otherwise streaming. `to_scannable` now sizes the scanner from an estimate of bytes-per-row: narrow datasets keep pyarrow's defaults (no throughput regression), while wide rows get a smaller batch size and reduced read-ahead so peak in-flight memory stays near a ~1 GiB budget. On a 10 GB embedding dataset this drops peak client RSS from ~11.7 GB to ~1.4 GB. Also exposes `write_parallelism` on `add()` (async, sync, and remote), forwarding to the existing Rust `AddDataBuilder` builder. This lets users cap the number of parallel write partitions, which each buffer data in flight, trading throughput for memory on large uploads. Fixes ENT-1883 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| progress, owns = _normalize_progress(progress) | ||
| try: | ||
| return await self._inner.add(data, mode or "append", progress=progress) | ||
| return await self._inner.add( |
There was a problem hiding this comment.
Could we route the sanitize path above through the same bounded Dataset scanner as to_scannable()? For mode="overwrite", non-default on_bad_vectors, and embedding-function tables, AsyncTable.add() converts a pyarrow.dataset.Dataset / lance.LanceDataset inside _sanitize_data() / _into_pyarrow_reader() before reaching this call. Those dataset branches still call bare data.scanner().to_reader(), so these table.add(dataset) cases keep PyArrow's default read-ahead and can still hit the wide-row OOM this PR is trying to avoid.
| try: | ||
| return max(1, dtype.bit_width // 8) | ||
| except (ValueError, AttributeError): | ||
| return _VARIABLE_WIDTH_ESTIMATE |
There was a problem hiding this comment.
Could we make the variable-length list estimate less optimistic here? For list<float32> / large_list<float32> columns, the schema does not always include the vector dimension, so the current 128-byte fallback can classify common embedding-shaped rows as "narrow" and leave PyArrow's default batch/read-ahead settings in place. A more robust approach might be to sample a small number of rows or the first scanner batch, compute the observed list lengths for numeric list columns, and use that sampled dimension to estimate bytes per row before deciding whether to bound batch_size, batch_readahead, and fragment_readahead.
There was a problem hiding this comment.
Yeah, I suppose we can get a small sample from a rescannable source. That would be the best approach.
There was a problem hiding this comment.
Pushed in 20ebde7: to_scannable now peeks up to 10 rows from Dataset/LanceDataset sources (tight batch size, no read-ahead, so the peek itself can't trigger the blowup we're fixing) and uses the observed average list length to size variable-length list/large_list columns, instead of falling back to the flat 128-byte estimate.
Schema alone can't tell us the width of a `list<float32>` column without a fixed size (e.g. embeddings not tagged with their dimension), so it fell back to a flat 128-byte estimate and could be misclassified as narrow, leaving pyarrow's default read-ahead in place for what's actually a wide-row dataset. Peek up to 10 rows from the source (with a tight batch size and no read-ahead, so the peek itself can't trigger the memory blowup this module exists to avoid) and use the observed average list length to refine the estimate. Addresses PR feedback: lancedb#3625 (comment)
|
Docs PR opened: lancedb/docs#312 Documented the new write_parallelism parameter on Python add() and the auto-bounded scan behavior for wide-row dataset ingestion. |
2 participants
Problem
table.add(dataset)with apyarrow.dataset.DatasetOOMs 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: aDatasetis scanned with pyarrow's default scanner settings (batch_size=131072rows,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_scannablenow sizes the scanner from an estimate of bytes-per-row derived from the schema: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/LanceDatasetscannables remain rescannable (retry-safe).Also: expose
write_parallelismonadd()AddDataBuilder::write_parallelismalready existed in Rust but was not exposed in Python. This PR forwards it through the async, sync, and remoteadd()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;Datasetreader streams bounded batches and stays rescannable.test_table.py:write_parallelismon sync and asyncadd(), and thatwrite_parallelism=0is rejected.Fixes ENT-1883
🤖 Generated with Claude Code