diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 11a5619a4d2..36098034f95 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -390,6 +390,7 @@ jobs: org.apache.comet.exec.CometExecSuite org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite + org.apache.comet.exec.CometMergeRowsSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 7b7fae8fbea..69e72474580 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -163,6 +163,7 @@ jobs: org.apache.comet.exec.CometExecSuite org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite + org.apache.comet.exec.CometMergeRowsSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index e039df23b76..13ce385c01b 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -85,6 +85,31 @@ runs natively; it is controlled by `spark.comet.exec.windowGroupLimit.enabled` ( - Signed-zero ordering (`-0.0` vs `+0.0`) diverges from Spark's `RankLimitIterator`; see [floating-point ordering](./floating-point.md#ordering-signed-zero-00-vs-00). +## MERGE INTO (MergeRowsExec) + +Comet can run `MergeRowsExec` (Spark's row-level `MERGE INTO` dispatch operator, Spark 3.5+) +natively, but it is disabled by default. Enable it with `spark.comet.exec.mergeRows.enabled=true`. + +**Missing per-clause row metrics:** Spark 4.x's `MergeRowsExec` exposes eight metrics -- +`numTargetRowsCopied`, `numTargetRowsInserted`, `numTargetRowsUpdated`, `numTargetRowsDeleted`, +`numTargetRowsMatchedUpdated`, `numTargetRowsMatchedDeleted`, `numTargetRowsNotMatchedBySourceUpdated`, +and `numTargetRowsNotMatchedBySourceDeleted` -- breaking down how many rows each `MERGE` clause +touched. Comet's native operator does not expose these; it only reports the generic `output_rows`, +`output_batches`, and `elapsed_compute` every native operator reports. EXPLAIN ANALYZE and the +Spark UI will not show a rows-inserted/updated/deleted breakdown for a native `MERGE`. (Spark +3.5.x's own `MergeRowsExec` does not have these metrics either -- they were added alongside a +`Context` field Spark only attaches to `MERGE` clauses starting in 4.x.) + +**Cardinality-violation error may differ from Spark's on rare inputs:** Spark validates cardinality +(rejecting an `ON` condition that matches one target row to multiple source rows, +`MERGE_CARDINALITY_VIOLATION`) row-at-a-time, interleaved with applying each `MATCHED` clause, so +whichever failure a given row hits first is the error Spark raises. Comet's native operator is +vectorized: it validates cardinality for an entire input batch before evaluating any clause. If a +single batch contains both a cardinality violation and an unrelated clause-evaluation error (for +example an ANSI divide-by-zero) on different rows, Comet may raise a different error than Spark +would for the same input, depending on which row each engine reaches first. The query fails either +way; only the specific error differs. + ## Round-Robin Partitioning Comet's native shuffle implementation of round-robin partitioning (`df.repartition(n)`) is not compatible with diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 3a18d86606d..eb2335a4b0e 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -116,9 +116,10 @@ omitted from the tables below and may be reconsidered based on demand: ## Writes -| Operator | Status | Notes | -| ------------------------ | ------ | ----------------------------------------------------------------- | -| `DataWritingCommandExec` | ⚠️ | Experimental native Parquet writes, disabled by default (opt-in). | +| Operator | Status | Notes | +| ------------------------ | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `DataWritingCommandExec` | ⚠️ | Experimental native Parquet writes, disabled by default (opt-in). | +| `MergeRowsExec` | ⚠️ | Row-level `MERGE INTO` dispatch (Spark 3.5+). Disabled by default; opt in with `spark.comet.exec.mergeRows.enabled=true`. See [Operator Compatibility](compatibility/operators.md). | ## Python and UDF diff --git a/native/common/src/error.rs b/native/common/src/error.rs index 81d095658e5..fb701616980 100644 --- a/native/common/src/error.rs +++ b/native/common/src/error.rs @@ -184,6 +184,12 @@ pub enum SparkError { #[error("[SCALAR_SUBQUERY_TOO_MANY_ROWS] Scalar subquery returned more than one row.")] ScalarSubqueryTooManyRows, + /// Mirrors Spark's `QueryExecutionErrors.mergeCardinalityViolationError()`, raised by + /// `MergeRowsExec.BitmapCardinalityValidator` when a MERGE's ON condition matches a single + /// target row against more than one source row. + #[error("[MERGE_CARDINALITY_VIOLATION] The ON search condition of the MERGE statement matched a single row from the target table with multiple rows of the source table. This could result in the target row being operated on more than once with an update or delete operation and is not allowed.")] + MergeCardinalityViolation, + #[error("{message}")] FileNotFound { message: String }, @@ -303,6 +309,7 @@ impl SparkError { SparkError::InvalidRegexGroupIndex { .. } => "InvalidRegexGroupIndex", SparkError::DatatypeCannotOrder { .. } => "DatatypeCannotOrder", SparkError::ScalarSubqueryTooManyRows => "ScalarSubqueryTooManyRows", + SparkError::MergeCardinalityViolation => "MergeCardinalityViolation", SparkError::FileNotFound { .. } => "FileNotFound", SparkError::DuplicateFieldCaseInsensitive { .. } => "DuplicateFieldCaseInsensitive", SparkError::DuplicateFieldByFieldId { .. } => "DuplicateFieldByFieldId", @@ -618,7 +625,8 @@ impl SparkError { | SparkError::UnexpectedPositiveValue { .. } | SparkError::UnexpectedNegativeValue { .. } | SparkError::InvalidRegexGroupIndex { .. } - | SparkError::ScalarSubqueryTooManyRows => "org/apache/spark/SparkRuntimeException", + | SparkError::ScalarSubqueryTooManyRows + | SparkError::MergeCardinalityViolation => "org/apache/spark/SparkRuntimeException", // DateTimeException SparkError::InvalidInputInCastToDatetime { .. } @@ -736,6 +744,9 @@ impl SparkError { // Subquery errors SparkError::ScalarSubqueryTooManyRows => Some("SCALAR_SUBQUERY_TOO_MANY_ROWS"), + // MERGE INTO errors + SparkError::MergeCardinalityViolation => Some("MERGE_CARDINALITY_VIOLATION"), + // File not found SparkError::FileNotFound { .. } => Some("_LEGACY_ERROR_TEMP_2055"), diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 4d4094b1472..bed0123fdb1 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -306,6 +306,7 @@ fn op_name(op: &OpStruct) -> &'static str { OpStruct::ShuffleScan(_) => "ShuffleScan", OpStruct::BroadcastNestedLoopJoin(_) => "BroadcastNestedLoopJoin", OpStruct::Sample(_) => "Sample", + OpStruct::MergeRows(_) => "MergeRows", OpStruct::ContribScan(_) => "ContribScan", OpStruct::WindowGroupLimit(_) => "WindowGroupLimit", } diff --git a/native/core/src/execution/operators/merge_rows.rs b/native/core/src/execution/operators/merge_rows.rs new file mode 100644 index 00000000000..3017784ce67 --- /dev/null +++ b/native/core/src/execution/operators/merge_rows.rs @@ -0,0 +1,1320 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{Array, ArrayRef, BooleanArray, Int64Array, RecordBatch}; +use arrow::compute::kernels::boolean::{and, and_not, not}; +use arrow::compute::{filter_record_batch, prep_null_mask_filter}; +use arrow::datatypes::{DataType, SchemaRef}; +use datafusion::common::utils::memory::estimate_memory_size; +use datafusion::common::{DataFusionError, HashSet, ScalarValue}; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::{ + execution::TaskContext, + physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, + }, +}; +use datafusion_comet_common::{cast_and_stamp_schema, SparkError}; +use futures::{Stream, StreamExt}; +use std::{ + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; + +/// A MergeRows instruction: condition plus zero (Discard), one (Keep), or two (Split) +/// output row projections. +#[derive(Debug, Clone)] +pub struct MergeInstructionExec { + pub condition: Arc, + pub outputs: Vec>>, +} + +/// Immutable MergeRows predicates, instruction groups, and optional cardinality row ID. +#[derive(Debug)] +struct MergeConfig { + is_source_row_present: Arc, + is_target_row_present: Arc, + matched_instructions: Vec, + not_matched_instructions: Vec, + not_matched_by_source_instructions: Vec, + /// `Some(ordinal)` when cardinality checking is requested; `None` turns it off. One field + /// instead of a `(bool, usize)` pair, since the ordinal is meaningless without the flag. + row_id_ordinal: Option, +} + +impl MergeConfig { + /// Validate the optional row ID against the current child schema. + fn validate(&self, child: &Arc) -> Result<(), DataFusionError> { + if let Some(ordinal) = self.row_id_ordinal { + let child_schema = child.schema(); + let child_fields = child_schema.fields().len(); + if ordinal >= child_fields { + return Err(DataFusionError::Internal(format!( + "MergeRows: row id ordinal {ordinal} is out of range for a child with \ + {child_fields} columns" + ))); + } + let data_type = child_schema.field(ordinal).data_type(); + if data_type != &DataType::Int64 { + return Err(DataFusionError::Internal(format!( + "MergeRows: row id column at ordinal {ordinal} must be Int64, got {data_type}" + ))); + } + } + Ok(()) + } +} + +/// Native implementation of Spark's row-level `MergeRowsExec` dispatch operator. +#[derive(Debug)] +pub struct MergeRowsExec { + config: Arc, + child: Arc, + schema: SchemaRef, + cache: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl MergeRowsExec { + #[allow(clippy::too_many_arguments)] + pub fn try_new( + is_source_row_present: Arc, + is_target_row_present: Arc, + matched_instructions: Vec, + not_matched_instructions: Vec, + not_matched_by_source_instructions: Vec, + row_id_ordinal: Option, + child: Arc, + schema: SchemaRef, + ) -> Result { + let config = Arc::new(MergeConfig { + is_source_row_present, + is_target_row_present, + matched_instructions, + not_matched_instructions, + not_matched_by_source_instructions, + row_id_ordinal, + }); + config.validate(&child)?; + + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::UnknownPartitioning(1), + // One output batch per input batch -- nothing is buffered until the input ends, so + // this is `Incremental`, not `Final`. + EmissionType::Incremental, + Boundedness::Bounded, + )); + + Ok(Self { + config, + child, + schema, + cache, + metrics: ExecutionPlanMetricsSet::new(), + }) + } +} + +impl DisplayAs for MergeRowsExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "CometMergeRowsExec") + } + DisplayFormatType::TreeRender => unimplemented!(), + } + } +} + +impl ExecutionPlan for MergeRowsExec { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> datafusion::common::Result> { + let [child] = children.as_slice() else { + return Err(DataFusionError::Internal(format!( + "MergeRows expects exactly one child, got {}", + children.len() + ))); + }; + let child = Arc::clone(child); + // A replacement child may have a different row-ID schema. + self.config.validate(&child)?; + Ok(Arc::new(MergeRowsExec { + config: Arc::clone(&self.config), + child, + schema: Arc::clone(&self.schema), + cache: Arc::clone(&self.cache), + metrics: self.metrics.clone(), + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> datafusion::common::Result { + let reservation = MemoryConsumer::new(format!("CometMergeRowsExec[{partition}]")) + .register(&context.runtime_env().memory_pool); + let child_stream = self.child.execute(partition, Arc::clone(&context))?; + Ok(Box::pin(MergeRowsStream { + config: Arc::clone(&self.config), + child_stream, + schema: Arc::clone(&self.schema), + // Cardinality state is partition-scoped and must survive batch boundaries. + seen: HashSet::new(), + reservation, + baseline: BaselineMetrics::new(&self.metrics, partition), + })) + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn name(&self) -> &str { + "CometMergeRowsExec" + } +} + +pub struct MergeRowsStream { + config: Arc, + child_stream: SendableRecordBatchStream, + schema: SchemaRef, + /// Target row ids already seen in a matched pair. Accumulated across *every* batch polled + /// from this stream (i.e. for the lifetime of the partition), not reset per batch -- a + /// cardinality violation where the two matching source rows land in different Arrow batches + /// must still be caught. Mirrors Spark's `MergeRowsExec.BitmapCardinalityValidator`, which is + /// task-scoped, not batch-scoped. + seen: HashSet, + /// Pool accounting for [`MergeRowsStream::seen`]. Held for the life of the stream and + /// released on drop. + reservation: MemoryReservation, + /// `elapsed_compute` / `output_rows` / `output_batches`. Without these the merge operator is + /// invisible in the Spark UI and in benchmarking, so its share of a slow MERGE cannot be + /// separated from the upstream join/scan or the downstream write. `record_poll` (called at + /// the end of every `poll_next`) increments `output_rows` and `output_batches` itself for + /// every emitted batch -- do not additionally track either metric alongside `baseline`, or + /// the pair double-counts. + /// + /// `output_rows / output_batches` is this operator's average output batch size -- a + /// fragmented merge output slows the downstream writer even when the writer itself is fast, + /// so this is the number to check first when a MERGE's write phase is slow. + baseline: BaselineMetrics, +} + +/// Fixed overhead of the `seen` `HashSet` control struct itself (not its bucket array), passed as +/// the `fixed_size` term to [`estimate_memory_size`] so the reservation covers the whole +/// collection rather than only its bucket array. +const SEEN_FIXED_BYTES: usize = std::mem::size_of::>(); + +/// Conservative correction for details DataFusion's generic hash-table estimator deliberately +/// does not model exactly. hashbrown's raw table stores an additional mirrored control group and, +/// for very small tables, enforces a minimum bucket count so failed lookups always terminate. +/// `estimate_memory_size::(1, ...)`, for example, models one bucket while hashbrown allocates +/// four i64 buckets plus its control group. 64 bytes covers that first-allocation gap as well as +/// control-group/alignment overhead at subsequent rehashes on current 32/64-bit targets. The tests +/// below compare the reservation against hashbrown's own `allocation_size()` so a future layout +/// change that exceeds this correction fails loudly rather than silently weakening the pool limit. +const SEEN_HASH_TABLE_SLACK_BYTES: usize = 64; + +fn estimate_seen_memory_size(num_elements: usize) -> Result { + estimate_memory_size::(num_elements, SEEN_FIXED_BYTES)? + .checked_add(SEEN_HASH_TABLE_SLACK_BYTES) + .ok_or_else(|| { + DataFusionError::ResourcesExhausted( + "MergeRows: cardinality memory estimate overflow".to_string(), + ) + }) +} + +/// Rewrites NULL slots to `false`. Every boolean in this operator goes through Spark's +/// `BasePredicate.eval`, which collapses a NULL predicate result to `false`, but Arrow's +/// `and`/`and_not` kernels propagate NULL -- left unflattened, a NULL condition would poison +/// `run_group`'s shrinking `remaining` mask and silently drop the row from every later +/// instruction in the group, including the catch-all `Keep(TrueLiteral, ...)` Spark's +/// `RewriteMergeIntoTable` appends. `arrow::compute::prep_null_mask_filter` does the flattening +/// but panics when there are no nulls, hence the guard. +fn null_to_false(array: &BooleanArray) -> BooleanArray { + if array.null_count() == 0 { + array.clone() + } else { + prep_null_mask_filter(array) + } +} + +fn eval_bool( + expr: &Arc, + batch: &RecordBatch, +) -> Result { + let array: ArrayRef = expr.evaluate(batch)?.into_array(batch.num_rows())?; + array + .as_any() + .downcast_ref::() + .map(null_to_false) + .ok_or_else(|| DataFusionError::Internal("MergeRows: expected boolean array".to_string())) +} + +fn project( + batch: &RecordBatch, + exprs: &[Arc], + schema: &SchemaRef, +) -> Result { + let mut columns = Vec::with_capacity(exprs.len()); + for expr in exprs { + columns.push(expr.evaluate(batch)?.into_array(batch.num_rows())?); + } + // A clause's projected nested types (e.g. a struct built by `named_struct`) reflect only + // that expression's own nullability, not the wider nullability `self.schema` carries to + // accommodate every instruction's output -- a later `Keep` referencing the target schema's + // nullable field, for instance. `cast_and_stamp_schema` reconciles each column with the + // declared schema the way `ExpandStream::expand` does for the same reason. + cast_and_stamp_schema("MergeRows", schema, columns, batch.num_rows()) +} + +/// Filters `batch` to `mask`, skipping the copy when every row is already selected. +fn filter_or_pass_through( + batch: &RecordBatch, + mask: &BooleanArray, +) -> Result { + if mask.true_count() == batch.num_rows() { + Ok(batch.clone()) + } else { + filter_record_batch(batch, mask).map_err(|e| e.into()) + } +} + +/// Runs one instruction group (matched / not_matched / not_matched_by_source) over the rows +/// selected by `group_mask`, producing zero or more output batches. Reproduces Spark's ordered, +/// first-match-wins clause evaluation (`MergeRows`: "the first matching expression is used") +/// by physically shrinking the working batch to the rows still unclaimed after each instruction. +/// +/// Output rows come out grouped by the instruction that produced them rather than in input row +/// order -- this operator is set-at-a-time where Spark's is row-at-a-time. That is safe because +/// nothing downstream depends on this operator's row order: Iceberg applies its required +/// distribution and ordering to the *write's* input, so `DistributionAndOrderingUtils` places the +/// repartition and sort above `MergeRows`, not below it. A partitioned `ClusteredWriter` therefore +/// still receives partition-clustered input. Do not wire a writer directly to this operator's +/// output without preserving that sort. +fn run_group( + batch: &RecordBatch, + group_mask: &BooleanArray, + instructions: &[MergeInstructionExec], + schema: &SchemaRef, +) -> Result, DataFusionError> { + if instructions.is_empty() || group_mask.true_count() == 0 { + return Ok(vec![]); + } + + // Narrow to the group's rows *before* evaluating any condition. Spark reaches + // `applyInstructions` only after a row has been routed to a group, so a clause condition is + // never evaluated against a row belonging to another group. Evaluating over the whole batch + // would additionally expose rows the clause was never meant to see -- e.g. a NOT MATCHED + // condition `s.a / s.b > 1` evaluated on matched rows, where `s.b` is a real value and may + // be 0, raising an ANSI divide-by-zero that Spark would never produce. + let mut current = filter_or_pass_through(batch, group_mask)?; + let mut out = Vec::new(); + let last = instructions.len() - 1; + + for (idx, instr) in instructions.iter().enumerate() { + if current.num_rows() == 0 { + // Nothing left in this group can fire. + break; + } + + // A row already claimed by an earlier instruction must never reach a later one's + // condition -- not just have its result masked out, but be physically absent from + // `current` -- since evaluating the condition itself (e.g. `s.a / s.b > 1`) can raise + // under ANSI for a row Spark would never have reevaluated. This is why `current` shrinks + // every iteration instead of narrowing a same-sized mask alongside a stable batch. + // + // Spark's `RewriteMergeIntoTable` appends an unconditional catch-all + // `Keep(TrueLiteral, ...)` as the last instruction of the matched / not-matched-by-source + // groups. A literal condition evaluates to a `ColumnarValue::Scalar`, so handle it + // without materializing an all-true same-value array. + let fire = match instr.condition.evaluate(¤t)? { + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => { + BooleanArray::from(vec![true; current.num_rows()]) + } + ColumnarValue::Scalar(ScalarValue::Boolean(Some(false) | None)) => continue, + value => value + .into_array(current.num_rows())? + .as_any() + .downcast_ref::() + .map(null_to_false) + .ok_or_else(|| { + DataFusionError::Internal("MergeRows: expected boolean array".to_string()) + })?, + }; + + if fire.true_count() == 0 { + continue; + } + + let filtered = filter_or_pass_through(¤t, &fire)?; + for output_exprs in &instr.outputs { + out.push(project(&filtered, output_exprs, schema)?); + } + + if idx != last { + current = if fire.true_count() == current.num_rows() { + current.slice(0, 0) + } else { + filter_record_batch(¤t, ¬(&fire)?)? + }; + } + } + + Ok(out) +} + +/// Detects a target row matched by more than one source row (Spark's +/// `MERGE_CARDINALITY_VIOLATION`), mirroring `MergeRowsExec.BitmapCardinalityValidator`: track +/// row ids seen within the matched group and fail on the first repeat. +fn check_cardinality( + batch: &RecordBatch, + matched_mask: &BooleanArray, + row_id_ordinal: usize, + seen: &mut HashSet, + reservation: &mut MemoryReservation, +) -> Result<(), DataFusionError> { + // Read the row-id column in place and walk only the positions the mask selects. Filtering + // first would allocate a copy of the column on every poll purely to iterate it, and + // `filter_record_batch` over the whole batch would copy every other column too -- neither is + // needed, since this check reads one column and keeps nothing. + let row_ids = batch + .column(row_id_ordinal) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal("MergeRows: row id column must be Int64".to_string()) + })?; + + // Walk the matched rows one at a time, and only ask the pool for more memory at the moment + // `seen` is about to grow its bucket array -- never speculatively for the whole batch. + // + // Two invariants this ordering preserves: + // 1. A duplicate id is detected *before* any allocation it would have triggered, so a MERGE + // that is already a cardinality violation reports `MERGE_CARDINALITY_VIOLATION` rather + // than a memory error, matching Spark's row-at-a-time `BitmapCardinalityValidator`. + // 2. The DataFusion pool gets to reject the growth before hashbrown asks the allocator for + // the enlarged table (charging afterwards is too late -- the allocation already happened). + // + // DataFusion's estimator intentionally models hashbrown approximately, so + // `estimate_seen_memory_size` adds a small conservative correction for the mirrored control + // group / minimum small-table buckets. The reservation only grows, since `seen` never shrinks. + for i in matched_mask.values().set_indices() { + // Spark's `BitmapCardinalityValidator.validate` reads `InternalRow.getLong(ordinal)` + // unconditionally, with no null check (confirmed via bytecode): a null long field reads + // as 0 (`UnsafeRow.setNullAt` zeroes the value slot; `GenericInternalRow`'s boxed-null + // unboxes to 0 via Scala's `null.asInstanceOf[Long]`). Mirror that exactly rather than + // skipping null row ids -- skipping would miss a real cardinality violation where two + // matched rows both carry a null row id, and reading the Arrow value buffer's raw byte + // content at a null slot (`row_ids.value(i)` without the null check) would not, since + // Arrow does not guarantee null slots are zero-filled. + let id = if row_ids.is_null(i) { + 0 + } else { + row_ids.value(i) + }; + + if seen.len() < seen.capacity() { + // `insert` cannot grow the table here, so a single lookup covers both the + // duplicate check and the insert. + if !seen.insert(id) { + return Err(DataFusionError::External(Box::new( + SparkError::MergeCardinalityViolation, + ))); + } + } else { + // The table is full: an `insert` of a new id would rehash. Check for the duplicate + // first (no allocation), then reserve for the growth before it happens. The extra + // `contains` lookup only runs at a rehash boundary -- O(log n) times over the set. + if seen.contains(&id) { + return Err(DataFusionError::External(Box::new( + SparkError::MergeCardinalityViolation, + ))); + } + let next_len = seen.len().checked_add(1).ok_or_else(|| { + DataFusionError::ResourcesExhausted( + "MergeRows: cardinality set length overflow".to_string(), + ) + })?; + let projected_bytes = estimate_seen_memory_size(next_len)?; + let additional = projected_bytes.saturating_sub(reservation.size()); + reservation.try_grow(additional)?; + if let Err(e) = seen.try_reserve(1) { + // Pool admission happened before allocation by design. If the allocator itself + // then fails, return that admission immediately instead of leaving the pool + // artificially charged until the whole stream is dropped. + if additional != 0 { + reservation.shrink(additional); + } + return Err(DataFusionError::ResourcesExhausted(format!( + "MergeRows: failed to allocate cardinality set: {e}" + ))); + } + seen.insert(id); + } + } + Ok(()) +} + +fn process_batch( + batch: RecordBatch, + config: &MergeConfig, + // Caller-owned and threaded across every batch of the partition -- must NOT be created + // fresh per call, or a cardinality violation split across two batches goes undetected. + seen: &mut HashSet, + reservation: &mut MemoryReservation, + schema: &SchemaRef, +) -> Result { + let source_present = eval_bool(&config.is_source_row_present, &batch)?; + let target_present = eval_bool(&config.is_target_row_present, &batch)?; + + let matched_mask = and(&target_present, &source_present)?; + let not_matched_mask = and_not(&source_present, &target_present)?; + let not_matched_by_source_mask = and_not(&target_present, &source_present)?; + + // Checks cardinality for every matched row in the batch before evaluating any instruction, + // whereas Spark validates and applies instructions row-at-a-time, interleaved in scan order. + // When a single batch contains both a cardinality violation and an unrelated + // instruction-evaluation error (e.g. an ANSI divide-by-zero) on different rows, this can + // surface a different error than Spark would for the same input, depending on which row + // comes first. Reordering these two phases would only flip which case diverges, not fix it -- + // a true fix needs row-at-a-time evaluation, which conflicts with this operator's vectorized + // design (see `run_group`'s doc comment). Accepted as a known limitation: both paths still + // fail the query, just with a different error. + if let Some(row_id_ordinal) = config.row_id_ordinal { + check_cardinality(&batch, &matched_mask, row_id_ordinal, seen, reservation)?; + } + + let mut batches = Vec::new(); + for (mask, instructions) in [ + (&matched_mask, &config.matched_instructions), + (¬_matched_mask, &config.not_matched_instructions), + ( + ¬_matched_by_source_mask, + &config.not_matched_by_source_instructions, + ), + ] { + batches.extend(run_group(&batch, mask, instructions, schema)?); + } + + if batches.is_empty() { + return Ok(RecordBatch::new_empty(Arc::clone(schema))); + } + + arrow::compute::concat_batches(schema, &batches).map_err(|e| e.into()) +} + +/// Upper bound on how many consecutive all-discarded input batches `poll_next` will absorb +/// within a single call before yielding to the executor. Without this, a MERGE dominated by +/// DELETE clauses against an upstream that resolves synchronously (e.g. an already-materialized +/// child) could loop indefinitely inside one `poll_next` call without ever returning +/// `Poll::Pending`, starving other tasks on the same worker thread. 128 mirrors the budget Tokio +/// itself applies to cooperative scheduling. +const MAX_DISCARDED_BATCHES_PER_POLL: u32 = 128; + +impl Stream for MergeRowsStream { + type Item = datafusion::common::Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + // `MergeRowsStream` is structurally `Unpin` (every field is), so projecting a plain + // `&mut Self` out of the `Pin` is sound. Doing so lets us split disjoint field borrows -- + // `&mut this.seen` alongside the other `&this.*` borrows -- so cardinality state + // accumulates across every batch polled from this stream instead of resetting per batch. + let this = self.get_mut(); + // Loop rather than return the empty result: an input batch whose rows are all discarded + // (a copy-on-write DELETE clause, say) produces no output rows, and forwarding a zero-row + // batch makes every downstream stage pay for nothing -- an FFI export/import pair into + // the write pipeline, and in the partitioned case a full `RecordBatchPartitionSplitter` + // pass. Keep pulling until there is something to emit or the child is done. + let mut discarded_budget = MAX_DISCARDED_BATCHES_PER_POLL; + loop { + let poll = match this.child_stream.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => { + // Times only this operator's own work; the upstream poll above is + // deliberately outside the timer so `elapsed_compute` is not the whole + // pipeline's wall clock. + let _timer = this.baseline.elapsed_compute().timer(); + let result = process_batch( + batch, + &this.config, + &mut this.seen, + &mut this.reservation, + &this.schema, + ); + match result { + Ok(batch) if batch.num_rows() == 0 => { + discarded_budget -= 1; + if discarded_budget == 0 { + // Give the executor a chance to run other tasks before pulling + // more batches; re-polling this stream is what drives progress + // here, so wake immediately rather than waiting on the child. + cx.waker().wake_by_ref(); + return Poll::Pending; + } + continue; + } + other => Poll::Ready(Some(other)), + } + } + other => other, + }; + return this.baseline.record_poll(poll); + } + } +} + +impl RecordBatchStream for MergeRowsStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int32Array, StructArray}; + use arrow::datatypes::{Field, Schema}; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; + use datafusion::logical_expr::Operator as DFOperator; + use datafusion::physical_expr::expressions::{binary, col, lit}; + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("row_id", DataType::Int64, true), + Field::new("val", DataType::Int32, true), + Field::new("target_present", DataType::Boolean, false), + Field::new("source_present", DataType::Boolean, false), + ])) + } + + fn test_batch( + row_ids: Vec, + vals: Vec, + target: Vec, + source: Vec, + ) -> RecordBatch { + RecordBatch::try_new( + test_schema(), + vec![ + Arc::new(Int64Array::from(row_ids)), + Arc::new(Int32Array::from(vals)), + Arc::new(BooleanArray::from(target)), + Arc::new(BooleanArray::from(source)), + ], + ) + .unwrap() + } + + /// Pool accounting is not what these tests exercise, so they run against an unbounded pool. + fn test_reservation() -> MemoryReservation { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("test").register(&pool) + } + + fn out_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("val", DataType::Int32, true)])) + } + + fn keep_all() -> MergeInstructionExec { + MergeInstructionExec { + condition: lit(true), + outputs: vec![vec![col("val", &test_schema()).unwrap()]], + } + } + + fn discard_all() -> MergeInstructionExec { + MergeInstructionExec { + condition: lit(true), + outputs: vec![], + } + } + + fn test_config( + matched_instructions: Vec, + not_matched_instructions: Vec, + not_matched_by_source_instructions: Vec, + row_id_ordinal: Option, + ) -> MergeConfig { + MergeConfig { + is_source_row_present: col("source_present", &test_schema()).unwrap(), + is_target_row_present: col("target_present", &test_schema()).unwrap(), + matched_instructions, + not_matched_instructions, + not_matched_by_source_instructions, + row_id_ordinal, + } + } + + #[test] + fn keep_matched_discard_rest() { + let batch = test_batch( + vec![1, 2, 3], + vec![10, 20, 30], + vec![true, false, true], + vec![true, true, false], + ); + let config = test_config( + vec![keep_all()], + vec![keep_all()], + vec![discard_all()], + None, + ); + let out = process_batch( + batch, + &config, + &mut HashSet::new(), + &mut test_reservation(), + &out_schema(), + ) + .unwrap(); + let vals = out.column(0).as_any().downcast_ref::().unwrap(); + let mut got: Vec = vals.iter().flatten().collect(); + got.sort(); + assert_eq!(got, vec![10, 20]); + } + + #[test] + fn first_match_wins_ordering() { + let batch = test_batch(vec![1], vec![5], vec![true], vec![true]); + let cond_false = MergeInstructionExec { + condition: binary( + col("val", &test_schema()).unwrap(), + DFOperator::Gt, + lit(100i32), + &test_schema(), + ) + .unwrap(), + outputs: vec![vec![lit(1i32)]], + }; + let cond_true = MergeInstructionExec { + condition: lit(true), + outputs: vec![vec![lit(2i32)]], + }; + let config = test_config(vec![cond_false, cond_true], vec![], vec![], None); + let out = process_batch( + batch, + &config, + &mut HashSet::new(), + &mut test_reservation(), + &out_schema(), + ) + .unwrap(); + let vals = out.column(0).as_any().downcast_ref::().unwrap(); + assert_eq!(vals.value(0), 2); + } + + #[test] + fn later_condition_is_not_evaluated_on_an_already_claimed_row() { + let batch = test_batch(vec![1, 2], vec![0, 2], vec![true, true], vec![true, true]); + let claims_zero = MergeInstructionExec { + condition: binary( + col("val", &test_schema()).unwrap(), + DFOperator::Eq, + lit(0i32), + &test_schema(), + ) + .unwrap(), + outputs: vec![vec![lit(111i32)]], + }; + let divides_by_val = MergeInstructionExec { + condition: binary( + binary( + lit(2i32), + DFOperator::Divide, + col("val", &test_schema()).unwrap(), + &test_schema(), + ) + .unwrap(), + DFOperator::Gt, + lit(0i32), + &test_schema(), + ) + .unwrap(), + outputs: vec![vec![lit(222i32)]], + }; + let config = test_config(vec![claims_zero, divides_by_val], vec![], vec![], None); + let out = process_batch( + batch, + &config, + &mut HashSet::new(), + &mut test_reservation(), + &out_schema(), + ) + .unwrap(); + let vals = out.column(0).as_any().downcast_ref::().unwrap(); + let mut got: Vec = vals.iter().flatten().collect(); + got.sort(); + assert_eq!(got, vec![111, 222]); + } + + #[test] + fn null_condition_falls_through_to_next_instruction() { + let batch = RecordBatch::try_new( + test_schema(), + vec![ + Arc::new(Int64Array::from(vec![1i64])), + Arc::new(Int32Array::from(vec![None::])), + Arc::new(BooleanArray::from(vec![true])), + Arc::new(BooleanArray::from(vec![true])), + ], + ) + .unwrap(); + let cond_null = MergeInstructionExec { + condition: binary( + col("val", &test_schema()).unwrap(), + DFOperator::Gt, + lit(100i32), + &test_schema(), + ) + .unwrap(), + outputs: vec![vec![lit(1i32)]], + }; + let keep_catch_all = MergeInstructionExec { + condition: lit(true), + outputs: vec![vec![lit(2i32)]], + }; + let config = test_config(vec![cond_null, keep_catch_all], vec![], vec![], None); + let out = process_batch( + batch, + &config, + &mut HashSet::new(), + &mut test_reservation(), + &out_schema(), + ) + .unwrap(); + assert_eq!( + out.num_rows(), + 1, + "row with a NULL clause condition must fall through to the catch-all Keep, not \ + disappear from the rewritten data file" + ); + let vals = out.column(0).as_any().downcast_ref::().unwrap(); + assert_eq!(vals.value(0), 2); + } + + #[test] + fn null_row_presence_flag_treated_as_false() { + let schema = Arc::new(Schema::new(vec![ + Field::new("row_id", DataType::Int64, true), + Field::new("val", DataType::Int32, true), + Field::new("target_present", DataType::Boolean, true), + Field::new("source_present", DataType::Boolean, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![1i64])), + Arc::new(Int32Array::from(vec![10])), + Arc::new(BooleanArray::from(vec![Some(true)])), + Arc::new(BooleanArray::from(vec![None::])), + ], + ) + .unwrap(); + let config = MergeConfig { + is_source_row_present: col("source_present", &schema).unwrap(), + is_target_row_present: col("target_present", &schema).unwrap(), + matched_instructions: vec![], + not_matched_instructions: vec![], + not_matched_by_source_instructions: vec![MergeInstructionExec { + condition: lit(true), + outputs: vec![vec![col("val", &schema).unwrap()]], + }], + row_id_ordinal: None, + }; + let out = process_batch( + batch, + &config, + &mut HashSet::new(), + &mut test_reservation(), + &out_schema(), + ) + .unwrap(); + assert_eq!(out.num_rows(), 1); + let vals = out.column(0).as_any().downcast_ref::().unwrap(); + assert_eq!(vals.value(0), 10); + } + + #[test] + fn condition_not_evaluated_outside_its_group() { + let batch = test_batch(vec![1, 2], vec![0, 5], vec![true, false], vec![true, true]); + let div_cond = MergeInstructionExec { + condition: binary( + binary( + lit(10i32), + DFOperator::Divide, + col("val", &test_schema()).unwrap(), + &test_schema(), + ) + .unwrap(), + DFOperator::Gt, + lit(1i32), + &test_schema(), + ) + .unwrap(), + outputs: vec![vec![col("val", &test_schema()).unwrap()]], + }; + assert!( + eval_bool(&div_cond.condition, &batch).is_err(), + "test is only meaningful if batch-wide evaluation of this condition errors" + ); + let config = test_config(vec![keep_all()], vec![div_cond], vec![], None); + let out = process_batch( + batch, + &config, + &mut HashSet::new(), + &mut test_reservation(), + &out_schema(), + ) + .expect("not-matched condition must not be evaluated against the matched row"); + let vals = out.column(0).as_any().downcast_ref::().unwrap(); + let mut got: Vec = vals.iter().flatten().collect(); + got.sort(); + assert_eq!(got, vec![0, 5]); + } + + fn bounded_reservation(limit: usize) -> MemoryReservation { + let pool: Arc = Arc::new(GreedyMemoryPool::new(limit)); + MemoryConsumer::new("test").register(&pool) + } + + fn run_cardinality( + n: usize, + batch_rows: usize, + reservation: &mut MemoryReservation, + ) -> Result, DataFusionError> { + let mut seen = HashSet::new(); + let mut next = 0i64; + while (next as usize) < n { + let end = ((next as usize) + batch_rows).min(n) as i64; + let ids: Vec = (next..end).collect(); + let len = ids.len(); + let batch = test_batch(ids, vec![0; len], vec![true; len], vec![true; len]); + let mask = BooleanArray::from(vec![true; len]); + check_cardinality(&batch, &mask, 0, &mut seen, reservation)?; + next = end; + } + Ok(seen) + } + + fn assert_seen_fully_reserved(seen: &HashSet, reservation: &MemoryReservation) { + let actual = seen.allocation_size().saturating_add(SEEN_FIXED_BYTES); + assert!( + reservation.size() >= actual, + "reserved {} bytes < actual HashSet footprint {} bytes (len={}, capacity={})", + reservation.size(), + actual, + seen.len(), + seen.capacity() + ); + } + + #[test] + fn cardinality_state_is_accounted_to_the_memory_pool() { + let mut reservation = test_reservation(); + let seen = run_cardinality(9, 4, &mut reservation).unwrap(); + assert!( + reservation.size() > 0, + "`seen` must be visible to the memory pool" + ); + assert_seen_fully_reserved(&seen, &reservation); + } + + #[test] + fn cardinality_reservation_covers_actual_hashbrown_allocations() { + for &(n, batch_rows) in &[ + (1usize, 1usize), + (2, 1), + (3, 1), + (7, 1), + (8, 1), + (9, 1), + (17, 1), + (64, 8), + (200, 16), + ] { + let mut reservation = test_reservation(); + let seen = run_cardinality(n, batch_rows, &mut reservation).unwrap(); + assert_seen_fully_reserved(&seen, &reservation); + } + } + + #[test] + fn cardinality_state_cannot_exceed_a_bounded_pool() { + let n = 917_505; + let mut reservation = bounded_reservation(16 * 1024 * 1024); + let err = run_cardinality(n, 4096, &mut reservation).unwrap_err(); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "expected the pool to reject the oversized cardinality table, got {err}" + ); + + let needed = estimate_seen_memory_size(n).unwrap(); + let mut ok_reservation = bounded_reservation(needed + 8 * 1024 * 1024); + let seen = run_cardinality(n, 4096, &mut ok_reservation).unwrap(); + assert_eq!(seen.len(), n); + assert_seen_fully_reserved(&seen, &ok_reservation); + } + + #[test] + fn cardinality_violation_wins_over_memory_exhaustion() { + let mut seen = HashSet::new(); + let mut reservation = bounded_reservation(0); + let batch = test_batch(vec![1], vec![0], vec![true], vec![true]); + check_cardinality( + &batch, + &BooleanArray::from(vec![true]), + 0, + &mut seen, + &mut test_reservation(), + ) + .unwrap(); + let dup = test_batch(vec![1, 2], vec![0, 0], vec![true, true], vec![true, true]); + let err = check_cardinality( + &dup, + &BooleanArray::from(vec![true, true]), + 0, + &mut seen, + &mut reservation, + ) + .unwrap_err(); + assert!( + err.to_string().contains("MERGE_CARDINALITY_VIOLATION"), + "expected cardinality violation, got {err}" + ); + } + + #[test] + fn cardinality_violation_detected() { + let batch = test_batch(vec![1, 1], vec![10, 20], vec![true, true], vec![true, true]); + let matched_mask = BooleanArray::from(vec![true, true]); + let mut seen = HashSet::new(); + let result = + check_cardinality(&batch, &matched_mask, 0, &mut seen, &mut test_reservation()); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("MERGE_CARDINALITY_VIOLATION")); + } + + #[test] + fn split_produces_two_rows() { + let batch = test_batch(vec![1], vec![7], vec![true], vec![true]); + let split = MergeInstructionExec { + condition: lit(true), + outputs: vec![vec![lit(1i32)], vec![lit(2i32)]], + }; + let config = test_config(vec![split], vec![], vec![], None); + let out = process_batch( + batch, + &config, + &mut HashSet::new(), + &mut test_reservation(), + &out_schema(), + ) + .unwrap(); + let vals = out.column(0).as_any().downcast_ref::().unwrap(); + let got: Vec = vals.iter().flatten().collect(); + assert_eq!(got, vec![1, 2]); + } + + #[test] + fn cardinality_violation_detected_for_null_row_ids() { + let batch = RecordBatch::try_new( + test_schema(), + vec![ + Arc::new(Int64Array::from(vec![None, None])), + Arc::new(Int32Array::from(vec![10, 20])), + Arc::new(BooleanArray::from(vec![true, true])), + Arc::new(BooleanArray::from(vec![true, true])), + ], + ) + .unwrap(); + let matched_mask = BooleanArray::from(vec![true, true]); + let result = check_cardinality( + &batch, + &matched_mask, + 0, + &mut HashSet::new(), + &mut test_reservation(), + ); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("MERGE_CARDINALITY_VIOLATION")); + } + + #[tokio::test] + async fn all_discarded_batch_is_not_emitted() { + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::prelude::SessionContext; + + let discarded = test_batch(vec![1], vec![10], vec![true], vec![true]); + let kept = test_batch(vec![2], vec![20], vec![false], vec![true]); + let source = + MemorySourceConfig::try_new_exec(&[vec![discarded, kept]], test_schema(), None) + .unwrap(); + + let exec = MergeRowsExec::try_new( + col("source_present", &test_schema()).unwrap(), + col("target_present", &test_schema()).unwrap(), + vec![discard_all()], + vec![keep_all()], + vec![], + None, + source, + out_schema(), + ) + .unwrap(); + + let ctx = SessionContext::new(); + let mut stream = exec.execute(0, ctx.task_ctx()).unwrap(); + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); + } + + #[test] + fn out_of_range_row_id_ordinal_is_rejected() { + use datafusion::datasource::memory::MemorySourceConfig; + let source = MemorySourceConfig::try_new_exec(&[vec![]], test_schema(), None).unwrap(); + let err = MergeRowsExec::try_new( + col("source_present", &test_schema()).unwrap(), + col("target_present", &test_schema()).unwrap(), + vec![keep_all()], + vec![], + vec![], + Some(99), + source, + out_schema(), + ) + .unwrap_err(); + assert!(err.to_string().contains("row id ordinal")); + } + + #[test] + fn non_int64_row_id_is_rejected_at_plan_construction() { + use datafusion::datasource::memory::MemorySourceConfig; + let wrong_schema = Arc::new(Schema::new(vec![ + Field::new("row_id", DataType::Int32, true), + Field::new("val", DataType::Int32, true), + Field::new("target_present", DataType::Boolean, false), + Field::new("source_present", DataType::Boolean, false), + ])); + let source = + MemorySourceConfig::try_new_exec(&[vec![]], Arc::clone(&wrong_schema), None).unwrap(); + let err = MergeRowsExec::try_new( + col("source_present", &wrong_schema).unwrap(), + col("target_present", &wrong_schema).unwrap(), + vec![], + vec![], + vec![], + Some(0), + source, + out_schema(), + ) + .unwrap_err(); + assert!(err.to_string().contains("must be Int64")); + } + + #[test] + fn with_new_children_rejects_ordinal_out_of_range_for_new_child() { + use datafusion::datasource::memory::MemorySourceConfig; + let original_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + Field::new("row_id", DataType::Int64, true), + ])); + let source = + MemorySourceConfig::try_new_exec(&[vec![]], Arc::clone(&original_schema), None) + .unwrap(); + let exec = Arc::new( + MergeRowsExec::try_new( + lit(true), + lit(true), + vec![], + vec![], + vec![], + Some(3), + source, + out_schema(), + ) + .unwrap(), + ); + + let narrow_schema = Arc::new(Schema::new(vec![Field::new( + "row_id", + DataType::Int64, + true, + )])); + let narrow_child = + MemorySourceConfig::try_new_exec(&[vec![]], narrow_schema, None).unwrap(); + let err = exec.with_new_children(vec![narrow_child]).unwrap_err(); + assert!( + err.to_string().contains("out of range"), + "expected an out-of-range ordinal error, got: {err}" + ); + } + + #[test] + fn with_new_children_rejects_wrong_arity() { + use datafusion::datasource::memory::MemorySourceConfig; + let source = MemorySourceConfig::try_new_exec(&[vec![]], test_schema(), None).unwrap(); + let exec = Arc::new( + MergeRowsExec::try_new( + col("source_present", &test_schema()).unwrap(), + col("target_present", &test_schema()).unwrap(), + vec![keep_all()], + vec![], + vec![], + Some(0), + source, + out_schema(), + ) + .unwrap(), + ); + let no_children = Arc::clone(&exec).with_new_children(vec![]).unwrap_err(); + assert!(no_children.to_string().contains("exactly one child")); + + let child_a = MemorySourceConfig::try_new_exec(&[vec![]], test_schema(), None).unwrap(); + let child_b = MemorySourceConfig::try_new_exec(&[vec![]], test_schema(), None).unwrap(); + let two_children = exec.with_new_children(vec![child_a, child_b]).unwrap_err(); + assert!(two_children.to_string().contains("exactly one child")); + } + + #[test] + fn with_new_children_revalidates_row_id_schema() { + use datafusion::datasource::memory::MemorySourceConfig; + let source = MemorySourceConfig::try_new_exec(&[vec![]], test_schema(), None).unwrap(); + let exec = Arc::new( + MergeRowsExec::try_new( + col("source_present", &test_schema()).unwrap(), + col("target_present", &test_schema()).unwrap(), + vec![keep_all()], + vec![], + vec![], + Some(0), + source, + out_schema(), + ) + .unwrap(), + ); + + let narrow_schema = Arc::new(Schema::new(vec![Field::new( + "only_col", + DataType::Int32, + true, + )])); + let narrow_child = + MemorySourceConfig::try_new_exec(&[vec![]], narrow_schema, None).unwrap(); + let err = exec.with_new_children(vec![narrow_child]).unwrap_err(); + assert!(err.to_string().contains("must be Int64")); + } + + #[test] + fn cardinality_violation_detected_across_batches() { + let mut seen = HashSet::new(); + let batch1 = test_batch(vec![1], vec![10], vec![true], vec![true]); + let batch2 = test_batch(vec![1], vec![20], vec![true], vec![true]); + let config = test_config(vec![keep_all()], vec![], vec![], Some(0)); + + let first = process_batch( + batch1, + &config, + &mut seen, + &mut test_reservation(), + &out_schema(), + ); + assert!(first.is_ok()); + + let second = process_batch( + batch2, + &config, + &mut seen, + &mut test_reservation(), + &out_schema(), + ); + assert!(second.is_err()); + assert!(second + .unwrap_err() + .to_string() + .contains("MERGE_CARDINALITY_VIOLATION")); + } + + #[test] + fn project_reconciles_nested_struct_nullability_with_declared_schema() { + let source_field = Field::new("n", DataType::Boolean, false); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Struct(vec![source_field.clone()].into()), + true, + )])); + let inner_values: ArrayRef = Arc::new(BooleanArray::from(vec![true])); + let payload_array: ArrayRef = Arc::new(StructArray::new( + vec![source_field].into(), + vec![inner_values], + None, + )); + let batch = RecordBatch::try_new(Arc::clone(&source_schema), vec![payload_array]).unwrap(); + + let target_field = Field::new("n", DataType::Boolean, true); + let out_schema = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Struct(vec![target_field].into()), + true, + )])); + + let out = project( + &batch, + &[col("payload", &source_schema).unwrap()], + &out_schema, + ) + .expect("project must reconcile projected nested nullability with the declared schema"); + assert_eq!( + out.schema().field(0).data_type(), + out_schema.field(0).data_type() + ); + } +} diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index 1d3c0749fbb..b636ea5eb59 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -31,6 +31,8 @@ pub use expand::ExpandExec; mod explode; pub use explode::ExplodeExec; mod iceberg_scan; +mod merge_rows; +pub use merge_rows::{MergeInstructionExec, MergeRowsExec}; mod parquet_writer; pub use parquet_writer::{ParquetCompression, ParquetWriterExec}; mod csv_scan; diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 45c30f4779b..6a4f31f33c1 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -39,8 +39,8 @@ use crate::execution::{ expressions::list_positions::ListPositionsExpr, expressions::subquery::Subquery, operators::{ - ExecutionError, ExpandExec, ExplodeExec, ParquetCompression, ParquetWriterExec, SampleExec, - ScanExec, ShuffleScanExec, + ExecutionError, ExpandExec, ExplodeExec, MergeInstructionExec, MergeRowsExec, + ParquetCompression, ParquetWriterExec, SampleExec, ScanExec, ShuffleScanExec, }, planner::expression_registry::ExpressionRegistry, planner::operator_registry::OperatorRegistry, @@ -2021,6 +2021,107 @@ impl PhysicalPlanner { Arc::new(SparkPlan::new(spark_plan.plan_id, expand, vec![child])), )) } + OpStruct::MergeRows(merge) => { + let [child] = children.as_slice() else { + return Err(ExecutionError::GeneralError(format!( + "MergeRows expects exactly one child, got {}", + children.len() + ))); + }; + let (scans, shuffle_scans, child) = + self.create_plan(child, inputs, partition_count)?; + + let missing_field = |name: &str| { + ExecutionError::GeneralError(format!("MergeRows proto missing `{name}`")) + }; + let is_source_row_present = self.create_expr( + merge + .is_source_row_present + .as_ref() + .ok_or_else(|| missing_field("is_source_row_present"))?, + child.schema(), + )?; + let is_target_row_present = self.create_expr( + merge + .is_target_row_present + .as_ref() + .ok_or_else(|| missing_field("is_target_row_present"))?, + child.schema(), + )?; + + let compile_instructions = |instrs: &[spark_operator::MergeInstruction]| -> Result< + Vec, + ExecutionError, + > { + instrs + .iter() + .map(|instr| { + let condition = self.create_expr( + instr + .condition + .as_ref() + .ok_or_else(|| missing_field("instruction condition"))?, + child.schema(), + )?; + let outputs = instr + .outputs + .iter() + .map(|row| { + row.exprs + .iter() + .map(|e| self.create_expr(e, child.schema())) + .collect::, _>>() + }) + .collect::, _>>()?; + Ok(MergeInstructionExec { condition, outputs }) + }) + .collect() + }; + + let matched_instructions = compile_instructions(&merge.matched_instructions)?; + let not_matched_instructions = + compile_instructions(&merge.not_matched_instructions)?; + let not_matched_by_source_instructions = + compile_instructions(&merge.not_matched_by_source_instructions)?; + + // Derive projected types like Expand so encoded/nested expression types match the + // batches we emit. If every instruction discards, use Spark's declared output types. + let output_rows: Vec>> = matched_instructions + .iter() + .chain(¬_matched_instructions) + .chain(¬_matched_by_source_instructions) + .flat_map(|instr| instr.outputs.iter().cloned()) + .collect(); + let schema = if output_rows.is_empty() { + let fields: Vec = merge + .output_types + .iter() + .map(to_arrow_datatype) + .enumerate() + .map(|(idx, dt)| Field::new(format!("col_{idx}"), dt, true)) + .collect(); + Arc::new(Schema::new(fields)) + } else { + ExpandExec::build_schema(&output_rows, &child.schema())? + }; + + let exec = Arc::new(MergeRowsExec::try_new( + is_source_row_present, + is_target_row_present, + matched_instructions, + not_matched_instructions, + not_matched_by_source_instructions, + merge.row_id_ordinal.map(|ord| ord as usize), + Arc::clone(&child.native_plan), + schema, + )?); + + Ok(( + scans, + shuffle_scans, + Arc::new(SparkPlan::new(spark_plan.plan_id, exec, vec![child])), + )) + } OpStruct::Explode(explode) => { assert_eq!(children.len(), 1); let (scans, shuffle_scans, child) = diff --git a/native/core/src/execution/planner/operator_registry.rs b/native/core/src/execution/planner/operator_registry.rs index c644741894c..719c78f90de 100644 --- a/native/core/src/execution/planner/operator_registry.rs +++ b/native/core/src/execution/planner/operator_registry.rs @@ -152,7 +152,8 @@ fn get_operator_type(spark_operator: &Operator) -> Option { OpStruct::CsvScan(_) => Some(OperatorType::CsvScan), OpStruct::ShuffleScan(_) => None, // Not yet in OperatorType enum OpStruct::BroadcastNestedLoopJoin(_) => None, - OpStruct::Sample(_) => None, // Not yet in OperatorType enum + OpStruct::Sample(_) => None, // Not yet in OperatorType enum + OpStruct::MergeRows(_) => None, // Not yet in OperatorType enum // Generic extension point for out-of-tree contrib scans (Delta, Lance, ...); not in // OperatorType enum. The arm stays unconditional even in non-contrib builds because the // proto enum is generated regardless of cargo features and Rust requires an exhaustive diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index acac87d53ea..21daa15c7b1 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -67,6 +67,7 @@ message Operator { BroadcastNestedLoopJoin broadcast_nested_loop_join = 117; Sample sample = 118; WindowGroupLimit window_group_limit = 119; + MergeRows merge_rows = 120; // Extension point for optional, out-of-tree contrib scans (Delta, Lance, ...). The concrete // scan message (e.g. `DeltaScan`) is packed into this envelope on the JVM side and dispatched // by `type_url` on the native side. Using a single permanent field -- rather than a new oneof @@ -763,6 +764,39 @@ message Expand { int32 num_expr_per_project = 3; } +// Native counterpart of Spark's `MergeRowsExec` (row-level MERGE dispatch). Mirrors the real +// Spark 4.x bytecode shape: `isSourceRowPresent` / `isTargetRowPresent` are predicates (not +// column ordinals), and each instruction is uniformly `condition + outputs`, where the number +// of output row projections (0/1/2) distinguishes Discard/Keep/Split -- there is no separate +// instruction-kind enum on the Spark side, so we don't invent one here either. +message MergeRows { + spark.spark_expression.Expr is_source_row_present = 1; + spark.spark_expression.Expr is_target_row_present = 2; + + repeated MergeInstruction matched_instructions = 3; + repeated MergeInstruction not_matched_instructions = 4; + repeated MergeInstruction not_matched_by_source_instructions = 5; + + // Ordinal (into the child row) of the target row-id column used for cardinality dedup. + // Mirrors `MergeRowsExec.checkCardinality`: present iff the check is on, in which case + // MERGE_CARDINALITY_VIOLATION is raised if a target row matches more than one source row. + optional int32 row_id_ordinal = 6; + + // Schema of the emitted rows, matching `MergeRowsExec.output`. + repeated spark.spark_expression.DataType output_types = 7; +} + +message MergeInstruction { + // Always present on the Spark side (Keep/Discard/Split.condition are non-optional). + spark.spark_expression.Expr condition = 1; + // 0 output rows = Discard, 1 = Keep, 2 = Split. + repeated MergeOutputRow outputs = 2; +} + +message MergeOutputRow { + repeated spark.spark_expression.Expr exprs = 1; +} + message Explode { // The array expression to explode into multiple rows spark.spark_expression.Expr child = 1; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a7d7b80db1f..1b36c96fdd7 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -257,6 +257,8 @@ object CometConf extends ShimCometConf { createExecEnabledConfig("localTableScan", defaultValue = false) val COMET_EXEC_SAMPLE_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("sample", defaultValue = true) + val COMET_EXEC_MERGE_ROWS_ENABLED: ConfigEntry[Boolean] = + createExecEnabledConfig("mergeRows", defaultValue = false) val COMET_EXEC_IN_MEMORY_CACHE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.inMemoryCache.enabled") diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 8c842e68abe..84d68c3a5f5 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -55,6 +55,7 @@ import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometConf.{COMET_SPARK_TO_ARROW_ENABLED, COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST} import org.apache.comet.CometSparkSessionExtensions._ import org.apache.comet.rules.CometExecRule.allExecs +import org.apache.comet.rules.shims.ShimCometMergeRows import org.apache.comet.serde._ import org.apache.comet.serde.operator._ import org.apache.comet.shims.{ShimCometStreaming, ShimCometWindowGroupLimit, ShimSubqueryBroadcast} @@ -92,7 +93,9 @@ object CometExecRule { classOf[SampleExec] -> CometSampleExec, classOf[WindowExec] -> CometWindowExec) ++ // WindowGroupLimitExec exists only on Spark 3.5+; the shim returns None on 3.4. - ShimCometWindowGroupLimit.windowGroupLimitClass.map(_ -> CometWindowGroupLimitExec) + ShimCometWindowGroupLimit.windowGroupLimitClass.map(_ -> CometWindowGroupLimitExec) ++ + // MergeRowsExec exists only on Spark 3.5+; the shim is empty on 3.4. + ShimCometMergeRows.nativeExecs /** * Sinks that have a native plan of ScanExec. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 3700e97642b..7256b2489c8 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1433,6 +1433,89 @@ case class CometExpandExec( override lazy val metrics: Map[String, SQLMetric] = Map.empty } +/** + * Native wrapper for Spark's `MergeRowsExec`. Predicates and instruction groups are real + * case-class fields so Catalyst can discover their references and scalar subqueries. + */ +case class CometMergeRowsExec( + override val nativeOp: Operator, + override val originalPlan: SparkPlan, + override val output: Seq[Attribute], + isSourceRowPresent: Expression, + isTargetRowPresent: Expression, + matchedInstructions: Seq[Expression], + notMatchedInstructions: Seq[Expression], + notMatchedBySourceInstructions: Seq[Expression], + checkCardinality: Boolean, + rowIdOrdinal: Option[Int], + child: SparkPlan, + override val serializedPlanOpt: SerializedPlan) + extends CometUnaryExec { + // Match Spark's MergeRowsExec partitioning contract. + override def outputPartitioning: Partitioning = UnknownPartitioning(0) + + // Only attributes not already supplied by the child are produced here. + override def producedAttributes: AttributeSet = + AttributeSet(output.filterNot(child.outputSet.contains)) + + // Cardinality checking also reads Spark's synthetic ROW_ID column. + @transient + override lazy val references: AttributeSet = { + val rowIdExprs = rowIdOrdinal.flatMap(child.output.lift).toSeq + val expressionReferences = + AttributeSet.fromAttributeSets((rowIdExprs ++ expressions).map(_.references)) + expressionReferences -- producedAttributes + } + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + this.copy(child = newChild) + + override def stringArgs: Iterator[Any] = + Iterator( + output, + matchedInstructions, + notMatchedInstructions, + notMatchedBySourceInstructions, + checkCardinality, + child) + + override def equals(obj: Any): Boolean = { + obj match { + case other: CometMergeRowsExec => + this.output == other.output && + this.isSourceRowPresent == other.isSourceRowPresent && + this.isTargetRowPresent == other.isTargetRowPresent && + this.matchedInstructions == other.matchedInstructions && + this.notMatchedInstructions == other.notMatchedInstructions && + this.notMatchedBySourceInstructions == other.notMatchedBySourceInstructions && + this.checkCardinality == other.checkCardinality && + this.rowIdOrdinal == other.rowIdOrdinal && + this.child == other.child && + this.serializedPlanOpt == other.serializedPlanOpt + case _ => + false + } + } + + override def hashCode(): Int = + Objects.hashCode( + output, + isSourceRowPresent, + isTargetRowPresent, + matchedInstructions, + notMatchedInstructions, + notMatchedBySourceInstructions, + Boolean.box(checkCardinality), + rowIdOrdinal, + child) + + // Spark 4.x per-clause metrics require instruction context that Spark 3.5 lacks. + // Expose baseline metrics until that context is version-gated through native serde. + override lazy val metrics: Map[String, SQLMetric] = + CometMetricNode.baselineMetrics(sparkContext) ++ Map( + "output_batches" -> SQLMetrics.createMetric(sparkContext, "number of output batches")) +} + object CometExplodeExec extends CometOperatorSerde[GenerateExec] { override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( diff --git a/spark/src/main/spark-3.4/org/apache/comet/rules/shims/ShimCometMergeRows.scala b/spark/src/main/spark-3.4/org/apache/comet/rules/shims/ShimCometMergeRows.scala new file mode 100644 index 00000000000..79057f09068 --- /dev/null +++ b/spark/src/main/spark-3.4/org/apache/comet/rules/shims/ShimCometMergeRows.scala @@ -0,0 +1,33 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.rules.shims + +import org.apache.spark.sql.execution.SparkPlan + +import org.apache.comet.serde.CometOperatorSerde + +/** + * Spark 3.4 predates `MergeRowsExec` (it was moved from Iceberg extensions into Spark core in + * Iceberg 1.4.0 / SPARK-52403, first shipping in Spark 3.5). Nothing to register here; CoW MERGE + * on 3.4 continues to run via Iceberg's own extension-provided operator, unconverted. + */ +object ShimCometMergeRows { + val nativeExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = Map.empty +} diff --git a/spark/src/main/spark-3.5/org/apache/comet/rules/shims/ShimCometMergeRows.scala b/spark/src/main/spark-3.5/org/apache/comet/rules/shims/ShimCometMergeRows.scala new file mode 100644 index 00000000000..8e1c59dae6c --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/comet/rules/shims/ShimCometMergeRows.scala @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.rules.shims + +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.MergeRowsExec + +import org.apache.comet.serde.CometOperatorSerde +import org.apache.comet.serde.operator.CometMergeRows + +/** + * `MergeRowsExec` (the row-level MERGE dispatch operator) exists on Spark 3.5+; this registers it + * for Spark 3.5. See `org.apache.comet.serde.operator.CometMergeRows` for the conversion logic. + */ +object ShimCometMergeRows { + val nativeExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = + Map(classOf[MergeRowsExec] -> CometMergeRows) +} diff --git a/spark/src/main/spark-3.5/org/apache/comet/serde/operator/CometMergeRows.scala b/spark/src/main/spark-3.5/org/apache/comet/serde/operator/CometMergeRows.scala new file mode 100644 index 00000000000..eba6b39f74d --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/comet/serde/operator/CometMergeRows.scala @@ -0,0 +1,174 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.serde.operator + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.plans.logical.MergeRows +import org.apache.spark.sql.comet.{CometMergeRowsExec, SerializedPlan} +import org.apache.spark.sql.execution.datasources.v2.MergeRowsExec +import org.apache.spark.sql.types.LongType + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.ConfigEntry +import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.{MergeInstruction, MergeOutputRow, Operator} +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType} + +/** + * Serde for Spark's `MergeRowsExec` (Spark 3.5+). An instruction is encoded as a condition plus + * zero, one, or two output rows for Discard, Keep, or Split. + */ +object CometMergeRows extends CometOperatorSerde[MergeRowsExec] { + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_EXEC_MERGE_ROWS_ENABLED) + + override def getSupportLevel(op: MergeRowsExec): SupportLevel = { + if (!cardinalityCheckSatisfied(op)) { + Unsupported(Some(cardinalityCheckFallbackReason)) + } else { + Compatible(None) + } + } + + override def convert( + op: MergeRowsExec, + builder: Operator.Builder, + childOp: OperatorOuterClass.Operator*): Option[Operator] = { + val input = op.child.output + + def convertInstruction(instr: MergeRows.Instruction): Option[MergeInstruction] = { + val condition = exprToProto(instr.condition, input) + val outputs = instr.outputs.map { row => + val exprs = row.map(exprToProto(_, input)) + if (exprs.forall(_.isDefined)) { + Some(MergeOutputRow.newBuilder().addAllExprs(exprs.map(_.get).asJava).build()) + } else { + None + } + } + + if (condition.isDefined && outputs.forall(_.isDefined)) { + Some( + MergeInstruction + .newBuilder() + .setCondition(condition.get) + .addAllOutputs(outputs.map(_.get).asJava) + .build()) + } else { + None + } + } + + val matched = op.matchedInstructions.map(convertInstruction) + val notMatched = op.notMatchedInstructions.map(convertInstruction) + val notMatchedBySource = op.notMatchedBySourceInstructions.map(convertInstruction) + + val isSourcePresent = exprToProto(op.isSourceRowPresent, input) + val isTargetPresent = exprToProto(op.isTargetRowPresent, input) + + val outputTypes = op.output.map(a => serializeDataType(a.dataType)) + + // Only wired when Spark asked for a cardinality check; if `checkCardinality` is false this + // must stay unset rather than default to 0. `row_id_ordinal` is `optional int32` (not a plain + // `int32`) precisely because 0 is a legitimate ordinal -- the row-id column could be the + // child's first column -- so presence must be distinguishable from a real value. + val rowIdOrd: Option[Int] = if (op.checkCardinality) rowIdOrdinal(op) else None + + // `childOp` is empty when the child did not itself convert to a native operator -- + // `CometExecRule.convertToComet` still calls `convert` in that case. Without this guard a + // childless `MergeRows` reaches the planner, where a missing child would otherwise be an + // invalid native plan rather than a clean JVM fallback. + if (childOp.nonEmpty && matched.forall(_.isDefined) && notMatched.forall(_.isDefined) && + notMatchedBySource.forall(_.isDefined) && isSourcePresent.isDefined && + isTargetPresent.isDefined && outputTypes.forall(_.isDefined) && + cardinalityCheckSatisfied(op)) { + val mergeBuilder = OperatorOuterClass.MergeRows + .newBuilder() + .setIsSourceRowPresent(isSourcePresent.get) + .setIsTargetRowPresent(isTargetPresent.get) + .addAllMatchedInstructions(matched.map(_.get).asJava) + .addAllNotMatchedInstructions(notMatched.map(_.get).asJava) + .addAllNotMatchedBySourceInstructions(notMatchedBySource.map(_.get).asJava) + .addAllOutputTypes(outputTypes.map(_.get).asJava) + rowIdOrd.foreach(mergeBuilder.setRowIdOrdinal) + Some(builder.setMergeRows(mergeBuilder).build()) + } else if (childOp.isEmpty) { + withFallbackReason(op, "No child operator") + None + } else if (!cardinalityCheckSatisfied(op)) { + withFallbackReason(op, cardinalityCheckFallbackReason) + None + } else { + withFallbackReason(op, "Unsupported expression in MERGE instructions") + None + } + } + + override def createExec(nativeOp: Operator, op: MergeRowsExec): CometMergeRowsExec = { + // Carry each instruction group and predicate across as its own field. `MergeRows.Instruction` + // is an `Expression`, so `Seq[Instruction]` widens to `Seq[Expression]` and Catalyst's + // expression machinery (subquery registration in particular) still sees every instruction -- + // see `CometMergeRowsExec`'s scaladoc for why the groups must not be flattened together. + CometMergeRowsExec( + nativeOp, + op, + op.output, + op.isSourceRowPresent, + op.isTargetRowPresent, + op.matchedInstructions.map(i => i: Expression), + op.notMatchedInstructions.map(i => i: Expression), + op.notMatchedBySourceInstructions.map(i => i: Expression), + op.checkCardinality, + if (op.checkCardinality) rowIdOrdinal(op) else None, + op.child, + SerializedPlan(None)) + } + + private val cardinalityCheckFallbackReason: String = + s"MERGE cardinality check requires a resolvable, Long-typed '${MergeRows.ROW_ID}' column" + + /** + * True iff cardinality checking is off, or the target row-id column resolves to a usable + * ordinal. Shared by `getSupportLevel` (the planning-time gate) and `convert` (which must + * re-derive the same condition to pick its own fallback branch) so the two checks cannot desync + * -- `convert` is only reached after `getSupportLevel` already passed, so its branch for this + * case is otherwise unreachable and would silently stop guarding anything if the two drifted + * apart. + */ + private def cardinalityCheckSatisfied(op: MergeRowsExec): Boolean = + !op.checkCardinality || rowIdOrdinal(op).isDefined + + /** + * Locates the ordinal of Spark's target row-id column (`MergeRows.ROW_ID`) in the child output, + * using the same SQLConf resolver as upstream `MergeRowsExec.references`. Requiring `LongType` + * keeps the native `Int64Array` cardinality validator type-safe; an unexpected schema falls + * back during planning instead of failing inside native execution. + */ + private def rowIdOrdinal(op: MergeRowsExec): Option[Int] = { + val idx = op.child.output.indexWhere { attr => + op.conf.resolver(attr.name, MergeRows.ROW_ID) && attr.dataType == LongType + } + if (idx >= 0) Some(idx) else None + } +} diff --git a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 6b976e55de1..55c5b05c802 100644 --- a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -275,6 +275,9 @@ trait ShimSparkErrorConverter { // multipleRowScalarSubqueryError was renamed to multipleRowSubqueryError in Spark 3.x Some(QueryExecutionErrors.multipleRowSubqueryError(sqlCtx(context))) + case "MergeCardinalityViolation" => + Some(QueryExecutionErrors.mergeCardinalityViolationError()) + case "IntervalArithmeticOverflowWithSuggestion" => // Spark 3.x uses a single intervalArithmeticOverflowError method Some( diff --git a/spark/src/main/spark-4.x/org/apache/comet/rules/shims/ShimCometMergeRows.scala b/spark/src/main/spark-4.x/org/apache/comet/rules/shims/ShimCometMergeRows.scala new file mode 100644 index 00000000000..9aa6b839d4c --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/rules/shims/ShimCometMergeRows.scala @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.rules.shims + +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.MergeRowsExec + +import org.apache.comet.serde.CometOperatorSerde +import org.apache.comet.serde.operator.CometMergeRows + +/** + * `MergeRowsExec` (the row-level MERGE dispatch operator) exists on Spark 3.5+; this registers it + * for Spark 4.x. See `org.apache.comet.serde.operator.CometMergeRows` for the conversion logic. + */ +object ShimCometMergeRows { + val nativeExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = + Map(classOf[MergeRowsExec] -> CometMergeRows) +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/serde/operator/CometMergeRows.scala b/spark/src/main/spark-4.x/org/apache/comet/serde/operator/CometMergeRows.scala new file mode 100644 index 00000000000..eba6b39f74d --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/serde/operator/CometMergeRows.scala @@ -0,0 +1,174 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.serde.operator + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.plans.logical.MergeRows +import org.apache.spark.sql.comet.{CometMergeRowsExec, SerializedPlan} +import org.apache.spark.sql.execution.datasources.v2.MergeRowsExec +import org.apache.spark.sql.types.LongType + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.ConfigEntry +import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.{MergeInstruction, MergeOutputRow, Operator} +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType} + +/** + * Serde for Spark's `MergeRowsExec` (Spark 3.5+). An instruction is encoded as a condition plus + * zero, one, or two output rows for Discard, Keep, or Split. + */ +object CometMergeRows extends CometOperatorSerde[MergeRowsExec] { + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_EXEC_MERGE_ROWS_ENABLED) + + override def getSupportLevel(op: MergeRowsExec): SupportLevel = { + if (!cardinalityCheckSatisfied(op)) { + Unsupported(Some(cardinalityCheckFallbackReason)) + } else { + Compatible(None) + } + } + + override def convert( + op: MergeRowsExec, + builder: Operator.Builder, + childOp: OperatorOuterClass.Operator*): Option[Operator] = { + val input = op.child.output + + def convertInstruction(instr: MergeRows.Instruction): Option[MergeInstruction] = { + val condition = exprToProto(instr.condition, input) + val outputs = instr.outputs.map { row => + val exprs = row.map(exprToProto(_, input)) + if (exprs.forall(_.isDefined)) { + Some(MergeOutputRow.newBuilder().addAllExprs(exprs.map(_.get).asJava).build()) + } else { + None + } + } + + if (condition.isDefined && outputs.forall(_.isDefined)) { + Some( + MergeInstruction + .newBuilder() + .setCondition(condition.get) + .addAllOutputs(outputs.map(_.get).asJava) + .build()) + } else { + None + } + } + + val matched = op.matchedInstructions.map(convertInstruction) + val notMatched = op.notMatchedInstructions.map(convertInstruction) + val notMatchedBySource = op.notMatchedBySourceInstructions.map(convertInstruction) + + val isSourcePresent = exprToProto(op.isSourceRowPresent, input) + val isTargetPresent = exprToProto(op.isTargetRowPresent, input) + + val outputTypes = op.output.map(a => serializeDataType(a.dataType)) + + // Only wired when Spark asked for a cardinality check; if `checkCardinality` is false this + // must stay unset rather than default to 0. `row_id_ordinal` is `optional int32` (not a plain + // `int32`) precisely because 0 is a legitimate ordinal -- the row-id column could be the + // child's first column -- so presence must be distinguishable from a real value. + val rowIdOrd: Option[Int] = if (op.checkCardinality) rowIdOrdinal(op) else None + + // `childOp` is empty when the child did not itself convert to a native operator -- + // `CometExecRule.convertToComet` still calls `convert` in that case. Without this guard a + // childless `MergeRows` reaches the planner, where a missing child would otherwise be an + // invalid native plan rather than a clean JVM fallback. + if (childOp.nonEmpty && matched.forall(_.isDefined) && notMatched.forall(_.isDefined) && + notMatchedBySource.forall(_.isDefined) && isSourcePresent.isDefined && + isTargetPresent.isDefined && outputTypes.forall(_.isDefined) && + cardinalityCheckSatisfied(op)) { + val mergeBuilder = OperatorOuterClass.MergeRows + .newBuilder() + .setIsSourceRowPresent(isSourcePresent.get) + .setIsTargetRowPresent(isTargetPresent.get) + .addAllMatchedInstructions(matched.map(_.get).asJava) + .addAllNotMatchedInstructions(notMatched.map(_.get).asJava) + .addAllNotMatchedBySourceInstructions(notMatchedBySource.map(_.get).asJava) + .addAllOutputTypes(outputTypes.map(_.get).asJava) + rowIdOrd.foreach(mergeBuilder.setRowIdOrdinal) + Some(builder.setMergeRows(mergeBuilder).build()) + } else if (childOp.isEmpty) { + withFallbackReason(op, "No child operator") + None + } else if (!cardinalityCheckSatisfied(op)) { + withFallbackReason(op, cardinalityCheckFallbackReason) + None + } else { + withFallbackReason(op, "Unsupported expression in MERGE instructions") + None + } + } + + override def createExec(nativeOp: Operator, op: MergeRowsExec): CometMergeRowsExec = { + // Carry each instruction group and predicate across as its own field. `MergeRows.Instruction` + // is an `Expression`, so `Seq[Instruction]` widens to `Seq[Expression]` and Catalyst's + // expression machinery (subquery registration in particular) still sees every instruction -- + // see `CometMergeRowsExec`'s scaladoc for why the groups must not be flattened together. + CometMergeRowsExec( + nativeOp, + op, + op.output, + op.isSourceRowPresent, + op.isTargetRowPresent, + op.matchedInstructions.map(i => i: Expression), + op.notMatchedInstructions.map(i => i: Expression), + op.notMatchedBySourceInstructions.map(i => i: Expression), + op.checkCardinality, + if (op.checkCardinality) rowIdOrdinal(op) else None, + op.child, + SerializedPlan(None)) + } + + private val cardinalityCheckFallbackReason: String = + s"MERGE cardinality check requires a resolvable, Long-typed '${MergeRows.ROW_ID}' column" + + /** + * True iff cardinality checking is off, or the target row-id column resolves to a usable + * ordinal. Shared by `getSupportLevel` (the planning-time gate) and `convert` (which must + * re-derive the same condition to pick its own fallback branch) so the two checks cannot desync + * -- `convert` is only reached after `getSupportLevel` already passed, so its branch for this + * case is otherwise unreachable and would silently stop guarding anything if the two drifted + * apart. + */ + private def cardinalityCheckSatisfied(op: MergeRowsExec): Boolean = + !op.checkCardinality || rowIdOrdinal(op).isDefined + + /** + * Locates the ordinal of Spark's target row-id column (`MergeRows.ROW_ID`) in the child output, + * using the same SQLConf resolver as upstream `MergeRowsExec.references`. Requiring `LongType` + * keeps the native `Int64Array` cardinality validator type-safe; an unexpected schema falls + * back during planning instead of failing inside native execution. + */ + private def rowIdOrdinal(op: MergeRowsExec): Option[Int] = { + val idx = op.child.output.indexWhere { attr => + op.conf.resolver(attr.name, MergeRows.ROW_ID) && attr.dataType == LongType + } + if (idx >= 0) Some(idx) else None + } +} diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 7fb822b66bd..ef29e185d0a 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -293,6 +293,9 @@ trait ShimSparkErrorConverter { case "ScalarSubqueryTooManyRows" => Some(QueryExecutionErrors.multipleRowScalarSubqueryError(context.headOption.orNull)) + case "MergeCardinalityViolation" => + Some(QueryExecutionErrors.mergeCardinalityViolationError()) + case "IntervalArithmeticOverflowWithSuggestion" => Some( QueryExecutionErrors.withSuggestionIntervalArithmeticOverflowError( diff --git a/spark/src/test/scala/org/apache/comet/exec/CometMergeRowsSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometMergeRowsSuite.scala new file mode 100644 index 00000000000..8f50378dccd --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometMergeRowsSuite.scala @@ -0,0 +1,388 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.exec + +import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkThrowable} +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.SubqueryExpression +import org.apache.spark.sql.comet.CometMergeRowsExec +import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog +import org.apache.spark.sql.execution.{QueryExecution, ScalarSubquery} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus + +/** + * `CometMergeRowsExec` converts Spark's `MergeRowsExec` and the `MergeRows` logical node that + * feeds it, both defined in Spark core (`execution.datasources.v2` / `catalyst.plans.logical`), + * not in any connector module. Spark plans a `MergeRowsExec` for MERGE INTO against any + * `SupportsRowLevelOperations` V2 table using group-based (copy-on-write) planning, independent + * of which connector implements the table. + * + * This suite pins that contract against Spark's own `InMemoryRowLevelOperationTableCatalog` test + * catalog rather than Iceberg. `InMemoryRowLevelOperationTable` selects its write shape via the + * `supports-deltas` table property: unset (default `false`) plans through group-based + * `MergeRows`; `supports-deltas=true` plans a JVM `WriteDelta` whose child is *also* a + * `MergeRowsExec` (and, with `split-updates=true`, emits `Split` for update-as-delete+insert). + * Comet's bottom-up conversion makes that child eligible for native execution under the JVM + * write, so both shapes are exercised here (`deltaMergeCase`) and checked against a pure-Spark + * baseline. Broader native-write acceleration is tracked by umbrella issue #5122. See + * `CometIcebergWriteActionSuite` for MERGE INTO coverage against real Iceberg tables. + */ +class CometMergeRowsSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + private val catalog = "generic_rowlevel" + + override protected def sparkConf: SparkConf = { + super.sparkConf + .set(s"spark.sql.catalog.$catalog", classOf[InMemoryRowLevelOperationTableCatalog].getName) + .set("spark.sql.shuffle.partitions", "4") + } + + private def assumeMerge(): Unit = assume(isSpark35Plus, "MergeRowsExec requires Spark 3.5+") + + test("MERGE on a non-Iceberg SupportsRowLevelOperations table engages CometMergeRowsExec") { + assumeMerge() + val target = s"$catalog.default.rowlevel_target" + val source = s"$catalog.default.rowlevel_source" + + def resetTables(): Unit = { + sql(s"DROP TABLE IF EXISTS $target") + sql(s"DROP TABLE IF EXISTS $source") + sql(s"CREATE TABLE $target (id INT, region STRING, amount DOUBLE) USING parquet") + sql(s"CREATE TABLE $source (id INT, region STRING, amount DOUBLE) USING parquet") + sql( + s"INSERT INTO $target VALUES " + + (0 until 20).map(i => s"($i, 'r${i % 3}', ${i * 1.5})").mkString(", ")) + sql( + s"INSERT INTO $source VALUES " + + (10 until 30).map(i => s"($i, 's${i % 3}', ${i * 2.0})").mkString(", ")) + } + + val mergeSql = + s"""MERGE INTO $target t USING $source s ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.amount = s.amount, t.region = s.region + |WHEN NOT MATCHED THEN INSERT (id, region, amount) VALUES (s.id, s.region, s.amount) + |""".stripMargin + + val captured = scala.collection.mutable.ArrayBuffer[QueryExecution]() + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = + captured += qe + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = + () + } + spark.listenerManager.register(listener) + try { + resetTables() + captured.clear() + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_MERGE_ROWS_ENABLED.key -> "true") { + sql(mergeSql) + } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + // Tree-string rendering uses Spark's stripped nodeName ("CometMergeRows"), not the Scala + // class name ("CometMergeRowsExec"). + val engaged = captured.exists(_.executedPlan.toString.contains("CometMergeRows")) + val cometResult = + sql(s"SELECT id, region, amount FROM $target ORDER BY id").collect().map(_.toString) + + resetTables() + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql(mergeSql) + } + val sparkResult = + sql(s"SELECT id, region, amount FROM $target ORDER BY id").collect().map(_.toString) + + assert( + engaged, + "CometMergeRowsExec did not engage against a non-Iceberg SupportsRowLevelOperations " + + "table") + assert( + cometResult.toSeq == sparkResult.toSeq, + "native MergeRows output diverged from Spark's for a non-Iceberg source.\n" + + s"comet: ${cometResult.mkString(", ")}\nspark: ${sparkResult.mkString(", ")}") + } finally { + spark.listenerManager.unregister(listener) + } + } + + test( + "MERGE cardinality violation raises SparkRuntimeException, not a generic native exception") { + assumeMerge() + val target = s"$catalog.default.rowlevel_target_card" + val source = s"$catalog.default.rowlevel_source_card" + + // Two source rows both match target row id=1 -> MERGE_CARDINALITY_VIOLATION. + sql(s"DROP TABLE IF EXISTS $target") + sql(s"DROP TABLE IF EXISTS $source") + sql(s"CREATE TABLE $target (id INT, amount DOUBLE) USING parquet") + sql(s"CREATE TABLE $source (id INT, amount DOUBLE) USING parquet") + sql(s"INSERT INTO $target VALUES (1, 10.0)") + sql(s"INSERT INTO $source VALUES (1, 20.0), (1, 30.0)") + + val mergeSql = + s"""MERGE INTO $target t USING $source s ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.amount = s.amount + |""".stripMargin + + // `withSQLConf`'s block result isn't usable directly: under the Spark 3.5 / Scala 2.12 build + // its signature returns `Unit` regardless of the block's type (only the 4.x / Scala 2.13 + // build is generic), so `val x = withSQLConf(...) { expr }` silently infers `x: Unit` there. + // Mutate a `var` from inside the block instead, which is portable across both. + val failedExecutions = scala.collection.mutable.ArrayBuffer[QueryExecution]() + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = () + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = + failedExecutions += qe + } + spark.listenerManager.register(listener) + var cometEx: Throwable = null + try { + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_MERGE_ROWS_ENABLED.key -> "true") { + cometEx = intercept[Exception](sql(mergeSql).collect()) + } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + assert( + failedExecutions.exists(qe => + find(qe.executedPlan) { case _: CometMergeRowsExec => true; case _ => false }.nonEmpty), + "cardinality-error coverage must execute through CometMergeRowsExec, not Spark fallback") + } finally { + spark.listenerManager.unregister(listener) + } + var sparkEx: Throwable = null + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkEx = intercept[Exception](sql(mergeSql).collect()) + } + + // Both must be the same real Spark exception type/condition, not a generic + // CometNativeException/CometQueryExecutionException wrapping an opaque message. + val sparkRuntimeExceptionClass = "org.apache.spark.SparkRuntimeException" + assert( + sparkEx.getClass.getName == sparkRuntimeExceptionClass, + s"expected Spark's own baseline to be a $sparkRuntimeExceptionClass, got ${sparkEx.getClass}") + assert( + cometEx.getClass.getName == sparkRuntimeExceptionClass, + s"Comet's native MERGE cardinality violation must surface as a $sparkRuntimeExceptionClass " + + s"like Spark's own, got ${cometEx.getClass}: ${cometEx.getMessage}") + val cometCondition = cometEx.asInstanceOf[SparkThrowable].getErrorClass + assert( + cometCondition == "MERGE_CARDINALITY_VIOLATION", + s"expected error condition MERGE_CARDINALITY_VIOLATION, got $cometCondition") + } + + /** Runs `mergeSql` with Comet on and returns the `CometMergeRowsExec` it planned, if any. */ + private def runAndCaptureMergeExec(mergeSql: String): Option[CometMergeRowsExec] = { + val captured = scala.collection.mutable.ArrayBuffer[QueryExecution]() + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = + captured += qe + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = + () + } + spark.listenerManager.register(listener) + try { + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_MERGE_ROWS_ENABLED.key -> "true") { + sql(mergeSql) + } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + captured + .flatMap(qe => + find(qe.executedPlan) { case _: CometMergeRowsExec => true; case _ => false }) + .collectFirst { case m: CometMergeRowsExec => m } + } finally { + spark.listenerManager.unregister(listener) + } + } + + test("MERGE with a scalar-subquery assignment runs natively and matches Spark") { + assumeMerge() + val target = s"$catalog.default.subq_target" + val source = s"$catalog.default.subq_source" + + def resetTables(): Unit = { + sql(s"DROP TABLE IF EXISTS $target") + sql(s"DROP TABLE IF EXISTS $source") + sql(s"CREATE TABLE $target (id INT, amount DOUBLE) USING parquet") + sql(s"CREATE TABLE $source (id INT, amount DOUBLE) USING parquet") + sql(s"INSERT INTO $target VALUES (1, 10.0), (2, 20.0), (3, 30.0)") + sql(s"INSERT INTO $source VALUES (2, 200.0), (3, 300.0), (4, 400.0)") + } + + // Spark forbids subqueries in MERGE *conditions* but allows them in assignment values. The + // subquery lives on a `MergeRows.Instruction` (an Expression); `CometMergeRowsExec` must retain + // it as a node expression so `CometNativeExec.prepareSubqueries` registers it, otherwise + // execution fails with `Subquery ... not found for plan ...`. + val mergeSql = + s"""MERGE INTO $target t USING $source s ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.amount = (SELECT max(amount) FROM $source) + |WHEN NOT MATCHED THEN INSERT (id, amount) VALUES (s.id, s.amount) + |""".stripMargin + + resetTables() + val mergeExec = runAndCaptureMergeExec(mergeSql) + assert( + mergeExec.isDefined, + "expected the subquery MERGE to still run through CometMergeRowsExec") + // The node must expose the assignment's subquery so `CometNativeExec.prepareSubqueries` + // registers it -- either the planned `execution.ScalarSubquery` or, once AQE has run, a + // subquery-bearing expression tree. Absence here is what previously caused `Subquery not found`. + assert( + mergeExec.get.expressions.exists(e => + e.exists(_.isInstanceOf[ScalarSubquery]) || SubqueryExpression.hasSubquery(e)), + "CometMergeRowsExec must expose the assignment subquery in `expressions`, got: " + + mergeExec.get.expressions.mkString("; ")) + val cometResult = + sql(s"SELECT id, amount FROM $target ORDER BY id").collect().map(_.toString).toSeq + + resetTables() + withSQLConf(CometConf.COMET_ENABLED.key -> "false")(sql(mergeSql)) + val sparkResult = + sql(s"SELECT id, amount FROM $target ORDER BY id").collect().map(_.toString).toSeq + + assert( + cometResult == sparkResult, + "native subquery MERGE diverged from Spark.\n" + + s"comet: ${cometResult.mkString(", ")}\nspark: ${sparkResult.mkString(", ")}") + } + + // Executes a delta-backed MERGE with Comet enabled and requires the MergeRows child below the + // JVM WriteDelta to convert natively. The final table contents are then compared with a pure + // Spark run of the same statement. For split updates we also assert that Spark actually planned + // a Split instruction, so the test cannot silently stop exercising delete+reinsert semantics. + private def deltaMergeCase(name: String, splitUpdates: Boolean, matchedClause: String): Unit = { + test(s"delta-backed MERGE: $name") { + assumeMerge() + val target = s"$catalog.default.delta_t_${name.replaceAll("\\W", "_")}" + val source = s"$catalog.default.delta_s_${name.replaceAll("\\W", "_")}" + + def reset(): Unit = { + sql(s"DROP TABLE IF EXISTS $target") + sql(s"DROP TABLE IF EXISTS $source") + sql( + s"CREATE TABLE $target (pk INT NOT NULL, amount INT) USING parquet " + + "TBLPROPERTIES ('supports-deltas' = 'true'" + + (if (splitUpdates) ", 'split-updates' = 'true')" else ")")) + sql(s"CREATE TABLE $source (pk INT NOT NULL, amount INT) USING parquet") + sql(s"INSERT INTO $target VALUES (1, 10), (2, 20), (3, 30)") + sql(s"INSERT INTO $source VALUES (2, 200), (3, 300), (4, 400)") + } + + val mergeSql = + s"""MERGE INTO $target t USING $source s ON t.pk = s.pk + |$matchedClause + |WHEN NOT MATCHED THEN INSERT (pk, amount) VALUES (s.pk, s.amount) + |""".stripMargin + + reset() + val mergeExec = runAndCaptureMergeExec(mergeSql) + assert( + mergeExec.isDefined, + s"delta MERGE '$name' must execute its MergeRows child natively under WriteDelta") + if (splitUpdates) { + assert( + mergeExec.get.matchedInstructions.exists(_.getClass.getSimpleName == "Split"), + s"delta MERGE '$name' must retain Spark's Split instruction for delete+reinsert") + } + val cometRows = + sql(s"SELECT pk, amount FROM $target ORDER BY pk").collect().map(_.toString).toSeq + + reset() + withSQLConf(CometConf.COMET_ENABLED.key -> "false")(sql(mergeSql)) + val sparkRows = + sql(s"SELECT pk, amount FROM $target ORDER BY pk").collect().map(_.toString).toSeq + + assert( + cometRows == sparkRows, + s"native delta MERGE '$name' diverged from Spark.\n" + + s"comet: ${cometRows.mkString(", ")}\nspark: ${sparkRows.mkString(", ")}") + } + } + + deltaMergeCase( + "matched update", + splitUpdates = false, + "WHEN MATCHED THEN UPDATE SET t.amount = s.amount") + deltaMergeCase( + "split update (delete + reinsert)", + splitUpdates = true, + "WHEN MATCHED THEN UPDATE SET t.amount = s.amount") + deltaMergeCase("matched delete", splitUpdates = false, "WHEN MATCHED THEN DELETE") + + test("CometMergeRowsExec semantic identity reflects instruction groups and cardinality flag") { + assumeMerge() + val target = s"$catalog.default.eq_target" + val source = s"$catalog.default.eq_source" + sql(s"DROP TABLE IF EXISTS $target") + sql(s"DROP TABLE IF EXISTS $source") + sql(s"CREATE TABLE $target (id INT, amount DOUBLE) USING parquet") + sql(s"CREATE TABLE $source (id INT, amount DOUBLE) USING parquet") + sql(s"INSERT INTO $source VALUES (1, 100.0), (3, 300.0)") + + def merge(clauses: String): CometMergeRowsExec = { + sql(s"DROP TABLE IF EXISTS $target") + sql(s"CREATE TABLE $target (id INT, amount DOUBLE) USING parquet") + sql(s"INSERT INTO $target VALUES (1, 10.0), (2, 20.0)") + val m = + runAndCaptureMergeExec(s"MERGE INTO $target t USING $source s ON t.id = s.id\n$clauses\n") + assert(m.isDefined) + m.get + } + + val update = merge("""WHEN MATCHED THEN UPDATE SET t.amount = s.amount + |WHEN NOT MATCHED THEN INSERT (id, amount) VALUES (s.id, s.amount) + |WHEN NOT MATCHED BY SOURCE THEN UPDATE SET t.amount = t.amount + 1""".stripMargin) + val delete = merge("""WHEN MATCHED THEN DELETE + |WHEN NOT MATCHED THEN INSERT (id, amount) VALUES (s.id, s.amount)""".stripMargin) + + // equal to a structural copy (hash contract), unequal to a different MERGE + val same = update.copy() + assert(update == same && update.hashCode() == same.hashCode()) + assert(update != delete, "different MERGE instructions must not compare equal") + assert( + update != update.copy(checkCardinality = !update.checkCardinality), + "nodes differing only in checkCardinality must not compare equal") + + // The dangerous case: identical flattened expressions, different group partitioning. Moving + // the MATCHED instruction into the NOT MATCHED BY SOURCE group leaves `expressions` (and so a + // naive flat-list equality) unchanged, but the operator means something entirely different. + // Move the last MATCHED instruction to the front of the (adjacent) NOT MATCHED group so the + // flattened expression order is byte-for-byte unchanged; only the group boundary shifts. + assume(update.matchedInstructions.nonEmpty) + val shuffled = update.copy( + matchedInstructions = update.matchedInstructions.dropRight(1), + notMatchedInstructions = + update.matchedInstructions.takeRight(1) ++ update.notMatchedInstructions) + assert( + update.expressions.toList == shuffled.expressions.toList, + "test premise: the two nodes must have the same flattened expressions") + assert(update != shuffled, "instruction-group boundaries must be part of equality") + assert( + update.canonicalized != shuffled.canonicalized, + "instruction-group boundaries must survive canonicalization") + } +}