feat(python): add parallel overwrite support to write_dataset#7850
feat(python): add parallel overwrite support to write_dataset#7850AdibaAdi wants to merge 1 commit into
Conversation
github-actions
Bot
added
A-python
📝 WalkthroughWalkthrough
ChangesParallel dataset writes
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant write_dataset
participant write_fragments
participant LanceOperation
Caller->>write_dataset: write Arrow table with n_jobs
write_dataset->>write_fragments: write table partitions concurrently
write_fragments-->>write_dataset: return fragment metadata
write_dataset->>LanceOperation: commit fragments as Overwrite
LanceOperation-->>Caller: produce dataset version
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
Hi @wjones127, this implements the I kept the first version deliberately narrow: parallel writes support in-memory I would especially appreciate your guidance on whether this Python-level orchestration is the right placement, or whether the worker count should eventually move into the Rust write path. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
python/python/tests/test_dataset.py (2)
6359-6367: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winParametrize instead of duplicating the assertion block.
Both
pytest.raisesblocks test the exact same error for different input types; this is precisely the "differ only by inputs" case the guideline calls out for@pytest.mark.parametrize.As per coding guidelines, "Use
@pytest.mark.parametrizefor Python tests that differ only by inputs."♻️ Proposed refactor
-def test_write_dataset_n_jobs_unsupported_input_type(tmp_path: Path): - table = _n_jobs_table(10) - - with pytest.raises(TypeError, match="only supported for in-memory pyarrow.Table"): - lance.write_dataset(table.to_reader(), str(tmp_path / "reader"), n_jobs=2) - - with pytest.raises(TypeError, match="only supported for in-memory pyarrow.Table"): - lance.write_dataset(table.to_pandas(), str(tmp_path / "pandas"), n_jobs=2) +@pytest.mark.parametrize("convert", [lambda t: t.to_reader(), lambda t: t.to_pandas()]) +def test_write_dataset_n_jobs_unsupported_input_type(tmp_path: Path, convert): + table = _n_jobs_table(10) + with pytest.raises(TypeError, match="only supported for in-memory pyarrow.Table"): + lance.write_dataset(convert(table), str(tmp_path / "ds"), n_jobs=2)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/python/tests/test_dataset.py` around lines 6359 - 6367, Refactor test_write_dataset_n_jobs_unsupported_input_type to use pytest.mark.parametrize for the reader and pandas input variants, passing each input and its corresponding output path into one shared pytest.raises assertion. Preserve the existing TypeError message match and n_jobs=2 behavior.Source: Coding guidelines
6369-6396: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winOnly 2 of ~14 unsupported options are covered by tests.
_ensure_parallel_write_supported'sunsupportedlist also rejectsnamespace_client,table_id,initial_bases,target_bases,target_all_bases,base_store_params,commit_lock,progress,auto_cleanup_options,transaction_properties,enable_stable_row_ids,external_blob_mode,allow_external_blob_outside_bases, andblob_pack_file_size_threshold, but onlyschemaandcommit_messageare exercised here. A parametrized test over the full option list would catch a future regression where a new/renamed option is silently dropped from the rejection list.As per coding guidelines, "Every bugfix and feature must have corresponding tests."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/python/tests/test_dataset.py` around lines 6369 - 6396, Expand the n_jobs unsupported-option coverage in test_write_dataset_n_jobs_unsupported_option and test_write_dataset_n_jobs_schema_not_silently_ignored by parametrizing over every option rejected by _ensure_parallel_write_supported, including namespace_client, table_id, initial_bases, target_bases, target_all_bases, base_store_params, commit_lock, progress, auto_cleanup_options, transaction_properties, enable_stable_row_ids, external_blob_mode, allow_external_blob_outside_bases, and blob_pack_file_size_threshold. Assert each option raises the existing ValueError, while preserving the separate schema and commit_message cases.python/python/lance/dataset.py (2)
7415-7431: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winWorker failures propagate without partition context.
_write_partitionre-raises whateverwrite_fragmentsthrows unchanged. When a failure happens withn_jobsworkers running concurrently, the resulting exception gives no indication of which partition (offset/length) failed, making debugging harder than necessary.As per coding guidelines, "Include full error context, including variable names, values, sizes, and types."
♻️ Proposed fix to add partition context on failure
- def _write_partition(partition: pa.Table) -> List[FragmentMetadata]: - return write_fragments( + def _write_partition( + partition_index: int, partition: pa.Table + ) -> List[FragmentMetadata]: + try: + return write_fragments( + partition, + uri, + mode="overwrite", + max_rows_per_file=max_rows_per_file, + max_rows_per_group=max_rows_per_group, + max_bytes_per_file=max_bytes_per_file, + data_storage_version=data_storage_version, + storage_options=storage_options, + ) + except Exception as e: + raise RuntimeError( + f"n_jobs parallel write failed on partition {partition_index} " + f"({partition.num_rows} rows): {e}" + ) from e🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/python/lance/dataset.py` around lines 7415 - 7431, Update the nested _write_partition function to catch write_fragments failures and re-raise them with the partition’s identifying context, including its offset and length (and preserving the original exception as the cause). Keep the existing fragment-writing arguments and executor.map ordering unchanged.Source: Coding guidelines
7404-7413: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winNo upper bound on effective worker/thread count.
num_workers = len(partitions)grows withn_jobsup tonum_rows. A user passing a very largen_jobson a large table could spawn a very large number of OS threads viaThreadPoolExecutor, risking resource exhaustion rather than throughput gains. Consider documenting the caveat prominently or capping the pool size relative toos.cpu_count()/IO parallelism defaults elsewhere in the codebase.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/python/lance/dataset.py` around lines 7404 - 7413, Cap the effective worker count assigned after partition creation in the surrounding dataset operation, using the repository’s existing CPU/IO parallelism limit or an os.cpu_count()-based bound. Preserve the non-empty-partition cap while preventing very large n_jobs values from creating an unbounded ThreadPoolExecutor.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/python/lance/dataset.py`:
- Around line 7305-7446: Move parallel-write validation and orchestration out of
the Python-only helpers _validate_n_jobs, _ensure_parallel_write_supported, and
_write_dataset_parallel into the shared Rust core, exposing a binding-neutral
API that Python and Java can reuse. Preserve the existing n_jobs bounds,
mode/option compatibility checks, partitioning, fragment commit, and failure
semantics while keeping the Python binding as a thin wrapper with consistent
parameter names.
---
Nitpick comments:
In `@python/python/lance/dataset.py`:
- Around line 7415-7431: Update the nested _write_partition function to catch
write_fragments failures and re-raise them with the partition’s identifying
context, including its offset and length (and preserving the original exception
as the cause). Keep the existing fragment-writing arguments and executor.map
ordering unchanged.
- Around line 7404-7413: Cap the effective worker count assigned after partition
creation in the surrounding dataset operation, using the repository’s existing
CPU/IO parallelism limit or an os.cpu_count()-based bound. Preserve the
non-empty-partition cap while preventing very large n_jobs values from creating
an unbounded ThreadPoolExecutor.
In `@python/python/tests/test_dataset.py`:
- Around line 6359-6367: Refactor
test_write_dataset_n_jobs_unsupported_input_type to use pytest.mark.parametrize
for the reader and pandas input variants, passing each input and its
corresponding output path into one shared pytest.raises assertion. Preserve the
existing TypeError message match and n_jobs=2 behavior.
- Around line 6369-6396: Expand the n_jobs unsupported-option coverage in
test_write_dataset_n_jobs_unsupported_option and
test_write_dataset_n_jobs_schema_not_silently_ignored by parametrizing over
every option rejected by _ensure_parallel_write_supported, including
namespace_client, table_id, initial_bases, target_bases, target_all_bases,
base_store_params, commit_lock, progress, auto_cleanup_options,
transaction_properties, enable_stable_row_ids, external_blob_mode,
allow_external_blob_outside_bases, and blob_pack_file_size_threshold. Assert
each option raises the existing ValueError, while preserving the separate schema
and commit_message cases.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 47184fb3-23d4-4e4f-a67f-aa8d21797c86
📒 Files selected for processing (2)
python/python/lance/dataset.pypython/python/tests/test_dataset.py
| # Parallel writes (n_jobs >= 2) currently support only mode="overwrite". | ||
| # | ||
| # "create" and "append" are intentionally excluded: the parallel path writes | ||
| # fragments and then commits them with a single LanceOperation.Overwrite. That | ||
| # operation ("overwrite or create", no read_version) can reproduce neither the | ||
| # atomic create-only guarantee the native writer enforces, nor append's | ||
| # conflict-detection/retry semantics, so supporting those modes here would | ||
| # silently weaken their concurrency guarantees. | ||
| # TODO(#1980): extend parallel writes to create/append once the fragment-commit | ||
| # path can preserve their transaction semantics (most likely a Rust-core impl, | ||
| # which owns these semantics and could be shared across all bindings). | ||
| _PARALLEL_WRITE_SUPPORTED_MODES = ("overwrite",) | ||
|
|
||
|
|
||
| def _validate_n_jobs(n_jobs: Optional[int]) -> None: | ||
| """Validate the ``n_jobs`` argument of :func:`write_dataset`. | ||
|
|
||
| ``None`` and ``1`` are always accepted and select the existing sequential | ||
| write path. Any other value must be an ``int`` >= 1. ``bool`` is rejected | ||
| explicitly because ``True``/``False`` are almost always a mistake for a | ||
| worker count. | ||
| """ | ||
| if n_jobs is None: | ||
| return | ||
| if isinstance(n_jobs, bool) or not isinstance(n_jobs, int): | ||
| raise TypeError( | ||
| f"n_jobs must be an int or None, got {type(n_jobs).__name__}: {n_jobs!r}" | ||
| ) | ||
| if n_jobs < 1: | ||
| raise ValueError(f"n_jobs must be >= 1, got {n_jobs}") | ||
|
|
||
|
|
||
| def _ensure_parallel_write_supported( | ||
| data_obj: ReaderLike, | ||
| mode: str, | ||
| *, | ||
| unsupported: List[tuple[str, bool]], | ||
| ) -> None: | ||
| """Reject requests that the parallel write path cannot faithfully honor. | ||
|
|
||
| Parallel writes (``n_jobs >= 2``) fan a single in-memory table out into | ||
| concurrent :func:`lance.fragment.write_fragments` calls followed by one | ||
| commit. That path only reproduces a subset of :func:`write_dataset` | ||
| behavior, so anything it cannot honor is rejected with a clear error rather | ||
| than being silently ignored or run sequentially. | ||
| """ | ||
| if not isinstance(data_obj, pa.Table): | ||
| raise TypeError( | ||
| "n_jobs >= 2 is only supported for in-memory pyarrow.Table input, got " | ||
| f"{type(data_obj).__name__}. Convert the input to a pyarrow.Table, or " | ||
| "use the default n_jobs=1 for other input types." | ||
| ) | ||
| if mode not in _PARALLEL_WRITE_SUPPORTED_MODES: | ||
| raise ValueError( | ||
| "n_jobs >= 2 currently only supports mode='overwrite' " | ||
| f"(create/append are not yet supported), got mode={mode!r}." | ||
| ) | ||
| set_options = sorted(name for name, is_set in unsupported if is_set) | ||
| if set_options: | ||
| raise ValueError( | ||
| "n_jobs >= 2 cannot be combined with the following option(s): " | ||
| f"{', '.join(set_options)}. Use the default n_jobs=1 to use them." | ||
| ) | ||
|
|
||
|
|
||
| def _write_dataset_parallel( | ||
| table: pa.Table, | ||
| uri: Union[str, Path, LanceDataset], | ||
| *, | ||
| n_jobs: int, | ||
| max_rows_per_file: int, | ||
| max_rows_per_group: int, | ||
| max_bytes_per_file: int, | ||
| data_storage_version: Optional[str], | ||
| storage_options: Optional[Dict[str, str]], | ||
| enable_v2_manifest_paths: bool, | ||
| ) -> LanceDataset: | ||
| """Write ``table`` with ``n_jobs`` parallel ``write_fragments`` calls and | ||
| commit the resulting fragments in a single ``LanceOperation.Overwrite``. | ||
|
|
||
| The caller guarantees ``table`` is a non-empty ``pyarrow.Table``, | ||
| ``n_jobs >= 2`` and ``mode == "overwrite"`` (the only mode parallel writes | ||
| currently support). Fragments are collected in partition order so the | ||
| committed layout is deterministic. | ||
|
|
||
| Failure semantics: if a worker fails, no commit is performed and the dataset | ||
| is never updated, but fragment files written by already-completed workers are | ||
| left uncommitted in storage. This matches the low-level | ||
| :func:`lance.fragment.write_fragments` API used for distributed writes. | ||
| """ | ||
| from concurrent.futures import ThreadPoolExecutor | ||
|
|
||
| from .fragment import write_fragments | ||
|
|
||
| if isinstance(uri, Path): | ||
| uri = os.fspath(uri) | ||
|
|
||
| num_rows = table.num_rows | ||
| # Ceil division so every row falls into exactly one contiguous partition. | ||
| partition_size = -(-num_rows // n_jobs) | ||
| partitions = [] | ||
| offset = 0 | ||
| while offset < num_rows: | ||
| length = min(partition_size, num_rows - offset) | ||
| partitions.append(table.slice(offset, length)) | ||
| offset += length | ||
| # Effective worker count is capped at the number of non-empty partitions, so | ||
| # a small table with a large n_jobs never spawns idle or empty workers. | ||
| num_workers = len(partitions) | ||
|
|
||
| def _write_partition(partition: pa.Table) -> List[FragmentMetadata]: | ||
| return write_fragments( | ||
| partition, | ||
| uri, | ||
| mode="overwrite", | ||
| max_rows_per_file=max_rows_per_file, | ||
| max_rows_per_group=max_rows_per_group, | ||
| max_bytes_per_file=max_bytes_per_file, | ||
| data_storage_version=data_storage_version, | ||
| storage_options=storage_options, | ||
| ) | ||
|
|
||
| with ThreadPoolExecutor(max_workers=num_workers) as executor: | ||
| # executor.map preserves input (partition) order and, when materialized, | ||
| # raises the first worker error before we reach the commit below — so a | ||
| # failed parallel write never produces a partial dataset version. | ||
| per_partition_fragments = list(executor.map(_write_partition, partitions)) | ||
|
|
||
| fragments = [ | ||
| fragment | ||
| for partition_fragments in per_partition_fragments | ||
| for fragment in partition_fragments | ||
| ] | ||
|
|
||
| operation = LanceOperation.Overwrite(table.schema, fragments) | ||
| return LanceDataset.commit( | ||
| uri, | ||
| operation, | ||
| storage_options=storage_options, | ||
| enable_v2_manifest_paths=enable_v2_manifest_paths, | ||
| ) | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
New parallel-write feature lives entirely in Python, not the Rust core.
_validate_n_jobs, _ensure_parallel_write_supported, and all partitioning/threading/commit orchestration in _write_dataset_parallel are pure Python. The repo's own TODO(#1980) comment (Lines 7313-7315) already flags that a Rust-core implementation is the more durable home for this logic ("most likely a Rust-core impl, which owns these semantics and could be shared across all bindings"). Today, if the Java binding wants the same n_jobs capability, it will have to re-derive the same validation rules (mode whitelist, unsupported-option list, worker-count math) independently, risking drift.
I recognize Lance's own "Distributed Write" guide demonstrates this exact write_fragments + ThreadPoolExecutor + LanceOperation.Overwrite pattern as a client-side technique, which somewhat justifies keeping the mechanics here — but the validation logic (which options are compatible with parallel writes, min/max n_jobs) is business logic that should ideally be centralized so it isn't reimplemented per binding.
As per coding guidelines, "Keep Python bindings as thin wrappers and centralize validation and logic in the Rust core" and "Keep Python and Java bindings thin, centralize validation and logic in Rust, and keep parameter names consistent across bindings."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/python/lance/dataset.py` around lines 7305 - 7446, Move parallel-write
validation and orchestration out of the Python-only helpers _validate_n_jobs,
_ensure_parallel_write_supported, and _write_dataset_parallel into the shared
Rust core, exposing a binding-neutral API that Python and Java can reuse.
Preserve the existing n_jobs bounds, mode/option compatibility checks,
partitioning, fragment commit, and failure semantics while keeping the Python
binding as a thin wrapper with consistent parameter names.
Source: Coding guidelines
1 participant
Why
lance.write_datasetcurrently consumes its input as a single stream and writes one file at a time. For large, fully in-memory tables, this can leave available CPU and I/O capacity underused. In #1980, the reporter observed low CPU and network utilization while writing a large in-memory table and improved throughput by manually distributing the work across threads.A maintainer confirmed the single-stream behavior, welcomed an
n_jobsparameter, and provided a reference approach based on parallelwrite_fragmentscalls followed by one commit. This PR brings that approach into the public API for the subset where the existing fragment-commit primitives can preserve correct behavior.What
This PR adds an optional
n_jobsparameter tolance.write_dataset.n_jobs=Noneandn_jobs=1preserve the existing sequential path.n_jobs >= 2partitions an in-memorypyarrow.Tableinto contiguous, non-empty slices.lance.fragment.write_fragments.LanceOperation.Overwrite.concurrent.futures.ThreadPoolExecutor.If a worker raises an exception, the final commit is not performed, so no partial dataset version is created.
Scope and backward compatibility
The parallel path currently supports:
pyarrow.Tableinputsmode="overwrite"Parallel
createandappendare explicitly rejected rather than silently weakening their existing transaction guarantees. The public fragment-commit path cannot currently reproduce the native writer's atomic create-only behavior or append conflict/retry semantics.Unsupported input types and incompatible options also raise descriptive errors.
Existing callers are unaffected because omitting
n_jobs, or setting it to1, follows the original sequential path.Tests
Added coverage for:
n_jobs=Noneandn_jobs=1n_jobs=2andn_jobs=4schema=with parallel modeBenchmark
Local macOS SSD benchmark using identical in-memory tables,
mode="overwrite", one warm-up, and five measured repetitions. Each run verified row count, schema, and full content.These results are local-storage measurements only, not Azure or production proof.
3,000,000 rows × 7 columns (~160 MB)
n_jobs1,500,000 rows × 61 columns (~698 MB)
n_jobsThe local results show a modest but consistent improvement that increases with worker count and workload size. The network-bound object-storage workload described in the issue was not reproduced here.
Known limitations
pyarrow.Tableinputs.mode="overwrite".Open questions for maintainers
n_jobs, as requested in the issue, or follow Lance conventions such asnum_workersor*_concurrency?Acceptance criteria
lance.write_datasetn_jobs=Noneandn_jobs=1Closes #1980