mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 12:08:35 +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:
+5
-1
@@ -625,12 +625,13 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (data, mode, progress=None))]
|
||||
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None))]
|
||||
pub fn add<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
data: PyScannable,
|
||||
mode: String,
|
||||
progress: Option<Py<PyAny>>,
|
||||
write_parallelism: Option<usize>,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let mut op = self_.inner_ref()?.add(data);
|
||||
if mode == "append" {
|
||||
@@ -640,6 +641,9 @@ impl Table {
|
||||
} else {
|
||||
return Err(PyValueError::new_err(format!("Invalid mode: {}", mode)));
|
||||
}
|
||||
if let Some(write_parallelism) = write_parallelism {
|
||||
op = op.write_parallelism(write_parallelism);
|
||||
}
|
||||
if let Some(progress_obj) = progress {
|
||||
let is_callable = Python::attach(|py| progress_obj.bind(py).is_callable());
|
||||
if is_callable {
|
||||
|
||||
Reference in New Issue
Block a user