From 0ee2907e383b1649f6c67ded013cbcefa5a72544 Mon Sep 17 00:00:00 2001 From: Sergei Turukin Date: Wed, 15 Apr 2026 16:07:06 -0500 Subject: [PATCH] fix(merge): preserve target rows when MERGE batch contains only target The MergeCOWFilterStream "no matches in this batch" fast path short-circuited on `matching_data_and_manifest_files.is_empty()` without checking the cumulative `all_matching_data_files` set. If a target file had been seen as matching in an earlier batch and a later batch contained only target rows for that file, the rows in the later batch were silently dropped. The downstream COW commit then overwrote the original file with the partial result, permanently losing the unmatched target rows whose batch hit the dead path. The fix tightens the guard to also require `all_matching_data_files` to be empty before taking the fast path. When a batch belongs to a file already in the overwrite set, it falls through to the main filter path which correctly emits target rows via `file_predicate OR source_exists`. Adds three unit tests against MergeCOWFilterStream covering the matching-then-target patterns, plus a SQL snapshot test that exercises the same shape end-to-end. Fixes #128 Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/datafusion/physical_plan/merge.rs | 28 ++++++++++++++++++- .../executor/src/tests/sql/ddl/merge_into.rs | 22 +++++++++++++++ ...mixed_unsorted_multi_row_no_data_loss.snap | 14 ++++++++++ 3 files changed, 63 insertions(+), 1 deletion(-) create mode 100644 crates/executor/src/tests/sql/ddl/snapshots/merge_into/query_merge_into_mixed_unsorted_multi_row_no_data_loss.snap diff --git a/crates/executor/src/datafusion/physical_plan/merge.rs b/crates/executor/src/datafusion/physical_plan/merge.rs index 22b63bc7..17682204 100644 --- a/crates/executor/src/datafusion/physical_plan/merge.rs +++ b/crates/executor/src/datafusion/physical_plan/merge.rs @@ -655,7 +655,13 @@ impl Stream for MergeCOWFilterStream { .push(filtered_batch); } - if matching_data_and_manifest_files.is_empty() { + // Only take the fast paths if the current batch references no target file + // that will be (or has been) overwritten. Otherwise the full filter path + // below is required so target rows belonging to `all_matching_data_files` + // are re-emitted into the rewritten data file. + if matching_data_and_manifest_files.is_empty() + && all_matching_data_files.is_empty() + { // Return early if all rows only come from source if matching_data_file_array.len() == source_exists_array.len() { return Poll::Ready(Some(Ok(batch))); @@ -1210,4 +1216,24 @@ mod tests { &[(0, 2), (0, 1), (0, 6), (0, 5)], 60 ); + // Regression test for https://github.com/Embucket/embucket/issues/128 + // + // If a target file has been seen as "matching" in an earlier batch and a subsequent + // batch contains only target rows (no `__source_exists` = true rows) for that same + // file, the rows in the later batch must still be passed through the filter so they + // land in the rewritten data file. Previously the "no matches, no source" fast path + // dropped them, causing silent data loss during `MERGE INTO` on unsorted inputs. + test_merge_cow_filter_stream!(matching_then_target, &[(0, 4), (0, 1)], 20); + test_merge_cow_filter_stream!( + matching_then_target_then_matching, + &[(0, 4), (0, 1), (0, 4)], + 30 + ); + // Mixed scenario: several target-only batches arriving AFTER the target file has + // been matched. + test_merge_cow_filter_stream!( + matching_then_multiple_target_batches, + &[(0, 4), (0, 1), (0, 1), (0, 1)], + 40 + ); } diff --git a/crates/executor/src/tests/sql/ddl/merge_into.rs b/crates/executor/src/tests/sql/ddl/merge_into.rs index e4acaeec..a67be45a 100644 --- a/crates/executor/src/tests/sql/ddl/merge_into.rs +++ b/crates/executor/src/tests/sql/ddl/merge_into.rs @@ -299,3 +299,25 @@ test_query!( ], snapshot_path = "merge_into" ); + +// Regression test for https://github.com/Embucket/embucket/issues/128. +// +// Target is one data file with many rows; source is a mix of updates (matches) and +// inserts (no match), and the target rows of the join land in the filter stream in +// batches where some contain source_exists=true rows and some only contain target +// rows. Previously the "no matches, no source" fast path would silently drop the +// target-only batches for a file that had already been marked as matching in an +// earlier batch, causing the final row count to be less than the expected +// (target_rows + new_source_rows). This test asserts that no target row is lost. +test_query!( + merge_into_mixed_unsorted_multi_row_no_data_loss, + "SELECT COUNT(*) as total_rows, COUNT(CASE WHEN description = 'updated row' THEN 1 END) as updated_rows, COUNT(CASE WHEN description = 'original row' THEN 1 END) as preserved_rows, COUNT(CASE WHEN description = 'new row' THEN 1 END) as inserted_rows FROM embucket.public.merge_target", + setup_queries = [ + "CREATE TABLE embucket.public.merge_target (id INTEGER, description VARCHAR)", + "CREATE TABLE embucket.public.merge_source (id INTEGER, description VARCHAR)", + "INSERT INTO embucket.public.merge_target VALUES (1, 'original row'), (2, 'original row'), (3, 'original row'), (4, 'original row'), (5, 'original row'), (6, 'original row'), (7, 'original row'), (8, 'original row'), (9, 'original row'), (10, 'original row')", + "INSERT INTO embucket.public.merge_source VALUES (3, 'updated row'), (7, 'updated row'), (11, 'new row'), (12, 'new row')", + "MERGE INTO merge_target t USING merge_source s ON t.id = s.id WHEN MATCHED THEN UPDATE SET t.description = s.description WHEN NOT MATCHED THEN INSERT (id, description) VALUES (s.id, s.description)", + ], + snapshot_path = "merge_into" +); diff --git a/crates/executor/src/tests/sql/ddl/snapshots/merge_into/query_merge_into_mixed_unsorted_multi_row_no_data_loss.snap b/crates/executor/src/tests/sql/ddl/snapshots/merge_into/query_merge_into_mixed_unsorted_multi_row_no_data_loss.snap new file mode 100644 index 00000000..3dd7a8b1 --- /dev/null +++ b/crates/executor/src/tests/sql/ddl/snapshots/merge_into/query_merge_into_mixed_unsorted_multi_row_no_data_loss.snap @@ -0,0 +1,14 @@ +--- +source: crates/executor/src/tests/sql/ddl/merge_into.rs +description: "\"SELECT COUNT(*) as total_rows, COUNT(CASE WHEN description = 'updated row' THEN 1 END) as updated_rows, COUNT(CASE WHEN description = 'original row' THEN 1 END) as preserved_rows, COUNT(CASE WHEN description = 'new row' THEN 1 END) as inserted_rows FROM embucket.public.merge_target\"" +info: "Setup queries: CREATE TABLE embucket.public.merge_target (id INTEGER, description VARCHAR); CREATE TABLE embucket.public.merge_source (id INTEGER, description VARCHAR); INSERT INTO embucket.public.merge_target VALUES (1, 'original row'), (2, 'original row'), (3, 'original row'), (4, 'original row'), (5, 'original row'), (6, 'original row'), (7, 'original row'), (8, 'original row'), (9, 'original row'), (10, 'original row'); INSERT INTO embucket.public.merge_source VALUES (3, 'updated row'), (7, 'updated row'), (11, 'new row'), (12, 'new row'); MERGE INTO merge_target t USING merge_source s ON t.id = s.id WHEN MATCHED THEN UPDATE SET t.description = s.description WHEN NOT MATCHED THEN INSERT (id, description) VALUES (s.id, s.description)" +--- +Ok( + [ + "+------------+--------------+----------------+---------------+", + "| total_rows | updated_rows | preserved_rows | inserted_rows |", + "+------------+--------------+----------------+---------------+", + "| 12 | 2 | 8 | 2 |", + "+------------+--------------+----------------+---------------+", + ], +)