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>

python/python/lancedb/_lancedb.pyi

@@ -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]

python/python/lancedb/remote/table.py

@@ -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:

python/python/lancedb/scannable.py

@@ -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(),
             )
 
 

python/python/lancedb/table.py

@@ -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)

python/python/tests/test_scannable.py

@@ -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

python/python/tests/test_table.py

@@ -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(