diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index def2816607..ca247643a0 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -337,6 +337,7 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + 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 bac1dcd15b..8a79f0d8d9 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -153,6 +153,7 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + 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 13aa6d943c..a418840cb3 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -74,6 +74,31 @@ incorrect result. When any single window expression in a `WindowExec` falls back `WindowGroupLimitExec` (window-based limit pushdown) is not yet supported and falls back to Spark ([#4837](https://github.com/apache/datafusion-comet/issues/4837)). +## 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 42235b1367..1fad6d828e 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 81d095658e..fb70161698 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 f798b80335..aaa0c82ad4 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -271,6 +271,7 @@ fn op_name(op: &OpStruct) -> &'static str { OpStruct::ShuffleScan(_) => "ShuffleScan", OpStruct::BroadcastNestedLoopJoin(_) => "BroadcastNestedLoopJoin", OpStruct::Sample(_) => "Sample", + OpStruct::MergeRows(_) => "MergeRows", } } 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 0000000000..2a9585ef89 --- /dev/null +++ b/native/core/src/execution/operators/merge_rows.rs @@ -0,0 +1,1058 @@ +// 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, RecordBatchOptions}; +use arrow::compute::kernels::boolean::{and, and_not}; +use arrow::compute::{filter_record_batch, prep_null_mask_filter}; +use arrow::datatypes::SchemaRef; +use datafusion::common::{DataFusionError, 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::SparkError; +use futures::{Stream, StreamExt}; +use std::collections::HashSet; +use std::{ + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; + +/// One `MergeRows.Instruction` (Keep / Discard / Split), expressed uniformly as a gating +/// condition plus zero, one, or two output row projections -- matching Spark's real +/// `condition: Expression, outputs: Seq[Seq[Expression]]` shape (Discard has zero output +/// projections, Keep has one, Split has two). +#[derive(Debug, Clone)] +pub struct MergeInstructionExec { + pub condition: Arc, + pub outputs: Vec>>, +} + +/// Configuration shared by `MergeRowsExec` and its `MergeRowsStream`: the row-presence +/// predicates, the three per-group instruction lists, and (when Spark's +/// `MergeRowsExec.checkCardinality` is on) the target row-id column's ordinal in the child +/// schema. Bundled into one struct, held behind an `Arc`, so `with_new_children` and `execute` +/// each clone one reference instead of threading seven fields by hand. +#[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 { + /// `row_id_ordinal` indexes directly into a child batch's columns, so a value out of range + /// for `child`'s schema would panic inside `check_cardinality` on the first batch. Called + /// from both `try_new` and `with_new_children`, since the latter can swap in a child whose + /// schema differs from the one this config was originally validated against. + fn validate(&self, child: &Arc) -> Result<(), DataFusionError> { + if let Some(ordinal) = self.row_id_ordinal { + 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" + ))); + } + } + Ok(()) + } +} + +/// A Comet native operator that reproduces Spark's `MergeRowsExec` (the row-level MERGE +/// dispatch operator introduced to Spark core in Iceberg 1.4.0 / SPARK-52403). Sits between the +/// target/source join and the write, deciding per row whether it becomes a kept row, is +/// discarded (a copy-on-write delete), or is split into two output rows. +#[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 = Arc::clone(&children[0]); + // Re-validate: an optimizer pass replacing the child here could hand back a schema the + // row-id ordinal no longer fits, and this path bypasses `try_new` entirely otherwise. + 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), + // One `seen` set per partition, created here and threaded through every batch this + // stream polls -- see the field doc on `MergeRowsStream::seen` for why it must not + // be reset per batch. + 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, +} + +/// Conservative per-entry cost of `seen`. hashbrown stores an 8-byte key plus a 1-byte control +/// slot at a ~87.5% load factor (~10.3 bytes/element) and doubles its table on growth; 16 bytes +/// per entry covers both without needing to observe the actual capacity. +const SEEN_ENTRY_BYTES: usize = 16; + +/// 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())?); + } + let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); + RecordBatch::try_new_with_options(Arc::clone(schema), columns, &options).map_err(|e| e.into()) +} + +/// 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") +/// via a shrinking `remaining` mask. +/// +/// 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 group_batch = filter_or_pass_through(batch, group_mask)?; + let mut remaining = BooleanArray::from(vec![true; group_batch.num_rows()]); + let mut out = Vec::new(); + let last = instructions.len() - 1; + + for (idx, instr) in instructions.iter().enumerate() { + if remaining.true_count() == 0 { + // Every later instruction's condition would AND against an all-false mask; nothing + // left in this group can fire. + break; + } + + // 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 (and then AND-ing against) a same-value n-row array. + let fire = match instr.condition.evaluate(&group_batch)? { + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => remaining.clone(), + ColumnarValue::Scalar(ScalarValue::Boolean(Some(false) | None)) => continue, + value => { + let cond = value + .into_array(group_batch.num_rows())? + .as_any() + .downcast_ref::() + .map(null_to_false) + .ok_or_else(|| { + DataFusionError::Internal("MergeRows: expected boolean array".to_string()) + })?; + and(&remaining, &cond)? + } + }; + + if fire.true_count() > 0 { + let filtered = filter_or_pass_through(&group_batch, &fire)?; + for output_exprs in &instr.outputs { + out.push(project(&filtered, output_exprs, schema)?); + } + } + + if idx != last { + remaining = and_not(&remaining, &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()) + })?; + + let mut new_entries = 0usize; + 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.insert(id) { + return Err(DataFusionError::External(Box::new( + SparkError::MergeCardinalityViolation, + ))); + } + new_entries += 1; + } + + // `seen` grows for the lifetime of the partition and is unbounded in the number of matched + // target rows, so it must be visible to the memory pool -- otherwise a large MERGE grows + // native memory with nothing to push back on it. Accounted after the fact (rather than + // reserving the batch's row count up front and releasing the remainder) since the overshoot + // is bounded by one batch. + reservation.try_grow(new_entries * SEEN_ENTRY_BYTES)?; + 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; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::execution::memory_pool::{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() { + // 3 rows: matched, not-matched (insert-only), not-matched-by-source. + 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(); + // row 1 (matched, kept) and row 2 (not-matched, inserted); row 3 discarded. + 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 null_condition_falls_through_to_next_instruction() { + // Regression test for the NULL-propagation data-loss bug. Models the plan + // `RewriteMergeIntoTable` actually builds for + // MERGE ... WHEN MATCHED AND s.val > 100 THEN UPDATE ... + // namely a two-instruction matched group whose second entry is the appended catch-all + // `Keep(TrueLiteral, target.output)`. With `val` NULL the first condition evaluates to + // NULL, which Spark treats as `false` and falls through to the catch-all; before the + // `null_to_false` normalization Arrow's NULL-propagating `and`/`not` poisoned the + // `remaining` mask and the row vanished from the output entirely. + 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() { + // The presence flags feed the same mask arithmetic as clause conditions, so a nullable + // `__row_from_source` / `__row_from_target` must also collapse to `false` rather than + // NULL -- otherwise the row falls out of all three group masks and is dropped. + 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(); + // source NULL -> false, target true => not-matched-by-source group. + 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() { + // Spark reaches `applyInstructions` only after a row is routed to a group, so a NOT + // MATCHED condition never sees a matched row. Row 1 is matched with `val = 0`; row 2 is + // the only not-matched row. Evaluating the not-matched condition `10 / val > 1` over the + // whole batch (the pre-fix behaviour) divides by row 1's zero and fails the query with an + // error Spark would never raise. + 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()]], + }; + // Guard the test's own premise: evaluated over the whole batch this condition really + // does fail, so a regression back to batch-wide evaluation cannot slip through silently. + 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(); + // row 1 kept by the matched group; row 2 kept by the not-matched group (10 / 5 > 1). + assert_eq!(got, vec![0, 5]); + } + + #[test] + fn cardinality_state_is_accounted_to_the_memory_pool() { + let mut reservation = test_reservation(); + let batch = test_batch(vec![1, 2], vec![10, 20], vec![true, true], vec![true, true]); + let matched_mask = BooleanArray::from(vec![true, true]); + check_cardinality( + &batch, + &matched_mask, + 0, + &mut HashSet::new(), + &mut reservation, + ) + .unwrap(); + assert_eq!( + reservation.size(), + 2 * SEEN_ENTRY_BYTES, + "`seen` must be visible to the memory pool so an unbounded MERGE has something \ + pushing back on it" + ); + } + + #[test] + fn cardinality_violation_detected() { + // Same target row_id (1) matched twice -> must error. + 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(); + // Asserting the ordered values (not just num_rows()) catches a regression that applies + // outputs[0] twice instead of outputs[0] then outputs[1]. + assert_eq!(got, vec![1, 2]); + } + + #[test] + fn cardinality_violation_detected_for_null_row_ids() { + // Two matched rows both carrying a NULL row id must trip the cardinality check, mirroring + // Spark's `BitmapCardinalityValidator`, which reads `InternalRow.getLong(ordinal)` + // unconditionally (no null check) -- a null long field reads as 0 on the Spark side, so + // two nulls collide exactly like two real zeros would. + 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(), + "two matched rows with a NULL row id must trip MERGE_CARDINALITY_VIOLATION, not be \ + silently skipped" + ); + assert!(result + .unwrap_err() + .to_string() + .contains("MERGE_CARDINALITY_VIOLATION")); + } + + /// A batch whose rows are all discarded must not surface as a zero-row batch: the stream + /// swallows it and pulls the next input instead, so downstream stages (the FFI hop into the + /// write pipeline, and `RecordBatchPartitionSplitter` for a partitioned table) never pay for + /// a batch with nothing in it. + #[tokio::test] + async fn all_discarded_batch_is_not_emitted() { + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::prelude::SessionContext; + + // Batch 1: matched-only rows, all discarded. Batch 2: one row that survives. + 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, + "the all-discarded batch must be swallowed, not forwarded as a zero-row batch" + ); + assert_eq!(batches[0].num_rows(), 1); + } + + /// A `row_id_ordinal` that does not address a real child column must be rejected at plan + /// construction. Reaching `check_cardinality` with it would panic on an out-of-bounds column + /// access instead of failing the query with a message. + #[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![], + // `test_schema()` has 4 columns, so 99 cannot be a row-id column. + Some(99), + source, + out_schema(), + ) + .unwrap_err(); + assert!( + err.to_string().contains("row id ordinal"), + "expected an out-of-range ordinal error, got: {err}" + ); + } + + /// `with_new_children` must re-run the same bounds check `try_new` does, not just copy the + /// old ordinal onto a new child. Constructs a valid `MergeRowsExec` against a 4-column child, + /// then swaps in a 1-column child the ordinal cannot possibly fit. + #[test] + fn with_new_children_rejects_ordinal_out_of_range_for_new_child() { + 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(3), // valid: test_schema() has 4 columns. + source, + out_schema(), + ) + .unwrap(), + ); + + let narrow_schema = Arc::new(Schema::new(vec![Field::new( + "only_col", + 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("row id ordinal"), + "expected an out-of-range ordinal error, got: {err}" + ); + } + + #[test] + fn cardinality_violation_detected_across_batches() { + // Regression test for the per-batch `seen` reset bug: `process_batch` must accept the + // caller's `seen` set and mutate it in place so state accumulates across the whole + // partition. Two separate batches each carry one of two source rows matching the same + // target row_id=1 -- exactly the case a fresh-per-batch `HashSet` would miss. + 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(), + "first batch introduces row_id=1 and should not trip the cardinality check" + ); + + let second = process_batch( + batch2, + &config, + &mut seen, + &mut test_reservation(), + &out_schema(), + ); + assert!( + second.is_err(), + "second batch reuses row_id=1 from a *different* Arrow batch and must trip \ + MERGE_CARDINALITY_VIOLATION" + ); + assert!(second + .unwrap_err() + .to_string() + .contains("MERGE_CARDINALITY_VIOLATION")); + } +} diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index b9b2b0fbd7..ca25b6fff7 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -29,6 +29,8 @@ mod copy; mod expand; pub use expand::ExpandExec; 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 43f066e76c..d5e6e97b30 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -29,8 +29,8 @@ use crate::execution::{ expressions::list_positions::ListPositionsExpr, expressions::subquery::Subquery, operators::{ - ExecutionError, ExpandExec, ParquetCompression, ParquetWriterExec, SampleExec, ScanExec, - ShuffleScanExec, + ExecutionError, ExpandExec, MergeInstructionExec, MergeRowsExec, ParquetCompression, + ParquetWriterExec, SampleExec, ScanExec, ShuffleScanExec, }, planner::expression_registry::ExpressionRegistry, planner::operator_registry::OperatorRegistry, @@ -1896,6 +1896,96 @@ impl PhysicalPlanner { Arc::new(SparkPlan::new(spark_plan.plan_id, expand, vec![child])), )) } + OpStruct::MergeRows(merge) => { + assert_eq!(children.len(), 1); + let (scans, shuffle_scans, child) = + self.create_plan(&children[0], inputs, partition_count)?; + + let is_source_row_present = self.create_expr( + merge.is_source_row_present.as_ref().unwrap(), + child.schema(), + )?; + let is_target_row_present = self.create_expr( + merge.is_target_row_present.as_ref().unwrap(), + 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().unwrap(), 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 the output schema from the compiled projection expressions rather than + // Spark's declared `output_types`, mirroring `Expand` above: a Comet scan can hand + // up a dictionary-encoded string where Spark declares a plain `Utf8`, and + // `merge_rows.rs`'s `project()` builds its `RecordBatch` from the expression + // outputs, so a schema taken from the proto types could mismatch at runtime. + // Every Keep/Split instruction across all three groups produces the same-shaped + // row, so `ExpandExec::build_schema` -- built for exactly this "N same-shaped + // projections" case -- applies unchanged. Falls back to the declared types when + // every instruction is a Discard, since there is then no projection to read the + // types off. + 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 d84a71abc3..b07c07ecd4 100644 --- a/native/core/src/execution/planner/operator_registry.rs +++ b/native/core/src/execution/planner/operator_registry.rs @@ -152,6 +152,7 @@ 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 } } diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index ced87262f3..6e236f5f39 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -66,6 +66,7 @@ message Operator { ShuffleScan shuffle_scan = 116; BroadcastNestedLoopJoin broadcast_nested_loop_join = 117; Sample sample = 118; + MergeRows merge_rows = 119; } } @@ -488,6 +489,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 6e7bb4271b..3afd75bc02 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -245,6 +245,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_NATIVE_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.native.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 ef2f37371c..5c8def2bac 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -53,6 +53,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, ShimSubqueryBroadcast} @@ -71,7 +72,7 @@ object CometExecRule { * Fully native operators. */ val nativeExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = - Map( + (Map( classOf[ProjectExec] -> CometProjectExec, classOf[FilterExec] -> CometFilterExec, classOf[LocalLimitExec] -> CometLocalLimitExec, @@ -87,7 +88,7 @@ object CometExecRule { classOf[SortExec] -> CometSortExec, classOf[LocalTableScanExec] -> CometLocalTableScanExec, classOf[SampleExec] -> CometSampleExec, - classOf[WindowExec] -> CometWindowExec) + classOf[WindowExec] -> CometWindowExec) ++ ShimCometMergeRows.nativeExecs).toMap /** * 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 74ef3af91c..07667e74ad 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 @@ -1429,6 +1429,55 @@ case class CometExpandExec( override lazy val metrics: Map[String, SQLMetric] = Map.empty } +case class CometMergeRowsExec( + override val nativeOp: Operator, + override val originalPlan: SparkPlan, + override val output: Seq[Attribute], + child: SparkPlan, + override val serializedPlanOpt: SerializedPlan) + extends CometUnaryExec { + // Matches Spark's `MergeRowsExec`, which does not override `outputPartitioning` and so inherits + // `SparkPlan`'s `UnknownPartitioning(0)`. Reporting the child's partitioning instead would be + // *stronger* than Spark's contract and could let `EnsureRequirements` skip a shuffle the JVM + // plan takes, so this must stay in step with Spark rather than be "improved". + override def outputPartitioning: Partitioning = UnknownPartitioning(0) + + override def producedAttributes: AttributeSet = outputSet + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + this.copy(child = newChild) + + override def stringArgs: Iterator[Any] = Iterator(output, child) + + override def equals(obj: Any): Boolean = { + obj match { + case other: CometMergeRowsExec => + this.output == other.output && + this.child == other.child && + this.serializedPlanOpt == other.serializedPlanOpt + case _ => + false + } + } + + override def hashCode(): Int = Objects.hashCode(output, child) + + // `output_rows / output_batches` is this operator's average output batch size -- the shape + // whatever native write follows is fed, whose per-batch cost scales with column count rather + // than row count. + // + // Does not expose Spark 4.x real `MergeRowsExec`'s 8 per-clause row counters + // (numTargetRowsCopied/Inserted/Updated/Deleted/MatchedUpdated/MatchedDeleted/ + // NotMatchedBySourceUpdated/NotMatchedBySourceDeleted). Spark 3.5.x's `MergeRowsExec` has none + // of these either -- `MergeRows.Keep`'s `Context` field (Copy/Update/Insert/Delete), which + // Spark 4.x uses to attribute them, doesn't exist on 3.5. Wiring them correctly needs a + // version-gated native context tag threaded through the proto/native/JVM serde, not just a + // stub -- a metric reading 0 after rows were inserted is worse than the key being absent. + 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 0000000000..79057f0906 --- /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 0000000000..8e1c59dae6 --- /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 0000000000..656a4f032a --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/comet/serde/operator/CometMergeRows.scala @@ -0,0 +1,166 @@ +/* + * 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.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` (the row-level MERGE dispatch operator moved from Iceberg + * extensions into Spark core in Iceberg 1.4.0 / SPARK-52403). Only exists on Spark 3.5+, hence + * this file lives in a version-gated source root rather than the shared `serde/operator` package. + * + * Spark's real shape (verified against the 4.1.1 / 3.5.7 bytecode) is simpler than it might look: + * `isSourceRowPresent` / `isTargetRowPresent` are boolean-valued `Expression`s (not column + * ordinals), and each `MergeRows.Instruction` (Keep / Discard / Split) is uniformly a `condition: + * Expression` plus `outputs: Seq[Seq[Expression]]` -- zero output projections means Discard, one + * means Keep, two means Split. There is no separate instruction-kind enum on the Spark side, so + * none is introduced on the wire either. + */ +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 `assert_eq!(children.len(), 1)` panics + // instead of falling back to the JVM operator. + 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 = { + CometMergeRowsExec(nativeOp, op, op.output, 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 Iceberg's target row-id column (`MergeRows.ROW_ID`) in the child + * output, requiring it to be `LongType` since the native cardinality check reads it as an + * `Int64Array` and would otherwise fail at execution time instead of falling back to the JVM + * operator at planning time. `MergeRowsExec` tracks this internally (via its private + * `BitmapCardinalityValidator`) for the cardinality check but doesn't expose the ordinal + * publicly, so it's re-derived here from the well-known column name. + */ + private def rowIdOrdinal(op: MergeRowsExec): Option[Int] = { + val idx = + op.child.output.indexWhere(a => a.name == MergeRows.ROW_ID && a.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 6b976e55de..55c5b05c80 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 0000000000..9aa6b839d4 --- /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 0000000000..656a4f032a --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/serde/operator/CometMergeRows.scala @@ -0,0 +1,166 @@ +/* + * 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.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` (the row-level MERGE dispatch operator moved from Iceberg + * extensions into Spark core in Iceberg 1.4.0 / SPARK-52403). Only exists on Spark 3.5+, hence + * this file lives in a version-gated source root rather than the shared `serde/operator` package. + * + * Spark's real shape (verified against the 4.1.1 / 3.5.7 bytecode) is simpler than it might look: + * `isSourceRowPresent` / `isTargetRowPresent` are boolean-valued `Expression`s (not column + * ordinals), and each `MergeRows.Instruction` (Keep / Discard / Split) is uniformly a `condition: + * Expression` plus `outputs: Seq[Seq[Expression]]` -- zero output projections means Discard, one + * means Keep, two means Split. There is no separate instruction-kind enum on the Spark side, so + * none is introduced on the wire either. + */ +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 `assert_eq!(children.len(), 1)` panics + // instead of falling back to the JVM operator. + 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 = { + CometMergeRowsExec(nativeOp, op, op.output, 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 Iceberg's target row-id column (`MergeRows.ROW_ID`) in the child + * output, requiring it to be `LongType` since the native cardinality check reads it as an + * `Int64Array` and would otherwise fail at execution time instead of falling back to the JVM + * operator at planning time. `MergeRowsExec` tracks this internally (via its private + * `BitmapCardinalityValidator`) for the cardinality check but doesn't expose the ordinal + * publicly, so it's re-derived here from the well-known column name. + */ + private def rowIdOrdinal(op: MergeRowsExec): Option[Int] = { + val idx = + op.child.output.indexWhere(a => a.name == MergeRows.ROW_ID && a.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 7fb822b66b..ef29e185d0 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 0000000000..ded24ccf47 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometMergeRowsSuite.scala @@ -0,0 +1,173 @@ +/* + * 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.connector.catalog.InMemoryRowLevelOperationTableCatalog +import org.apache.spark.sql.execution.QueryExecution +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 `MergeRows`; `true` + * plans through `WriteDelta` instead, a distinct row-dispatch shape this operator does not + * target. 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. + var cometEx: Throwable = null + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_MERGE_ROWS_ENABLED.key -> "true") { + cometEx = intercept[Exception](sql(mergeSql).collect()) + } + 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") + } +}