feat: accept TableProvider write inputs for merge_insert and insert#7368
Conversation
github-actions
Bot
added
A-python
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Make `Arc<dyn TableProvider>` the canonical internal write input behind ergonomic wrappers. Re-readable sources are now replayed across retries without spilling to disk, and materialized sources report statistics that let DataFusion choose the merge-join build side. - merge_insert: add `execute_provider` (canonical) and `execute_batches` (multi-partition `MemTable`); `execute(stream)` becomes a wrapper. Retries re-scan the provider instead of the removed `new_source_iter`/ `SpillStreamIter` replay layer. A one-shot stream spills only when retries are enabled; `spill_for_retry(false)` fails fast instead of buffering. - Plan against the provider directly so a MemTable/file source's statistics reach the join; stream sources keep the one-shot path (no stats lost, and the source's original error type is preserved). - InsertBuilder gains `execute_provider`/`execute_uncommitted_provider`. - Python routes materialized inputs (pa.Table, RecordBatch, DataFrame, ...) through the in-memory path; streams and scanners keep spilling. Re-scannable Python providers (pa.dataset.Dataset/Scanner) and parallel data-file writes over provider partitions remain follow-ups. Issue: lance-format#4583 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Always plan against the source TableProvider directly (read_table), so every source — including spilled and one-shot streams — exposes its statistics to the merge join. The merge write node already requires a single-partition input, so the optimizer coalesces multi-partition providers; the previous target_partitions(1) hack and the scan_provider_directly branch are removed. Drop the provider_to_stream first-batch peek that preserved the source error's concrete type. Source errors are shared across join partitions by DataFusion (DataFusionError::Shared), so the type cannot be recovered, and no caller needs it — Python surfaces these errors by message. The error conversion now handles Shared (recursing when sole-owner, otherwise preserving the message under the execution category). Revert the InsertBuilder provider methods: they only adapt a provider back to a stream and add no parallelism until fragment fan-out exists. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Verify that merging an empty batch list is a no-op that leaves the target unchanged, exercising the empty-source partition and schema-fallback paths in batches_to_provider / batches_into_partitions. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
b340516 to
64b7576
Compare
📝 WalkthroughWalkthroughMerge insert execution now uses DataFusion table providers for retryable sources, adds replayable spilling and provider-based planning, exposes batch execution through Python bindings, and routes materialized Python inputs through batch paths. ChangesMerge insert provider flow
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant PythonClient
participant MergeInsertBuilder
participant TableProvider
participant MergeInsertJob
participant Dataset
PythonClient->>MergeInsertBuilder: execute materialized or streaming source
MergeInsertBuilder->>TableProvider: wrap source
MergeInsertBuilder->>MergeInsertJob: execute provider
MergeInsertJob->>TableProvider: re-scan source for retries
MergeInsertJob->>Dataset: commit transaction and merge statistics
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![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
🤖 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 `@rust/lance-datafusion/src/exec.rs`:
- Around line 880-884: The public helper documentation lacks runnable examples.
In rust/lance-datafusion/src/exec.rs lines 880-884, add a `# Examples` section
to `provider_to_stream` demonstrating scan-to-stream usage; in
rust/lance-datafusion/src/spill.rs lines 67-76, add a `# Examples` section to
`spilling_table_provider` showing wrapping a one-shot stream and rescanning.
Ensure both examples are valid and compile as documentation tests.
🪄 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: ASSERTIVE
Plan: Pro Plus
Run ID: 4b774129-56ff-48cb-aef4-f8e5a398c26b
📒 Files selected for processing (10)
python/python/lance/dataset.pypython/python/lance/lance/__init__.pyipython/python/lance/types.pypython/python/tests/test_dataset.pypython/src/dataset.rsrust/lance-core/src/error.rsrust/lance-datafusion/src/exec.rsrust/lance-datafusion/src/spill.rsrust/lance/src/dataset/write.rsrust/lance/src/dataset/write/merge_insert.rs
hamersaw
left a comment
hamersaw
left a comment
There was a problem hiding this comment.
Looks great! Just a few clarification questions more so for my understanding than anything.
| yield batch | ||
|
|
||
|
|
||
| def _is_materialized(data_obj: ReaderLike) -> bool: |
There was a problem hiding this comment.
Maintaining this list seems a little difficult. IIUC we support all ArrowArrayStreamReader parsable types (presumably huge?). And this function is attempting to decide which are oneshot and need to be spilled to disk. It may not be a huge concern because worst case scenario we just double store on ingest, but notably we are missing things like Scanner (Lances), pa.dataset.Dataset, or LanceDataset which can be reread without spilling.
There was a problem hiding this comment.
Good point on the drift. The classifier answers "can I cheaply wrap this in an in-memory MemTable?" — a Scanner/LanceDataset is re-readable but not in-memory, so materializing it just to skip a spill could be worse than the spill. But you're right that re-scanning it (no spill, no double-store) is strictly better than what we do today. Filed #7895 to make source classification a singledispatch registry and add a re-scannable provider for exactly these types.
| .map_err(|e| DataFusionError::Execution(format!("Failed to spawn temp dir task: {e}")))? | ||
| .map_err(|e| DataFusionError::Execution(format!("Failed to create temp dir: {e}")))?; | ||
| let tmp_path = tmp_dir.std_path().join("spill.arrows"); | ||
| let (mut sender, receiver) = create_replay_spill(tmp_path, schema.clone(), memory_limit); |
There was a problem hiding this comment.
So we create the persistent spill recovery (either complete in-memory or + spill file) before beginning any of the insert / merge_insert logic. I don't think it's useful to explore right now, but it feels like some level of parallelization could really help performance here.
There was a problem hiding this comment.
Worth noting the spill is already streaming: a reader can start pulling batches as soon as they land (it doesn't block on the full drain), and it stays in memory entirely unless it exceeds the 100 MB budget. So the "create spill before insert logic" step isn't a serial full-buffer barrier. Agreed there's more parallelism to chase later, though. (I've pulled these two properties up into the spilling_table_provider doc so they're visible at the call site.)
| // When the source is a one-shot stream and `conflict_retries > 0`, the source | ||
| // is spilled (memory, then disk) so it can be replayed on each retry. Set to | ||
| // false to fail fast on contention instead of buffering the stream. Has no | ||
| // effect on re-scannable sources (materialized batches, files), which never | ||
| // spill. | ||
| spill_for_retry: bool, |
There was a problem hiding this comment.
It looks like the previous new_source_iter had an enable_retries function that was set if retries were configured. So users didn't have to manually set this, instead it was auto-inferred. I don't have great context, but would it be possible to use the same "enable on retries" and, similar to the python API, chose not to spill internally when we know the source is not oneshot? This could take the onus off users to define spill vs not.
There was a problem hiding this comment.
It's already automatic in that direction — spilling only kicks in when conflict_retries > 0, and re-scannable providers (materialized batches, files) never spill regardless. spill_for_retry is only an opt-out (default true) to fail-fast instead of buffering a genuine one-shot stream. The gap is that a plain SendableRecordBatchStream is opaque, so we can't tell re-scannable from one-shot at that boundary — which is what #7895 addresses on the Python side.
Add compiling doctests to the two new public helpers `provider_to_stream` and `spilling_table_provider`, per the public-API-needs-examples guideline. Also surface two spill properties on `spilling_table_provider` that were only documented on the lower-level `create_replay_spill`: it buffers in memory first and only touches disk past `memory_limit`, and a scan can stream batches as they land rather than blocking on the full drain. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@rust/lance-datafusion/src/spill.rs`:
- Around line 69-84: Update the public documentation block for
create_replay_spill to add intra-doc links to SpillPartition, SpillReceiver,
create_replay_spill, and StreamingTable wherever those concepts are referenced.
Keep the existing behavioral explanation unchanged and ensure each link resolves
to the actual visible Rust item.
🪄 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: ASSERTIVE
Plan: Pro Plus
Run ID: 196f7a8c-b007-41cc-99b5-05f12d0dc80e
📒 Files selected for processing (2)
rust/lance-datafusion/src/exec.rsrust/lance-datafusion/src/spill.rs
| /// The source is drained in the background into a replayable spill. Two properties | ||
| /// keep this cheap for the common case: | ||
| /// | ||
| /// - **Memory-first.** Up to `memory_limit` bytes are buffered in memory; the spill | ||
| /// only touches disk once that budget is exceeded. A source that fits under the | ||
| /// limit never hits the filesystem. | ||
| /// - **Streaming replay.** A scan can start consuming batches as soon as they land, | ||
| /// before the source has finished draining — the first reader is not blocked | ||
| /// waiting for the whole source to buffer. | ||
| /// | ||
| /// Each scan of the returned provider replays the full source, which is what makes a | ||
| /// one-shot stream usable in the write retry loop. | ||
| /// | ||
| /// The provider reports no statistics — the source size is not known until it has | ||
| /// been fully drained — so callers that need source statistics (e.g. to drive join | ||
| /// ordering) should prefer a materialized or file-backed provider instead. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Add intra-doc links to referenced types.
The new doc block explains behavior well but never links to SpillPartition, SpillReceiver, create_replay_spill, or StreamingTable, which readers would want to jump to.
As per coding guidelines: "Document all public APIs with examples and links to relevant structs and methods; keep examples synchronized with actual signatures."
📝 Example link additions
-/// Each scan of the returned provider replays the full source, which is what makes a
+/// Each scan of the returned provider (backed by [`SpillPartition`]) replays the full source,
+/// which is what makes a
/// one-shot stream usable in the write retry loop.🤖 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 `@rust/lance-datafusion/src/spill.rs` around lines 69 - 84, Update the public
documentation block for create_replay_spill to add intra-doc links to
SpillPartition, SpillReceiver, create_replay_spill, and StreamingTable wherever
those concepts are referenced. Keep the existing behavioral explanation
unchanged and ensure each link resolves to the actual visible Rust item.
Source: Coding guidelines
2 participants
Implements #4583: make
Arc<dyn TableProvider>the canonical internal write input behind ergonomic wrappers, so re-readable sources replay across retries without spilling and materialized sources expose statistics to the merge join.Changes
execute_provider(canonical entry) andexecute_batches(multi-partitionMemTable);execute(stream)is now a thin wrapper. The retry loop re-scans the provider per attempt, replacing thenew_source_iter/SpillStreamIterreplay layer. A one-shot stream spills (memory→disk) only whenconflict_retries > 0;spill_for_retry(false)fails fast on contention instead of buffering the stream.execute_provider/execute_batchesplan against the provider viaread_table, so aMemTable/file source's exactnum_rows/total_byte_sizereach the optimizer and drive join build-side selection. Stream sources keep the one-shot path — they carry no statistics anyway, and this preserves the source's original error type.execute_provider/execute_uncommitted_providerfor input-shape parity (no retry/spill today).pa.Table,pa.RecordBatch, pandas/polars frames, dict/list-of-dict) route through the new in-memoryexecute_batchespath; streams, readers, scanners, and datasets keep the spilling path.Tests
execute_provider,execute_batches, source-stats-in-plan,spill_for_retry(false)fail-fast, andInsertBuilder::execute_provider; fulldataset::write::suite + spill tests green; clippy/fmt clean.Follow-ups (out of scope here)
pa.dataset.Dataset/Scanner(currently still spill — needs a Python-callbackTableProvider).🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes