Reimplement upsert as a single-pass, file-at-a-time operation with metadata-driven pruning

[kevinjqliu]

Summary

Transaction.upsert should be reimplemented as a single-pass, file-at-a-time operation. The current implementation encodes the source keyset as a BooleanExpression up front and hands that expression to every downstream stage. That expression is expensive to build, is evaluated many times, and does not reduce the data read below what per-file Iceberg metadata already allows. Every reported upsert bottleneck (slow filter construction, PyArrow segfaults on composite keys, full-table scans in delete(), double reads, unbounded memory, multiple snapshots) traces back to that design.

The proposal keeps the public API unchanged and treats pruning as an optional optimization layer derived directly from the source keyset's statistics, never from an n-term predicate.

Background: how upsert works today

1. create_match_filter turns every source key into a node of a Python expression tree. For composite keys it emits one And(...) disjunct per distinct key tuple.
2. That expression is bound, projected, walked by the manifest evaluator, walked again by the strict and inclusive metrics evaluators for every data file, converted to an Arrow expression, canonicalized by Acero, and evaluated per row group.
3. Matched rows are read with every column materialized, then compared to the source in a Python loop (slice + .as_py() per cell) because Arrow cannot join or compare nested types.
4. overwrite() calls delete(), which re-plans the same files against the same expression, loads each candidate file fully into memory with no row-filter pushdown, filters it, and rewrites it sequentially.
5. The result is up to three snapshots: DELETE, OVERWRITE, APPEND.

What the community has measured

Stage	Reported cost
Build the expression, 1M keys	10 s before any I/O
Acero canonicalization, 2 or 3 join columns	SIGSEGV at ~10k keys
Scan with the expression, 2M-row table	382 s
txn.delete() with the expression, 1M keys	1054 s of 1066 s total
Same table, successive upserts	12 min, 31 min, 51 min as file count grows
Python row diff, 20k rows x 200 columns	20 s (0.15 s vectorized)
Memory, 18.5 GB table on Lambda	> 10 GB, OOM

Why the fixes so far have not landed

Every proposed fix optimizes one stage of the pipeline while keeping the pipeline: shrink or reshape the expression, read fewer columns, speed up the Python diff, or drop the diff. None of them changes the three structural costs:

• the keyset is a predicate, so every stage scales with source rows x target files
• matched files are read twice, the second time without pushdown
• delete() rewrites whole files, one at a time, fully in memory

Key insight

The simplest correct upsert compares the source against every live data file in the table. The boolean expression was only ever a way to skip files, and skipping files can be done directly from the source keyset with Iceberg metadata: partition values and column bounds. Pruning is an optimization on top of the baseline, not the mechanism for matching.

Proposed design

Baseline algorithm

Inputs: source table cast to the current table schema, join columns, when_matched_update_all, when_not_matched_insert_all.

1. Plan candidate files. Start from every data file on the target branch with delete files attached (DataScan.plan_files with ALWAYS_TRUE). The pruning layer filters this list.
2. Read each candidate file once, projected to the current table schema via ArrowScan.to_table(tasks=[one file]), the same way delete() reads files today. One worker holds one file, which also sidesteps the eager prefetch behavior of the batch reader.
3. Match in Arrow, not through a predicate. Inner-join the file's key columns against the source key columns with row-index markers. Join columns must be primitive, non-null, non-float types; reject anything else up front with a clear error. Nested types are fine in non-key columns because they are never joined.
4. Diff vectorized. Per non-key column, not_equal with Kleene null handling, recursing into struct fields. Python fallback only for list and map columns.
5. Decide per file. If no matched row changed, the file is untouched. Otherwise the replacement file is the target rows minus the changed matched rows, written via _dataframe_to_data_files. Record matched source indices in a per-source-row counter.
6. After all files. Source rows with a nonzero counter and a detected change form the update set. Source rows with a zero counter form the insert set. A counter above one means the target has duplicate keys and the upsert should abort. Write update and insert rows together through the normal partitioned writer once.
7. Commit one OVERWRITE snapshot using the existing _OverwriteFiles producer: delete_data_file for each replaced original, append_data_file for each rewritten and new file. Existing replaced-file validation covers the commit-retry path.

When when_matched_update_all=False, step 2 only needs the join columns projected and step 5 never rewrites, so the operation degenerates to a pure APPEND of unmatched rows.

Pruning as an optimization layer

All three levels are computed once from the source keyset and never build an expression.

Level	Signal	How to skip	Cost per unit
Manifest	partition field summaries	For each partition field whose source column is a join column, apply the partition transform to the source values to get a set of partition values. Skip the manifest if its summary range for that field is disjoint from the set's min and max.	O(partition fields)
Data file, partition	DataFile.partition	Exact membership check of the file's partition tuple against the transformed set. Compare in partition space, not by projecting a row-space predicate.	O(partition fields)
Data file, column bounds	lower_bounds / upper_bounds for join column field IDs	Decode the bounds, binary-search the sorted source values for that column. Skip if no value lands in the range. Composite keys are checked per column independently, which is conservative but never explodes.	O(join cols x log n)

Two rules keep this correct:

• Partition pruning may only use partition fields derived from join columns. A source row's own partition value says nothing about where the old row lives; using it silently turns an update into a duplicate insert when a key moves partitions.
• Pruning must be a pure filter over the candidate list. Turning it off yields the baseline, and tests should run both ways.

The column-bounds check is what _InclusiveMetricsEvaluator computes for an In predicate, minus walking n nodes per file.

Expected cost

Stage	Today	Proposed
Before I/O	O(n) Python nodes, then bind and project	sort join columns, transform partition values
File selection	walk n-node tree per manifest and per file, three evaluators	O(cols x log n) per file
Row filtering	Arrow canonicalizes n disjuncts, segfaults at ~10k composite keys	hash join per file
Data reads	every candidate read twice, second time without pushdown	once
Rewrite	sequential, whole file in memory	per file, parallel by worker count
Snapshots	DELETE, OVERWRITE, APPEND	one OVERWRITE

What this fixes

• Keyset-as-predicate construction and evaluation are gone.
• Double scan and full-file loads in delete() are gone.
• Python cell loops are gone except for list and map columns.
• Memory is bounded by workers x one file + source.
• Snapshot count drops to one, which restores atomicity from the reader's perspective and works with catalogs that reject multi-snapshot commits.
• Schema evolution works because every file is projected to the current schema and the source is cast once.
• Cross-file duplicate detection in the target becomes possible via the per-source-row counter.

What this does not fix

• A key whose partition column is not a join column still requires scanning every file that column-bounds pruning cannot exclude. Only equality deletes remove that limit. The design leaves a seam: step 5 could emit an equality delete file instead of a rewrite in the future.
• Nested or null join keys are rejected, not supported.
• Source keys must fit in memory, as today.

Alternatives considered

• Patch in place (vectorize the diff, fix schema-evolution projection, regroup the composite-key filter). Cheap and worth doing regardless, but leaves the source x files scaling, the double read, and the whole-file loads untouched.
• Engine delegation (DataFusion / iceberg-rust). Right long-term direction, not available now. The per-file step in this design (source + target file in, rewritten rows + matched indices out) is the natural pluggable boundary and keeps all PyArrow calls in io/pyarrow.py.
• Equality-delete write path. Solves partition migration and rewrite cost but adds compaction and read-side obligations. Defer, keep the seam.

Open decisions

• One OVERWRITE snapshot vs. the current three. Separate snapshots were argued to help other clients detect conflicts. A single snapshot fixes atomicity and multi-snapshot catalog rejections, and the manifest entries still record every deleted and added file. Recommendation: single snapshot.
• Null join keys. Reject with a clear error, or silently never match. Recommendation: reject, since silent non-matching creates duplicates on the next run.
• Fate of pyiceberg.table.upsert_util. It is public by accident. Dissolve it into io/pyarrow.py with deprecation shims.

Suggested PR sequence

1. Keyset pruning helpers (column-bounds and partition-space) with unit tests covering every partition transform.
2. Per-file match-and-diff function in io/pyarrow.py, with the vectorized compare and the join-column type guard.
3. New Transaction.upsert body on the _OverwriteFiles producer, with pruning switchable off in tests. Fold in existing regression tests for schema evolution, composite-key segfaults, wide tables, and partition transforms.
4. Deprecate upsert_util.

[kevinjqliu]

Proposal: single-pass upsert with metadata-driven pruning

Goal

Replace the body of Transaction.upsert with an implementation that never encodes the source keyset as a BooleanExpression. The baseline compares the source against every live data file. Pruning is a separate, optional layer that skips files using Iceberg metadata derived from the keyset. The public signature and UpsertResult do not change.

Algorithm

upsert(df, join_cols, when_matched_update_all, when_not_matched_insert_all)

  source      = cast df to table schema
  validate    join_cols are primitive, non-null, non-float
  keys        = sorted source values per join column
  parts       = transformed partition values, for partition fields sourced from join_cols
  matched     = int32 array of zeros, len(source)

  candidates  = plan_files(ALWAYS_TRUE) on the target branch
  candidates  = prune(candidates, keys, parts)          # optional layer

  for each file in candidates (parallel, one file per worker):
      target   = read file once, projected to current table schema
      pairs    = join(target[join_cols], source[join_cols], inner, with row-index markers)
      matched[pairs.source_idx] += 1
      if when_matched_update_all:
          changed  = vectorized diff of target rows vs source rows for pairs
          if any(changed):
              keep     = target minus changed target rows
              replaced.append((file, write(keep)))

  if any(matched > 1): raise duplicate keys in target
  updates = source rows where matched == 1 and changed          # only if when_matched_update_all
  inserts = source rows where matched == 0                      # only if when_not_matched_insert_all
  new_files = write(concat(updates, inserts))                   # normal partitioned writer, once

  with update_snapshot().overwrite() as snap:
      for original, rewritten in replaced:
          snap.delete_data_file(original)
          snap.append_data_file(*rewritten)
      snap.append_data_file(*new_files)

  return UpsertResult(rows_updated=len(updates), rows_inserted=len(inserts))

Notes on each step

• Read once. Each candidate file is read a single time with ArrowScan.to_table(tasks=[file]), projected to the current table schema. Positional delete files attached by plan_files are applied during the read. This replaces the current pattern of one scan to find matches and a second full read inside delete().
• Match by join, not by predicate. The file's key columns are inner-joined against the source key columns with row-index marker columns on both sides. Arrow's hash join handles primitive key types; nested types are only ever in non-key columns, so they never reach the join.
• Diff vectorized. For each non-key column, not_equal combined with is_null on either side using Kleene logic. Structs recurse into fields by name. Lists and maps fall back to a per-row Python compare. A row is changed if any column is changed.
• Rewrite only what changed. A file with matched rows but no changes is not touched. A file with changed rows is replaced by its unchanged rows. Updated source rows go through the normal partitioned writer with the inserts, so a row whose partition value changed lands in the right partition.
• Duplicate detection across files. The matched counter over source rows detects a key that appears in two target files or twice in one file, which the batch-based implementation cannot do.
• One snapshot. The existing _OverwriteFiles producer takes delete_data_file and append_data_file calls and emits a single OVERWRITE snapshot. Replaced-file validation on commit retry already exists.
• when_matched_update_all=False. Project only join columns in the read, skip the diff, never rewrite. The operation becomes an APPEND of unmatched rows.
• when_not_matched_insert_all=False. Skip the insert set.

Pruning layer

Computed once from keys and parts. Applied as a pure filter over the candidate list, so disabling it yields the baseline.

Level	Input	Skip when	Cost
Manifest	partition field summaries	for a partition field sourced from a join column, the summary [lower, upper] is disjoint from [min(parts), max(parts)]	O(partition fields)
File, partition	DataFile.partition	the file's value for a join-sourced partition field is not in parts	O(partition fields)
File, bounds	lower_bounds / upper_bounds for join column field IDs	for any join column, binary search of keys[col] finds no value within [lower, upper]	O(join cols x log n)

Rules:

• Only partition fields whose source column is a join column may prune. A source row's own partition value does not say where the old row lives.
• Partition comparisons happen in partition space: transform the source values with the field's transform and compare to DataFile.partition directly. Never project a row-space predicate.
• Bounds are decoded with from_bytes on the field's primitive type. A file missing bounds for a join column is never pruned.
• Composite keys are checked per column independently. This is conservative (a file may survive pruning when no full tuple matches) but never grows with the number of keys.

Validation up front

• join_cols resolved from identifier fields when not given, as today.
• Every join column must be a primitive type that is not float or double. Reject nested, dictionary-encoded, extension, and null types with a ValueError.
• Null values in join columns raise a ValueError. Silent non-matching would insert a duplicate on the next run.
• Duplicate keys in the source raise, as today.
• Source schema is checked against the table schema, then cast once.

Memory and parallelism

• Peak memory is workers x (one target file + its rewritten rows) + source. The source and its sorted keys must fit in memory, as today.
• File processing is embarrassingly parallel. Use the existing executor with the configured worker count. Each task returns (file, rewritten data files, matched source indices, changed source indices) and the driver merges.
• Written files are only referenced once the snapshot commits. A failure mid-run leaves orphan files, same as overwrite today.

Code placement

• pyiceberg/io/pyarrow.py: upsert_file(source, target, join_cols, detect_changes) -> (kept_rows, matched_idx, changed_idx) and the vectorized diff. This is the only place that touches PyArrow compute, and it is the boundary a future execution engine would replace.
• pyiceberg/table/upsert_util.py (or a private module): keyset pruning helpers and the join-column type guard. No PyArrow compute beyond sorting and transform application.
• pyiceberg/table/__init__.py: the new Transaction.upsert body, orchestration only.
• Existing create_match_filter, has_duplicate_rows, get_rows_to_update remain as deprecated shims for one release.

Testing

• Baseline vs pruned: every test runs with pruning disabled and enabled and asserts identical results and identical snapshots.
• Partition transforms: identity, bucket, truncate, year, month, day, hour, and a mixed spec, each with a key that moves partitions and a key that does not.
• Composite keys at 100k rows with two and three join columns, run in a subprocess to catch a segfault.
• Wide table: 200 columns, 20k matched rows, unchanged and changed variants.
• Nested non-key columns: struct, list, map, and struct-of-list.
• Schema evolution: add_column, union_by_name, rename of a join column, on main and on a diverged branch.
• Duplicate keys in the target across two files.
• Both flag combinations that disable a side, and the no-op case that must not create a snapshot.
• Concurrent writer between plan and commit, to exercise the retry path.

Rollout

1. Pruning helpers and type guard, with unit tests.
2. upsert_file and the vectorized diff in io/pyarrow.py, with unit tests.
3. New Transaction.upsert body behind the existing signature, with the integration tests above.
4. Deprecation shims for the old helpers.

Out of scope, kept as seams

• Equality-delete output in place of a rewrite, for keys whose partition column is not a join column.
• A pluggable execution engine behind upsert_file.
• Streaming the source as a RecordBatchReader.