Skip to content

feat(datafusion): accept a caller-supplied memory pool - #9801

Draft
wjones127 wants to merge 1 commit into
mainfrom
claude/lance-issue-9799-61a8c8
Draft

wjones127 wants to merge 1 commit into
mainfrom
claude/lance-issue-9799-61a8c8

Conversation

@wjones127

Copy link
Copy Markdown
Contributor

Previously, callers could only set the size of the DataFusion memory pool (mem_pool_size / LANCE_MEM_POOL_SIZE). Lance decided how that pool was shared, and since #9408 each spilling execution gets its own fresh pool. As a result, N concurrent merge_insert or delta operations can use N times the configured budget, and a server has no way to enforce a single node-wide limit.

This PR adds LanceExecutionOptions::memory_pool. When it is set, executions use that pool as given, whether or not spilling is enabled, and mem_pool_size is ignored. Sessions built with a supplied pool skip the session cache. The sort spill reservation comes from the pool's limit when that limit is finite. Callers who don't set it get the same behavior as before.

The option can also be set per operation:

  • Scanner::memory_pool covers scans, count_rows, and analyze_plan.
  • MergeInsertBuilder::memory_pool covers the plan-based merge, the indexed and full-table joins, and the fragment-patch sort.
  • DatasetDeltaBuilder::with_memory_pool covers the delta scans and the deleted-row anti join.

Example

let memory_pool = Arc::new(GreedyMemoryPool::new(512 * 1024 * 1024));
let (dataset, _stats) = MergeInsertBuilder::try_new(dataset, vec!["id".to_string()])?
    .when_matched(WhenMatched::UpdateAll)
    .memory_pool(memory_pool)
    .try_build()?
    .execute(new_data)
    .await?;

Not included

Scalar index training does not take the pool yet. That needs changes to CreateIndexBuilder and the optimize path, plus the trainer sorts in lance-index (btree, label_list, json, hilbert). It will come in a follow-up PR.

Part of #9799

🤖 Generated with Claude Code

Add `LanceExecutionOptions::memory_pool` and per-operation setters on
`Scanner`, `MergeInsertBuilder`, and `DatasetDeltaBuilder` so embedders
can control the DataFusion memory pool's sizing and sharing policy.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@github-actions github-actions Bot added the enhancement New feature or request label Oct 8, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant