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 |", + "+------------+--------------+----------------+---------------+", + ], +)