diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index d6e7f3407..3b94c5d45 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/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] diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 1006942ea..8c702aace 100644 --- a/python/python/lancedb/remote/table.py +++ b/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: diff --git a/python/python/lancedb/scannable.py b/python/python/lancedb/scannable.py index b26aac2fa..baf8e0c2b 100644 --- a/python/python/lancedb/scannable.py +++ b/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` +# instead of `list`), 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(), ) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 20ad2e742..6739d5864 100644 --- a/python/python/lancedb/table.py +++ b/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) diff --git a/python/python/tests/test_scannable.py b/python/python/tests/test_scannable.py new file mode 100644 index 000000000..2ec44de63 --- /dev/null +++ b/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` instead of + # `list`) 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 diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index ac331838e..456b900ee 100644 --- a/python/python/tests/test_table.py +++ b/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( diff --git a/python/src/table.rs b/python/src/table.rs index c9784facb..703c64c53 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -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>, + write_parallelism: Option, ) -> PyResult> { 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 {