sunchao commented on code in PR #5318: URL: https://github.com/apache/datafusion-comet/pull/5318#discussion_r3868461677
########## spark/src/main/scala/org/apache/spark/sql/comet/operators.scala: ########## @@ -1428,6 +1428,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) Review Comment: [P2] Retain MERGE assignment subqueries Could you preserve the instruction expressions as fields on this node, or fall back for scalar-subquery assignments? An assignment such as `WHEN MATCHED THEN UPDATE SET amount = (SELECT max(amount) FROM source)` is accepted by the scalar-subquery serializer, but Spark's expression traversal does not recurse into `originalPlan`. Consequently, `CometNativeExec.prepareSubqueries` and `collectSubqueries` cannot discover or register that subquery, leaving the native lookup without an entry (`Subquery ... not found for plan ...`). I verified that the corresponding Spark 4.1.3 MERGE succeeds with `[1,3]`; a probe compiled from this new class found one scalar subquery on the original `MergeRowsExec`, zero on `CometMergeRowsExec`, and zero returned by `collectSubqueries`. ########## native/core/src/execution/operators/merge_rows.rs: ########## @@ -0,0 +1,1164 @@ +// 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::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::{cast_and_stamp_schema, 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<dyn PhysicalExpr>, + pub outputs: Vec<Vec<Arc<dyn PhysicalExpr>>>, +} + +/// 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<dyn PhysicalExpr>, + is_target_row_present: Arc<dyn PhysicalExpr>, + matched_instructions: Vec<MergeInstructionExec>, + not_matched_instructions: Vec<MergeInstructionExec>, + not_matched_by_source_instructions: Vec<MergeInstructionExec>, + /// `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<usize>, +} + +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<dyn ExecutionPlan>) -> 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<MergeConfig>, + child: Arc<dyn ExecutionPlan>, + schema: SchemaRef, + cache: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +impl MergeRowsExec { + #[allow(clippy::too_many_arguments)] + pub fn try_new( + is_source_row_present: Arc<dyn PhysicalExpr>, + is_target_row_present: Arc<dyn PhysicalExpr>, + matched_instructions: Vec<MergeInstructionExec>, + not_matched_instructions: Vec<MergeInstructionExec>, + not_matched_by_source_instructions: Vec<MergeInstructionExec>, + row_id_ordinal: Option<usize>, + child: Arc<dyn ExecutionPlan>, + schema: SchemaRef, + ) -> Result<Self, DataFusionError> { + 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<dyn ExecutionPlan>> { + vec![&self.child] + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> { + 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<TaskContext>, + ) -> datafusion::common::Result<SendableRecordBatchStream> { + 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<PlanProperties> { + &self.cache + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn name(&self) -> &str { + "CometMergeRowsExec" + } +} + +pub struct MergeRowsStream { + config: Arc<MergeConfig>, + 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<i64>, + /// 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<dyn PhysicalExpr>, + batch: &RecordBatch, +) -> Result<BooleanArray, DataFusionError> { + let array: ArrayRef = expr.evaluate(batch)?.into_array(batch.num_rows())?; + array + .as_any() + .downcast_ref::<BooleanArray>() + .map(null_to_false) + .ok_or_else(|| DataFusionError::Internal("MergeRows: expected boolean array".to_string())) +} + +fn project( + batch: &RecordBatch, + exprs: &[Arc<dyn PhysicalExpr>], + schema: &SchemaRef, +) -> Result<RecordBatch, DataFusionError> { + 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<RecordBatch, DataFusionError> { + 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<Vec<RecordBatch>, 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::<BooleanArray>() + .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<i64>, + 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::<Int64Array>() + .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)?; Review Comment: [P2] Account for allocated hash-table capacity The fixed 16-byte charge per ID is not conservative immediately after `HashSet` growth. Calling this exact `check_cardinality` function with 917,505 distinct IDs in 4,096-row batches succeeded against a 16 MiB `GreedyMemoryPool`, while an allocator probe measured 18,874,384 bytes of live table allocation versus only 14,680,080 bytes reserved. Thus an 18 MiB persistent table is accepted by a 16 MiB pool. This uncharged capacity grows with partition size, so the overshoot is not bounded by the current batch and can bypass the memory budget, risking executor OOM. Could you account for allocated table capacity, for example using DataFusion's `estimate_memory_size`, instead of only the number of inserted IDs? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
