feat: configurable memory limit for merge_insert's DataFusion join

[yeung108]

Problem

The original ask in #1983: merge_insert matches source rows against the target table via a DataFusion join with no configured RAM limit, so it can use unbounded memory. This PR addresses that directly (see #7718 for a separate, more fundamental fix to the same issue's underlying cause — a cache-invalidation gap that made the problem look worse than it is and affects every writer, not just merge_insert).

What I found

merge_insert actually has two join code paths, and only one had any memory bound:

• Indexed-scan path (create_indexed_scan_joined_stream, used when every join column has a scalar index and there's no delete_not_matched_by_source clause): already runs through execute_plan() with use_spilling: true, so it gets a memory-pool-bounded, spilling-capable SessionContext.
• Full-table-scan path (create_full_table_joined_stream, used whenever the join key isn't fully indexed — including the common case of no index at all — or a delete_not_matched_by_source clause forces a full scan): built its SessionContext via a bare SessionContext::new_with_config(SessionConfig::default().with_target_partitions(1)). No memory limit, no spilling, nothing.

The full-table-scan path is the default for any merge_insert on a table without an index on the merge key — exactly the scenario in the original reproduction in #1983 (and the first repro in the linked investigation): a single call's RSS scaled with the existing target table size, unbounded.

Fix

Both join paths now build their execution context through a new MergeInsertJob::join_execution_options() helper, which always enables spilling (LanceExecutionOptions::use_spilling: true) and threads through two new caller-configurable knobs:

• MergeInsertBuilder::mem_pool_size(bytes) — memory pool limit for the join, defaulting to LanceExecutionOptions's existing default (150MB, or the LANCE_MEM_POOL_SIZE env var).
• MergeInsertBuilder::max_temp_directory_size(bytes) — max spill directory size, defaulting to the existing 100GB default.

Both default to None or LanceExecutionOptions' own defaults, so this is non-breaking — existing callers get the same defaults, but the full-table-scan path now actually gets a memory pool and spilling instead of none at all, and everyone gets the ability to tune it.

Exposed through to Python as MergeInsertBuilder.mem_pool_size() / .max_temp_directory_size(), following the existing conflict_retries/use_index pattern (Rust #[pymethods] + the documented Python wrapper class in python/python/lance/dataset.py).

Testing

• test_merge_insert_mem_pool_size_is_configurable: unit-level check that the builder methods set MergeInsertParams.mem_pool_size/max_temp_directory_size, and that join_execution_options() correctly derives LanceExecutionOptions from them (use_spilling: true, both sizes threaded through).
• test_merge_insert_full_table_join_with_small_mem_pool_size: end-to-end smoke test — a merge_insert forced onto the full-table-scan path via use_index(false), with a deliberately small mem_pool_size (256KB, vs. the 150MB default), completes successfully with correct row counts. This exercises the actual code path changed here, not just config plumbing.
• All 171 existing tests in dataset::write::merge_insert:: pass unmodified, including the plan-shape assertions (test_explain_plan*) and the various indexed/no-index composite-key tests — confirming this doesn't change join algorithm selection or plan structure, only the execution context's memory configuration.
• cargo check/fmt --check/clippy -D warnings clean on the lance crate and the pylance (python bindings) crate, ruff format/ruff check clean on the modified Python file — all run under the toolchain/tool versions pinned by this repo's rust-toolchain.toml and .pre-commit-config.yaml.

Summary by CodeRabbit

• New Features

  • Added merge operation configuration options to limit join memory usage and control maximum spill-to-disk size.
  • Available via the Python merge builder with chainable methods.

• Bug Fixes

  • Merge execution now applies the configured memory and temporary spill limits consistently across all join execution paths.
  • Improved correctness for constrained-memory scenarios, verified with new regression coverage.

[coderabbitai]

📝 Walkthrough

Walkthrough

MergeInsertBuilder now supports join memory-pool and temporary spill-directory limits. The settings are exposed through Python bindings, propagated across merge execution paths, and covered by propagation and end-to-end tests.

Changes

Merge insert memory configuration

Layer / File(s)	Summary
Configuration API and bindings rust/lance/src/dataset/write/merge_insert.rs, python/src/dataset.rs, python/python/lance/dataset.py	Adds optional memory and temporary-directory limits, Rust setters, and chainable Python-facing builder methods.
Join execution option propagation rust/lance/src/dataset/write/merge_insert.rs	Applies configured spilling and size limits to indexed-scan, full-table-scan, v2, and plan-analysis execution contexts.
Configuration and merge regression coverage rust/lance/src/dataset/write/merge_insert.rs	Tests option propagation and verifies a full-table merge with a small memory pool.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PythonBuilder
  participant RustMergeInsertBuilder
  participant JoinExecution
  PythonBuilder->>RustMergeInsertBuilder: set memory and spill limits
  RustMergeInsertBuilder->>JoinExecution: derive join execution options
  JoinExecution->>JoinExecution: execute indexed or full-table merge

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately captures the main change: configurable join memory limits for merge_insert.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 4

🤖 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 539-585: Add `Examples` sections to the `mem_pool_size` and
`max_temp_directory_size` docstrings in `MergeInsertBuilder`, demonstrating
configuration through the actual chained builder API with valid
Rust/PyO3-compatible signatures; include references or links to the relevant
builder methods and keep both examples consistent with the documented parameter
units.

In `@rust/lance/src/dataset/write/merge_insert.rs`:
- Around line 583-606: The public setters mem_pool_size and
max_temp_directory_size lack required examples, and their LanceExecutionOptions
references are not intra-doc links. Update both doc comments to link using
[`LanceExecutionOptions`] and add concise `# Examples` sections demonstrating
each setter.
- Around line 8770-8819: Update
test_merge_insert_full_table_join_with_small_mem_pool_size so the merge
configuration makes can_use_create_plan() return false, rather than relying on
use_index(false), which only disables the scalar-index path. Adjust the
matched/not-matched actions or other setup to force execute_uncommitted_impl()
through create_joined_stream() and create_full_table_joined_stream(), then
retain assertions verifying successful completion and correct results.
- Around line 8820-8839: Ensure the spill-path test actually exercises spilling:
enlarge the target and source datasets beyond what the 256 KiB pool can hold, or
add an explicit assertion that temporary spill files were created. Update the
expected inserted, updated, total, and value-filtered counts in the test
surrounding job.execute_reader and the subsequent count_rows assertions.

🪄 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: 4820dbf7-6f5a-4a88-8794-c45fa2e68101

📥 Commits

Reviewing files that changed from the base of the PR and between c5c3bf1 and 9f586df.

📒 Files selected for processing (3)

• python/python/lance/dataset.py
• python/src/dataset.rs
• rust/lance/src/dataset/write/merge_insert.rs