lance-format / lance-format/lance

Differentiate streaming and materialized inputs for add and merge insert

Open
#4,583 0 comments 1 reaction 1 assignee View on GitHub

@jbapple is already working on this.

Since Aug 28, 2025.

enhancement
Dominant language
Rust
Stars
7.1k
Forks
852
Avg merge
3d 18h
Merged PRs (30d)
272

Description

Most of our APIs assume inputs are a stream. This is nice in that it supports larger-than-memory writes. However, if data is fully materialized — or backed by a re-readable source like a file — we can often do things more optimally. To give a few examples:

  1. If we have 2 million rows in memory we want to insert, we can write multiple data files in parallel. Currently we write them sequentially.
  2. To support retries for write operations, we buffer data on disk. This is wasteful when the data is already re-readable (in memory, or already a file on disk).
  3. For merge_insert, we can compute basic statistics like num_rows and num_bytes, which DataFusion can use to optimize the join order. Currently the source side carries no statistics, so the optimizer can't choose the build side well.

This also supports downstream use cases: https://github.com/lancedb/lancedb/issues/2602

Design

The original proposal was a two-variant InputData enum (Stream / Materialized). The problem: it has no clean home for a re-readable source that isn't in memory — e.g. a file or set of files on disk. That source is both re-scannable and can describe its size, but it's neither a one-shot Stream nor Materialized. lancedb#2602 reached the same conclusion independently, landing on three cases: in-memory, one-shot stream, and "stream factory" (something re-startable).

DataFusion's TableProvider is the natural type for the re-startable case, and it strictly dominates a bare stream-factory closure because it carries statistics too:

  • scan() is callable repeatedly → replay / retry without spilling.
  • The plan from scan() reports statistics. MemTable in particular fills both num_rows and total_byte_size exactly (compute_record_batch_statistics), which is exactly use case 3.
  • MemTable parallelizes over its partitions (output_partitioning() = UnknownPartitioning(partitions.len())), giving the structure for use case 1 for free.
  • It composes directly with the DataFusion plan merge_insert already builds.

So rather than make TableProvider the single public input type (overkill for the simple append path, and dishonest for one-shot streams, which can only become a provider that errors on the second scan), we keep an ergonomic surface and make TableProvider the canonical internal representation.

Public API
impl MergeInsertBuilder {
    // Ergonomic wrappers
    async fn execute(self, stream: SendableRecordBatchStream) -> Result<...> { ... }
    async fn execute_batches(self, batches: Vec<RecordBatch>) -> Result<...> { ... }
    // Canonical entry point
    async fn execute_provider(self, provider: Arc<dyn TableProvider>) -> Result<...> { ... }
}

execute and execute_batches are thin wrappers over execute_provider:

  • execute_batches(batches) → build a multi-partition MemTable (use the balanced repartition helper / split by row budget — a single inner Vec is one partition and gets no parallelism) → execute_provider. In-memory replay, exact stats, never spills.
  • execute(stream) → wrap the one-shot stream in a provider → execute_provider. This is the only genuinely non-replayable case (see Spilling below).
  • execute_provider(provider) → re-scans the provider directly. Files (Lance, Parquet, a re-readable pa.dataset.Dataset/Scanner) re-read from their durable copy with no spill.

The same trio applies to add/InsertBuilder for consistency (it doesn't replay/spill today, but the input shape should match).

Spilling & retries

The commit-conflict retry loop re-runs the whole merge against the latest table version on each attempt, so every attempt must read the source again. Today (merge_insert only) new_source_iter manufactures replayability from a one-shot stream: 1 batch → clone in memory; otherwise → SpillStreamIter, which drains the source to a temp file (create_replay_spill, ~100 MB in memory then disk) and hands out a fresh replay stream per attempt.

That whole apparatus is a hand-rolled "make a one-shot stream re-scannable" — which is exactly what TableProvider::scan() already is. Under the new design:

  • The retry executor holds Arc<dyn TableProvider> and calls scan() fresh per attempt. new_source_iter is removed as the universal replay layer.
  • execute_batchesMemTable: re-scannable in memory, no spill (subsumes & generalizes today's 1-batch special case).
  • execute_provider(file) → re-read the file each attempt; the file is the durable copy, zero extra disk.
  • execute(stream) → the only case that needs a spill.

So spilling goes from "the default for any multi-batch source" to one opt-in adapter for the one-shot stream case:

  • Repackage SpillStreamIter as a SpillingTableProvider that wraps a one-shot stream — first scan() tees to the spill file, later scans replay from it. Reuse create_replay_spill (keep mem-then-disk).
  • Derive a replayable flag from how the source was built (stream → false, others → true) so the retry loop chooses re-scan vs. fail-fast and never scans a one-shot provider twice.
  • For execute(stream) with retries: preserve current behavior (wrap in SpillingTableProvider) to avoid breaking the public API, but add an explicit opt-out (e.g. .spill_for_retry(false)) for callers who'd rather fail fast with "can't retry with stream input" — matching lancedb#2602's principle of not silently buffering a stream to disk.
  • SpillingTableProvider reports Absent statistics (it can't know num_rows before draining); the real stats wins come from MemTable/file providers.
Bindings

Convert each input to the right provider in the thin binding layer:

Materialized → MemTable:

  • pa.Table
  • pd.DataFrame
  • pa.RecordBatch

Re-scannable → custom TableProvider:

  • pa.dataset.Dataset (re-invoke the scanner per scan(); push down projection/filter/limit; expose count_rows as stats)
  • pa.dataset.Scanner
  • stream-factory callable, e.g. Callable[[], pa.RecordBatchReader]

One-shot stream → spilling/one-shot provider:

  • pa.RecordBatchReader
  • raw C-stream / socket sources

TODO

  • Define execute_provider(Arc<dyn TableProvider>) as the canonical merge_insert entry; make execute(stream) and execute_batches(batches) thin wrappers.
  • execute_batches: build a multi-partition MemTable so it gets exact stats (use case 3) and parallel partitions (toward use case 1).
  • Replace new_source_iter/iterator-of-streams in the retry path with re-scanning the TableProvider per attempt; carry a replayable flag.
  • Repackage SpillStreamIter as SpillingTableProvider for the one-shot stream case; add .spill_for_retry(false) opt-out. (use case 2)
  • Confirm DataFusion's join-selection picks the build side from source/target statistics now that the source provider reports them. (use case 3)
  • Bindings: convert Python (and TS/Java) inputs to the appropriate provider per the mapping above.
  • Fan the fragment writer out over the provider's partitions so materialized inserts write data files in parallel. (use case 1 — the substantive remaining piece; do_write_fragments consumes a single stream today)

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.