sunchao commented on code in PR #5318: URL: https://github.com/apache/datafusion-comet/pull/5318#discussion_r3837598665
########## 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<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())?); + } + 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()) Review Comment: [P2] Normalize nested output types before stamping the batch Could this use `cast_and_stamp_schema`, as `ExpandStream::expand` does? `ExpandExec::build_schema` widens nested nullability across the instructions, but this direct stamp does not reconcile each projected array with that widened type. For a target column `payload STRUCT<n: BOOLEAN>`, `UPDATE SET payload = named_struct('n', s.d IS NULL)` produces a non-nullable inner `n`, while Spark's appended carryover `Keep` references the target's nullable `n`. The derived schema is therefore nullable, and `RecordBatch::try_new_with_options` rejects the update array with `column types must match schema types`. I reproduced that error using the exact-head `CreateNamedStruct`, schema builder, and `run_group`. The corresponding Spark 4.0.4 MERGE succeeds. Normalizing the array before stamping also succeeds and preserves the value. ########## 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<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())?); + } + 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<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") +/// 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<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 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)? { Review Comment: [P2] Evaluate later predicates only for the remaining rows Could you narrow to `remaining` before evaluating each clause condition? The mask is currently applied after `instr.condition.evaluate(&group_batch)`, so rows handled by an earlier clause still evaluate later predicates whenever another row remains. With ANSI enabled, two matched rows having `s.d = 0` and `2`, and clauses `WHEN MATCHED AND s.d = 0 THEN UPDATE SET v = 111` followed by `WHEN MATCHED AND 2 / s.d > 0 THEN UPDATE SET v = 222`, the second predicate divides by the already-handled row's zero. I ran this MERGE on Spark 4.0.4 and got `[1,111], [2,222]`. A two-row probe of this exact `run_group` fails with `DivideByZero`, while the same function succeeds row-at-a-time. Spark's `applyInstructions` returns immediately after the first matching clause, so this is a failure of an otherwise valid MERGE, not just a different error ordering. -- 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]
