[Bug]: Python 1.5.1 MERGE regression: small source batch uses substantially more RSS than 1.5.0

[bnewman-tech]

What happened?

After upgrading the Python package from deltalake==1.5.0 to deltalake==1.5.1, a Delta MERGE workload started using substantially more memory. In the original environment this eventually resulted in the process being killed by the OS.

I reduced this to a local synthetic repro with no production data. The workload is:

• Local Delta table
• String key column similar to a typical row key
• Large target table
• Small source update batch
• Merge predicate: source.ROWKEY = target.ROWKEY
• Unique source keys
• streamed_exec=True

Smaller local repro:

deltalake	pyarrow	target rows	source rows	max RSS increase	end RSS
1.5.0	24.0.0	300k	5k	+126 MB	303 MB
1.5.1	24.0.0	300k	5k	+651 MB	825 MB

Larger local repro:

deltalake	pyarrow	target rows	source rows	max RSS increase	end RSS
1.5.0	24.0.0	1M	10k	+294 MB	537 MB
1.5.1	24.0.0	1M	10k	+1,225 MB	1,442 MB

The compare from python-v1.5.0 to python-v1.5.1 is:

python-v1.5.0...python-v1.5.1

The most suspicious change I found is:

#4333

That PR says duplicate-match validation adds marker columns:

• __delta_rs_match_cardinality_class
• __delta_rs_target_row_index

It also says this validation state is kept in memory per validation stream:

target_row_state: HashMap<u64, (i32, i32)>

There is a follow-up merge validation semantics change here:

#4359

The release also includes the DataFusion/Arrow/object_store upgrade:

#4346

The memory jump follows deltalake==1.5.1 in the local repro matrix.

Expected behavior

I expected deltalake==1.5.1 merge memory usage to remain in the same general range as deltalake==1.5.0 for the same workload, especially with streamed_exec=True.

For this repro, deltalake==1.5.0 uses roughly +126 MB to +294 MB max RSS depending on table size. I would expect deltalake==1.5.1 to stay close to that range, or at least not require several times more RSS for the same small source batch and merge predicate.

If the additional memory is expected because of duplicate-match validation, I would expect a way to bound/spill that state or a documented mitigation for merge workloads where source keys are already unique.

Operating System

macOS

Binding

Python

Bindings Version

1.5.1

Steps to reproduce

1. Save this script as delta_merge_memory_repro.py.

from __future__ import annotations

import argparse
import gc
import os
import shutil
import tempfile
import threading
import time
from pathlib import Path

import psutil
import pyarrow as pa
from deltalake import DeltaTable, write_deltalake


def make_target_table(row_count: int) -> pa.Table:
    return pa.table(
        {
            "ROWKEY": [f"key-{i:010d}" for i in range(row_count)],
            "rowversion": list(range(row_count)),
            "amount": [i % 1000 for i in range(row_count)],
            "payload": [f"payload-{i % 1000:04d}" for i in range(row_count)],
        }
    )


def make_source_table(*, target_rows: int, source_rows: int) -> pa.Table:
    start = max(0, target_rows - source_rows)
    keys = [f"key-{i:010d}" for i in range(start, target_rows)]
    return pa.table(
        {
            "ROWKEY": keys,
            "rowversion": list(range(target_rows, target_rows + source_rows)),
            "amount": [10_000 + (i % 1000) for i in range(source_rows)],
            "payload": [f"updated-{i % 1000:04d}" for i in range(source_rows)],
        }
    )


def sample_peak_rss(*, stop: threading.Event, samples: list[float]) -> None:
    process = psutil.Process(os.getpid())
    while not stop.is_set():
        samples.append(process.memory_info().rss / 1024 / 1024)
        time.sleep(0.02)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--target-rows", type=int, default=300_000)
    parser.add_argument("--source-rows", type=int, default=5_000)
    parser.add_argument("--target-file-size", type=int, default=800_000)
    args = parser.parse_args()

    table_path = Path(tempfile.mkdtemp(prefix="delta-merge-memory-repro-"))
    process = psutil.Process(os.getpid())

    try:
        write_deltalake(
            str(table_path),
            make_target_table(args.target_rows),
            mode="overwrite",
            target_file_size=args.target_file_size,
        )

        gc.collect()
        source = make_source_table(
            target_rows=args.target_rows,
            source_rows=args.source_rows,
        )

        before_rss_mb = process.memory_info().rss / 1024 / 1024
        samples: list[float] = []
        stop = threading.Event()
        sampler = threading.Thread(
            target=sample_peak_rss,
            kwargs={"stop": stop, "samples": samples},
            daemon=True,
        )

        sampler.start()
        started = time.perf_counter()
        metrics = (
            DeltaTable(str(table_path))
            .merge(
                source=source,
                predicate="source.ROWKEY = target.ROWKEY",
                source_alias="source",
                target_alias="target",
                streamed_exec=True,
            )
            .when_matched_update(
                {
                    "rowversion": "source.rowversion",
                    "amount": "source.amount",
                    "payload": "source.payload",
                }
            )
            .when_not_matched_insert(
                {
                    "ROWKEY": "source.ROWKEY",
                    "rowversion": "source.rowversion",
                    "amount": "source.amount",
                    "payload": "source.payload",
                }
            )
            .execute()
        )
        duration_seconds = time.perf_counter() - started

        stop.set()
        sampler.join(timeout=1)
        after_rss_mb = process.memory_info().rss / 1024 / 1024
        peak_rss_mb = max(samples or [before_rss_mb])

        print(
            {
                "target_rows": args.target_rows,
                "source_rows": args.source_rows,
                "before_rss_mb": round(before_rss_mb, 1),
                "peak_rss_mb": round(peak_rss_mb, 1),
                "after_rss_mb": round(after_rss_mb, 1),
                "peak_rss_delta_mb": round(peak_rss_mb - before_rss_mb, 1),
                "duration_seconds": round(duration_seconds, 3),
                "metrics": metrics,
            }
        )
    finally:
        shutil.rmtree(table_path, ignore_errors=True)


if __name__ == "__main__":
    main()

2. Run the small comparison.

uv run --no-project --with deltalake==1.5.0 --with pyarrow==24.0.0 --with psutil python delta_merge_memory_repro.py --target-rows 300000 --source-rows 5000 --target-file-size 800000
uv run --no-project --with deltalake==1.5.1 --with pyarrow==24.0.0 --with psutil python delta_merge_memory_repro.py --target-rows 300000 --source-rows 5000 --target-file-size 800000

3. Optionally run the larger confirmation.

uv run --no-project --with deltalake==1.5.0 --with pyarrow==24.0.0 --with psutil python delta_merge_memory_repro.py --target-rows 1000000 --source-rows 10000 --target-file-size 1200000
uv run --no-project --with deltalake==1.5.1 --with pyarrow==24.0.0 --with psutil python delta_merge_memory_repro.py --target-rows 1000000 --source-rows 10000 --target-file-size 1200000

Relevant logs