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:
Will Jones
2026-07-17 07:41:22 -07:00
committed by GitHub
parent f05140f21c
commit 7813907eb7
7 changed files with 403 additions and 4 deletions
+1
View File
@@ -219,6 +219,7 @@ class Table:
data: pa.RecordBatchReader,
mode: Literal["append", "overwrite"],
progress: Optional[Any] = None,
write_parallelism: Optional[int] = None,
) -> AddResult: ...
async def update(
self, updates: Dict[str, str], where: Optional[str]
+8
View File
@@ -574,6 +574,7 @@ class RemoteTable(Table):
on_bad_vectors: str = "error",
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
) -> AddResult:
"""Add more data to the [Table](Table). It has the same API signature as
the OSS version.
@@ -599,6 +600,12 @@ class RemoteTable(Table):
progress: bool, callable, or tqdm-like, optional
A callback or tqdm-compatible progress bar. See
:meth:`Table.add` for details.
write_parallelism: int, optional
Number of partitions to write in parallel. Higher values increase
throughput but also peak memory use, since each partition buffers
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
Returns
-------
@@ -614,6 +621,7 @@ class RemoteTable(Table):
on_bad_vectors=on_bad_vectors,
fill_value=fill_value,
progress=progress,
write_parallelism=write_parallelism,
)
)
finally:
+154 -2
View File
@@ -7,10 +7,158 @@ import sys
from typing import Callable, Iterator, Optional
from lancedb.arrow import to_arrow
import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.dataset as ds
from .pydantic import LanceModel
# pyarrow's default scanner settings are tuned for narrow rows. For wide rows
# (e.g. embedding columns) they buffer a huge read-ahead window in host memory
# and can OOM the client during bulk ingestion. We size the scanner so the
# estimated in-flight memory stays within a budget, while leaving narrow
# datasets on pyarrow's defaults (no throughput regression).
_SCAN_MEMORY_BUDGET_BYTES = 1024 * 1024 * 1024 # ~1 GiB in-flight target
_TARGET_BATCH_BYTES = 16 * 1024 * 1024 # ~16 MiB per batch
_MIN_BATCH_ROWS = 512
# pyarrow defaults (see arrow/dataset ScanOptions); we never exceed these.
_PA_DEFAULT_BATCH_ROWS = 131_072
_PA_DEFAULT_BATCH_READAHEAD = 16
_PA_DEFAULT_FRAGMENT_READAHEAD = 4
# Read-ahead used for wide rows. pyarrow reads a whole parquet row group at a
# time and keeps `batch_readahead` of them resident, so read-ahead depth (not
# batch size) dominates peak memory for wide data; keep both small but leave a
# little prefetch for throughput. Tuned empirically against embedding datasets.
_WIDE_BATCH_READAHEAD = 2
_WIDE_FRAGMENT_READAHEAD = 1
# Estimate for variable-width columns (string/binary/list) whose true width is
# unknown from the schema alone. Only needs to be large enough to flag "wide".
_VARIABLE_WIDTH_ESTIMATE = 128
# Rows peeked from a rescannable source to refine the list-length guess for
# variable-length list columns (e.g. embeddings stored as `list<float32>`
# instead of `list<float32, N>`), whose per-row width the schema can't tell us.
_SAMPLE_ROWS = 10
def _observed_list_length(sample: pa.ChunkedArray) -> Optional[int]:
"""Average element count per row observed in a list/large_list sample."""
if len(sample) == 0:
return None
mean = pc.mean(pc.list_value_length(sample)).as_py()
return None if mean is None else max(1, round(mean))
def _estimate_field_width(
dtype: pa.DataType, sample: Optional[pa.ChunkedArray] = None
) -> int:
if pa.types.is_fixed_size_list(dtype):
return dtype.list_size * _estimate_field_width(dtype.value_type)
if pa.types.is_struct(dtype):
return sum(
_estimate_field_width(
dtype.field(i).type,
pc.struct_field(sample, [i]) if sample is not None else None,
)
for i in range(dtype.num_fields)
)
if pa.types.is_dictionary(dtype):
return _estimate_field_width(dtype.value_type)
if pa.types.is_fixed_size_binary(dtype):
return dtype.byte_width
if pa.types.is_boolean(dtype):
return 1
if (pa.types.is_list(dtype) or pa.types.is_large_list(dtype)) and (
sample is not None
):
observed_length = _observed_list_length(sample)
if observed_length is not None:
return observed_length * _estimate_field_width(dtype.value_type)
# Fixed-width scalars (ints, floats, temporal, decimal) expose bit_width;
# variable-width types (string, binary, list, map, ...) raise ValueError.
try:
return max(1, dtype.bit_width // 8)
except (ValueError, AttributeError):
return _VARIABLE_WIDTH_ESTIMATE
def _estimate_bytes_per_row(
schema: pa.Schema, sample: Optional[pa.Table] = None
) -> int:
return max(
1,
sum(
_estimate_field_width(
field.type, sample.column(field.name) if sample is not None else None
)
for field in schema
),
)
def _sample_head(head: Callable[..., pa.Table]) -> Optional[pa.Table]:
"""Best-effort peek at a few rows to refine the bytes-per-row estimate.
Uses a tight batch size and no read-ahead so the peek itself can't trigger
the wide-row memory blowup this module exists to avoid. Returns None (fall
back to the schema-only estimate) if sampling isn't possible for any
reason, e.g. an empty dataset.
"""
try:
sample = head(
_SAMPLE_ROWS,
batch_size=_SAMPLE_ROWS,
batch_readahead=1,
fragment_readahead=1,
)
except Exception:
return None
return sample if sample.num_rows > 0 else None
def _bounded_scanner_kwargs(
schema: pa.Schema, sample: Optional[pa.Table] = None
) -> dict:
"""Scanner kwargs that cap in-flight memory for wide rows.
Narrow datasets keep pyarrow's defaults unchanged (no throughput
regression). For wide rows (e.g. embedding columns) pyarrow's default
read-ahead buffers many large batches/row-groups at once, which can OOM the
client during bulk ingestion, so we shrink the batch size and read-ahead to
keep the estimated in-flight memory near the budget.
Read-ahead (not just batch size) has to drop: pyarrow reads a whole parquet
row group at a time and keeps `batch_readahead`/`fragment_readahead` of them
resident, so a small batch size alone still pins large row-group buffers.
`sample`, if given, is a small (see `_SAMPLE_ROWS`) table of rows from the
source used to refine the estimate for variable-length list columns (e.g.
embeddings stored without a fixed size), whose width the schema alone
can't tell us.
"""
bytes_per_row = _estimate_bytes_per_row(schema, sample)
# If pyarrow's defaults already stay within budget, leave them alone so
# narrow datasets keep their throughput. A "unit" of in-flight memory is one
# default-sized batch, held `batch_readahead + fragment_readahead` deep.
default_in_flight = (
_PA_DEFAULT_BATCH_ROWS
* bytes_per_row
* (_PA_DEFAULT_BATCH_READAHEAD + _PA_DEFAULT_FRAGMENT_READAHEAD)
)
if default_in_flight <= _SCAN_MEMORY_BUDGET_BYTES:
return {}
# Wide rows: cap batch bytes and pull read-ahead down so only a couple of
# large row-group buffers are resident at once.
batch_size = min(
_PA_DEFAULT_BATCH_ROWS,
max(_MIN_BATCH_ROWS, _TARGET_BATCH_BYTES // bytes_per_row),
)
return {
"batch_size": batch_size,
"batch_readahead": _WIDE_BATCH_READAHEAD,
"fragment_readahead": _WIDE_FRAGMENT_READAHEAD,
}
@dataclass
class Scannable:
@@ -56,10 +204,12 @@ def _from_table(data: pa.Table) -> Scannable:
@to_scannable.register(ds.Dataset)
def _from_dataset(data: ds.Dataset) -> Scannable:
sample = _sample_head(data.head)
scanner_kwargs = _bounded_scanner_kwargs(data.schema, sample)
return Scannable(
schema=data.schema,
num_rows=data.count_rows(),
reader=lambda: data.scanner().to_reader(),
reader=lambda: data.scanner(**scanner_kwargs).to_reader(),
)
@@ -206,10 +356,12 @@ def _register_optional_converters():
@to_scannable.register(lance.LanceDataset)
def _from_lance(data: lance.LanceDataset) -> Scannable:
sample = _sample_head(data.head)
scanner_kwargs = _bounded_scanner_kwargs(data.schema, sample)
return Scannable(
schema=data.schema,
num_rows=data.count_rows(),
reader=lambda: data.scanner().to_reader(),
reader=lambda: data.scanner(**scanner_kwargs).to_reader(),
)
+29 -1
View File
@@ -1199,6 +1199,7 @@ class Table(ABC):
on_bad_vectors: OnBadVectorsType = "error",
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
) -> AddResult:
"""Add more data to the [Table](Table).
@@ -1244,6 +1245,13 @@ class Table(ABC):
with tqdm() as pbar:
table.add(data, progress=pbar)
write_parallelism: int, optional
Number of partitions to write in parallel. Higher values increase
throughput but also peak memory use, since each partition buffers
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
Returns
-------
AddResult
@@ -3158,6 +3166,7 @@ class LanceTable(Table):
on_bad_vectors: OnBadVectorsType = "error",
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
) -> AddResult:
"""Add data to the table.
If vector columns are missing and the table
@@ -3179,6 +3188,12 @@ class LanceTable(Table):
progress: bool, callable, or tqdm-like, optional
A callback or tqdm-compatible progress bar. See
:meth:`Table.add` for details.
write_parallelism: int, optional
Number of partitions to write in parallel. Higher values increase
throughput but also peak memory use, since each partition buffers
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
Returns
-------
@@ -3194,6 +3209,7 @@ class LanceTable(Table):
on_bad_vectors=on_bad_vectors,
fill_value=fill_value,
progress=progress,
write_parallelism=write_parallelism,
)
)
finally:
@@ -4936,6 +4952,7 @@ class AsyncTable:
on_bad_vectors: Optional[OnBadVectorsType] = None,
fill_value: Optional[float] = None,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
) -> AddResult:
"""Add more data to the [Table](Table).
@@ -4960,6 +4977,12 @@ class AsyncTable:
progress: callable or tqdm-like, optional
A callback or tqdm-compatible progress bar. See
:meth:`Table.add` for details.
write_parallelism: int, optional
Number of partitions to write in parallel. Higher values increase
throughput but also peak memory use, since each partition buffers
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
"""
schema = await self.schema()
@@ -4991,7 +5014,12 @@ class AsyncTable:
data = to_scannable(data)
progress, owns = _normalize_progress(progress)
try:
return await self._inner.add(data, mode or "append", progress=progress)
return await self._inner.add(
data,
mode or "append",
progress=progress,
write_parallelism=write_parallelism,
)
except RuntimeError as e:
if "Cast error" in str(e):
raise ValueError(e)
+183
View File
@@ -0,0 +1,183 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import numpy as np
import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.parquet as pq
from lancedb.scannable import (
_PA_DEFAULT_BATCH_ROWS,
_SAMPLE_ROWS,
_VARIABLE_WIDTH_ESTIMATE,
_WIDE_BATCH_READAHEAD,
_WIDE_FRAGMENT_READAHEAD,
_bounded_scanner_kwargs,
_estimate_bytes_per_row,
_sample_head,
to_scannable,
)
def test_estimate_bytes_per_row():
# fixed-width scalars
assert _estimate_bytes_per_row(pa.schema([("a", pa.int64())])) == 8
assert (
_estimate_bytes_per_row(pa.schema([("a", pa.int32()), ("b", pa.float64())]))
== 12
)
assert _estimate_bytes_per_row(pa.schema([("a", pa.bool_())])) == 1
# fixed-size list (embedding) dominates
assert (
_estimate_bytes_per_row(pa.schema([("v", pa.list_(pa.float32(), 768))]))
== 768 * 4
)
# struct sums its children
struct = pa.struct([("x", pa.int32()), ("y", pa.int32())])
assert _estimate_bytes_per_row(pa.schema([("s", struct)])) == 8
# variable-width columns get a flat estimate, not zero
assert _estimate_bytes_per_row(pa.schema([("s", pa.string())])) > 0
def test_estimate_bytes_per_row_uses_sample_for_variable_length_lists():
# A vector column without a fixed size (e.g. `list<float32>` instead of
# `list<float32, 768>`) has no width the schema alone can tell us.
schema = pa.schema([("v", pa.list_(pa.float32()))])
assert _estimate_bytes_per_row(schema) == _VARIABLE_WIDTH_ESTIMATE
sample = pa.table({"v": pa.array([[0.0] * 768], type=pa.list_(pa.float32()))})
assert _estimate_bytes_per_row(schema, sample) == 768 * 4
def test_estimate_bytes_per_row_sample_ignores_missing_or_null_lists():
schema = pa.schema([("v", pa.list_(pa.float32()))])
# an all-null sample column can't tell us anything either
sample = pa.table({"v": pa.array([None], type=pa.list_(pa.float32()))})
assert _estimate_bytes_per_row(schema, sample) == _VARIABLE_WIDTH_ESTIMATE
def test_bounded_scanner_kwargs_narrow_uses_defaults():
# Narrow rows stay on pyarrow defaults (empty kwargs) so throughput is
# unchanged.
for schema in [
pa.schema([("a", pa.int64()), ("b", pa.int32()), ("c", pa.string())]),
pa.schema([("a", pa.int64()), ("t", pa.string()), ("u", pa.string())]),
# a 100-dim float32 vector is still under the per-row budget
pa.schema([("id", pa.int64()), ("v", pa.list_(pa.float32(), 100))]),
]:
assert _bounded_scanner_kwargs(schema) == {}, schema
def test_bounded_scanner_kwargs_wide_is_bounded():
schema = pa.schema(
[
("uid", pa.string()),
("img", pa.list_(pa.float32(), 768)),
("txt", pa.list_(pa.float32(), 768)),
]
)
kwargs = _bounded_scanner_kwargs(schema)
assert kwargs, "wide schema should be throttled"
assert kwargs["batch_readahead"] == _WIDE_BATCH_READAHEAD
assert kwargs["fragment_readahead"] == _WIDE_FRAGMENT_READAHEAD
# batch is capped well below the pyarrow default for wide rows
assert kwargs["batch_size"] < _PA_DEFAULT_BATCH_ROWS
def test_bounded_scanner_kwargs_variable_length_list_needs_sample():
# Without a sample, a variable-length (not fixed-size) vector column looks
# narrow because its true width is unknown from the schema alone.
schema = pa.schema([("uid", pa.string()), ("vec", pa.list_(pa.float32()))])
assert _bounded_scanner_kwargs(schema) == {}
sample = pa.table(
{
"uid": pa.array(["a"]),
"vec": pa.array([[0.0] * 768], type=pa.list_(pa.float32())),
}
)
kwargs = _bounded_scanner_kwargs(schema, sample)
assert kwargs, "sample should reveal the wide vector column"
assert kwargs["batch_size"] < _PA_DEFAULT_BATCH_ROWS
def _write_wide_dataset(
path, *, files=2, rows_per_file=20_000, dim=768, fixed_size=True
):
rng = np.random.default_rng(0)
for i in range(files):
emb = rng.standard_normal((rows_per_file, dim), dtype=np.float32)
vec_type = pa.list_(pa.float32(), dim) if fixed_size else pa.list_(pa.float32())
vec_array = (
pa.FixedSizeListArray.from_arrays(pa.array(emb.reshape(-1)), dim)
if fixed_size
else pa.array(emb.tolist(), type=vec_type)
)
table = pa.table(
{
"uid": pa.array([f"{i}_{j}" for j in range(rows_per_file)]),
"vec": vec_array,
}
)
pq.write_table(table, f"{path}/part-{i}.parquet")
def test_dataset_reader_respects_bounded_batch_size(tmp_path):
# The Dataset path should stream small batches for wide rows, not pyarrow's
# 131072-row default, and still return every row.
_write_wide_dataset(str(tmp_path))
dataset = ds.dataset(str(tmp_path), format="parquet")
expected = _bounded_scanner_kwargs(dataset.schema)["batch_size"]
scannable = to_scannable(dataset)
assert scannable.rescannable
assert scannable.num_rows == 40_000
total = 0
for batch in scannable.reader():
assert batch.num_rows <= expected
total += batch.num_rows
assert total == 40_000
# factory can be called again (rescannable)
assert sum(b.num_rows for b in scannable.reader()) == 40_000
def test_dataset_reader_samples_variable_length_list_width(tmp_path):
# A vector column stored without a fixed size (e.g. produced by tools that
# don't tag list columns with their length) is invisible to the
# schema-only estimate, so `to_scannable` must peek a sample of rows to
# detect that it's wide and bound the scanner accordingly.
_write_wide_dataset(str(tmp_path), fixed_size=False)
dataset = ds.dataset(str(tmp_path), format="parquet")
schema_only_kwargs = _bounded_scanner_kwargs(dataset.schema)
assert schema_only_kwargs == {}, "schema alone can't see the list width"
scannable = to_scannable(dataset)
assert scannable.rescannable
assert scannable.num_rows == 40_000
total = 0
for batch in scannable.reader():
assert batch.num_rows < _PA_DEFAULT_BATCH_ROWS
total += batch.num_rows
assert total == 40_000
def test_sample_head_is_bounded_rows(tmp_path):
# The peek itself must not read the whole dataset.
_write_wide_dataset(str(tmp_path), files=1, rows_per_file=1000, fixed_size=False)
dataset = ds.dataset(str(tmp_path), format="parquet")
sample = _sample_head(dataset.head)
assert sample.num_rows == _SAMPLE_ROWS
def test_sample_head_returns_none_for_empty_dataset(tmp_path):
table = pa.table({"v": pa.array([], type=pa.list_(pa.float32()))})
pq.write_table(table, f"{tmp_path}/empty.parquet")
dataset = ds.dataset(str(tmp_path), format="parquet")
assert _sample_head(dataset.head) is None
+23
View File
@@ -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(
+5 -1
View File
@@ -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 {