viirya commented on code in PR #7190: URL: https://github.com/apache/datafusion-comet/pull/7190#discussion_r4385725060
########## native/operators/src/membership_key_reduction.rs: ########## @@ -0,0 +1,302 @@ +// 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. + +//! Keep the first occurrence of each integer key in a duplicate-heavy shuffle batch. +//! A left-semi join only needs to know whether a right-side key exists. + +use std::fmt::Formatter; +use std::hash::Hash; +use std::sync::Arc; + +use arrow::array::{downcast_integer_array, Array, PrimitiveArray, UInt32Array}; +use arrow::compute::take_record_batch; +use arrow::datatypes::ArrowPrimitiveType; +use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::utils::memory::estimate_memory_size; +use datafusion::common::{DataFusionError, HashSet, Result}; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::{ + ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, +}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; +use futures::StreamExt; + +const PREFIX_ROWS: usize = 8192; +const MIN_BATCH_ROWS: usize = 1024; +const MAX_BATCH_ROWS: usize = 64 * 1024; + +#[derive(Debug)] +pub struct MembershipKeyReductionExec { + input: Arc<dyn ExecutionPlan>, + metrics: ExecutionPlanMetricsSet, +} + +impl MembershipKeyReductionExec { + /// The join planner supplies a fresh shuffle batch containing one signed integer key. + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + Self { + input, + metrics: ExecutionPlanMetricsSet::new(), + } + } +} + +impl DisplayAs for MembershipKeyReductionExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CometMembershipKeyReductionExec") + } +} + +impl ExecutionPlan for MembershipKeyReductionExec { + fn name(&self) -> &str { + "CometMembershipKeyReductionExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + self.input.properties() + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::LowerEqual + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + let [input]: [_; 1] = children.try_into().map_err(|_| { + DataFusionError::Internal("Membership key reduction requires one child".into()) + })?; + Ok(Arc::new(Self::new(input))) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let input = self.input.execute(partition, Arc::clone(&context))?; + let input_rows = + MetricBuilder::new(&self.metrics).counter("right_key_reduction_input_rows", partition); + let output_rows = + MetricBuilder::new(&self.metrics).counter("right_key_reduction_output_rows", partition); + let bypass_rows = + MetricBuilder::new(&self.metrics).counter("right_key_reduction_bypass_rows", partition); + let time = + MetricBuilder::new(&self.metrics).subset_time("right_key_reduction_time", partition); + let mut seen = 0usize; + let mut min_rows = MIN_BATCH_ROWS; + let stream = input.map(move |batch| { + let batch = batch?; + let rows = batch.num_rows(); + input_rows.add(rows); + // Avoid the setup cost on short partitions, including a single large batch. Review Comment: Thanks for the isolated control. The trade-off is small either way, and keeping the gate as a conservative policy for this version is fine with me. Resolving. ########## native/core/src/execution/planner/join.rs: ########## @@ -201,6 +206,64 @@ impl OperatorBuilder for HashJoinBuilder { /// Builder for `SortMergeJoin`. pub struct SortMergeJoinBuilder; +/// Spark identifies the shuffle input. Check the actual sort and its key before replacing its child. +fn reduce_before_sort( + planner: &PhysicalPlanner, + params: &mut JoinParameters, + sort_options: &[SortOptions], +) -> Result<Option<Arc<dyn ExecutionPlan>>, ExecutionError> { + if params.join_type != DFJoinType::LeftSemi || params.join_filter.is_some() { + return Ok(None); + } + let Some(sort) = params.right.native_plan.downcast_ref::<SortExec>() else { + return Ok(None); + }; + let input = sort.input(); + let ([(_, key)], [options]) = (params.join_on.as_slice(), sort_options) else { + return Ok(None); + }; + let schema = input.schema(); + let order = sort.expr().first(); + if !(input.is::<ScanExec>() || input.is::<ShuffleScanExec>()) + || schema.fields().len() != 1 + || !schema.field(0).data_type().is_signed_integer() + || !key.is::<Column>() + || sort.expr().len() != 1 + || order.options != *options + || !order.expr.eq(key) + { + return Ok(None); + } + // Check at first execution, after Spark has constructed the input iterators. Review Comment: Could you explain the mechanism this guards against? With the reservation gone, the reducer makes no Spark memory request, and its only effect is fewer rows reaching the sort. The description says that at the previously failing headroom both settings now write zero JVM spill bytes, but not why enabling reduction made a JVM consumer spill before. If the cause is how the fair pool's per-consumer limit interacts with Spark's `acquireExecutionMemory`, a comment here saying so would make it possible to judge later whether the guard is still needed. Two related points: - The check runs once, at the first execution. A JVM consumer that starts holding memory after that isn't covered, so the comment should describe this as a best-effort check. - `planner.task_context` is `None` in the native planner tests, so this block never runs there. In the SQL test it only ever returns `false`. ########## native/operators/src/membership_key_reduction.rs: ########## @@ -0,0 +1,316 @@ +// 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. + +//! Keep the first occurrence of each integer key in a duplicate-heavy shuffle batch. +//! A left-semi join only needs to know whether a right-side key exists. + +use std::fmt::Formatter; +use std::hash::Hash; +use std::sync::Arc; + +use arrow::array::{downcast_integer_array, Array, PrimitiveArray, UInt32Array}; +use arrow::compute::{take_record_batch, BatchCoalescer}; +use arrow::datatypes::ArrowPrimitiveType; +use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, HashSet, Result}; +use datafusion::execution::{async_try_stream, TaskContext}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; +use futures::{StreamExt, TryStreamExt}; + +const PREFIX_ROWS: usize = 8192; +const MIN_BATCH_ROWS: usize = 1024; +const MAX_BATCH_ROWS: usize = 64 * 1024; +const MIN_OUTPUT_ROWS: usize = 128; + +#[derive(Debug)] +pub struct MembershipKeyReductionExec { + input: Arc<dyn ExecutionPlan>, + metrics: ExecutionPlanMetricsSet, +} + +impl MembershipKeyReductionExec { + /// The join planner supplies a single signed integer key column. + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + Self { + input, + metrics: ExecutionPlanMetricsSet::new(), + } + } +} + +impl DisplayAs for MembershipKeyReductionExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CometMembershipKeyReductionExec") + } +} + +impl ExecutionPlan for MembershipKeyReductionExec { + fn name(&self) -> &str { + "CometMembershipKeyReductionExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + self.input.properties() + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::LowerEqual + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + let [input]: [_; 1] = children.try_into().map_err(|_| { + DataFusionError::Internal("Membership key reduction requires one child".into()) + })?; + Ok(Arc::new(Self::new(input))) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let input = self.input.execute(partition, context)?; + let input_rows = + MetricBuilder::new(&self.metrics).counter("right_key_reduction_input_rows", partition); + let output_rows = + MetricBuilder::new(&self.metrics).counter("right_key_reduction_output_rows", partition); + let bypass_rows = + MetricBuilder::new(&self.metrics).counter("right_key_reduction_bypass_rows", partition); + let time = + MetricBuilder::new(&self.metrics).subset_time("right_key_reduction_time", partition); + let mut seen = 0usize; + let mut min_rows = MIN_BATCH_ROWS; + let mut batches = input.map(move |batch| -> Result<_> { + let batch = batch?; + let rows = batch.num_rows(); + input_rows.add(rows); + // Avoid speculative hashing on short partitions, including a single large batch. + let attempt = seen >= PREFIX_ROWS && (min_rows..=MAX_BATCH_ROWS).contains(&rows); + seen = seen.saturating_add(rows); + if attempt { + let _timer = time.timer(); + if let Some(reduced) = try_reduce_batch(&batch)? { + return Ok((reduced.num_rows() < MIN_OUTPUT_ROWS, reduced)); + } + // A larger shuffle frame may still be worth reducing. With uniform + // frame sizes the rest of the partition goes straight to the sort. + min_rows = (rows + 1).next_power_of_two(); + } + bypass_rows.add(rows); + Ok((false, batch)) + }); + // The sort accounts for Arrow buffers, but each small RecordBatch also has + // metadata. Pack tiny reductions so this overhead stays small on long partitions. + let mut small = BatchCoalescer::new(self.schema(), MIN_OUTPUT_ROWS); + let stream = async_try_stream(|mut out| async move { + while let Some((pack, batch)) = batches.try_next().await? { + let forward = if pack { + small.push_batch(batch)?; + None + } else { + small.finish_buffered_batch()?; + Some(batch) + }; + for batch in [small.next_completed_batch(), forward] Review Comment: Each loop iteration pops at most one completed batch. That's correct today, because packed inputs are always under `MIN_OUTPUT_ROWS` and each iteration completes at most one batch. But the correctness depends on that invariant, which isn't stated anywhere. A `while let Some(batch) = small.next_completed_batch()` would drain everything at no extra cost and stays correct if the packing rule ever changes. ########## native/operators/src/membership_key_reduction.rs: ########## @@ -0,0 +1,316 @@ +// 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. + +//! Keep the first occurrence of each integer key in a duplicate-heavy shuffle batch. +//! A left-semi join only needs to know whether a right-side key exists. + +use std::fmt::Formatter; +use std::hash::Hash; +use std::sync::Arc; + +use arrow::array::{downcast_integer_array, Array, PrimitiveArray, UInt32Array}; +use arrow::compute::{take_record_batch, BatchCoalescer}; +use arrow::datatypes::ArrowPrimitiveType; +use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, HashSet, Result}; +use datafusion::execution::{async_try_stream, TaskContext}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; +use futures::{StreamExt, TryStreamExt}; + +const PREFIX_ROWS: usize = 8192; +const MIN_BATCH_ROWS: usize = 1024; +const MAX_BATCH_ROWS: usize = 64 * 1024; +const MIN_OUTPUT_ROWS: usize = 128; Review Comment: Could you add a line on why 128? For example, the per-batch overhead it is meant to amortize, or the measurement it came from. ########## spark/src/main/scala/org/apache/spark/sql/comet/operators.scala: ########## @@ -3554,6 +3555,26 @@ object CometSortMergeJoinExec extends CometOperatorSerde[SortMergeJoinExec] { override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( CometConf.COMET_EXEC_SORT_MERGE_JOIN_ENABLED) + // The JVM and direct shuffle readers both decode a fresh IPC body for each batch. Review Comment: Makes sense. Keeping the first version scoped to the inputs you validated is reasonable, and the new comment states the reason correctly. Resolving. ########## spark/src/main/scala/org/apache/comet/CometExecIterator.scala: ########## @@ -449,6 +449,17 @@ object CometExecIterator extends Logging { } } + /** Called by the native planner before reducing keys for an eligible semi join. */ + def hasOtherExecutionMemory(context: TaskContext): Boolean = { Review Comment: Could there be a committed test for the `true` branch? For example, allocate a JVM `MemoryConsumer` in the same task before the native plan is created, then assert that no reducer is inserted and the results still match. The 215 paired runs in the description are convincing, but none of them are in the suite. Also, native code calls this method by name, so a rename or signature change would compile cleanly and only fail at runtime. A test would catch that. ########## native/core/src/execution/planner/join.rs: ########## @@ -201,6 +206,64 @@ impl OperatorBuilder for HashJoinBuilder { /// Builder for `SortMergeJoin`. pub struct SortMergeJoinBuilder; +/// Spark identifies the shuffle input. Check the actual sort and its key before replacing its child. +fn reduce_before_sort( + planner: &PhysicalPlanner, + params: &mut JoinParameters, + sort_options: &[SortOptions], +) -> Result<Option<Arc<dyn ExecutionPlan>>, ExecutionError> { + if params.join_type != DFJoinType::LeftSemi || params.join_filter.is_some() { + return Ok(None); + } + let Some(sort) = params.right.native_plan.downcast_ref::<SortExec>() else { + return Ok(None); + }; + let input = sort.input(); + let ([(_, key)], [options]) = (params.join_on.as_slice(), sort_options) else { + return Ok(None); + }; + let schema = input.schema(); + let order = sort.expr().first(); + if !(input.is::<ScanExec>() || input.is::<ShuffleScanExec>()) + || schema.fields().len() != 1 + || !schema.field(0).data_type().is_signed_integer() + || !key.is::<Column>() + || sort.expr().len() != 1 + || order.options != *options + || !order.expr.eq(key) + { + return Ok(None); + } + // Check at first execution, after Spark has constructed the input iterators. + // Reducing the sort's request must not make an already-resident JVM consumer spill. + if let Some(context) = &planner.task_context { + let occupied = JVMClasses::with_env(|env| -> Result<bool, CometError> { + let result = env.call_static_method( + jni::jni_str!("org/apache/comet/CometExecIterator"), + jni::jni_str!("hasOtherExecutionMemory"), + jni::jni_sig!("(Lorg/apache/spark/TaskContext;)Z"), + &[JValue::Object(context.as_obj())], + ); + if let Some(error) = check_exception(env)? { + return Err(error); + } + Ok(result?.z()?) + })?; + if occupied { Review Comment: When this path skips reduction, the user can't see why. The Scala side declares the reduction metrics whenever the config is on, so the join shows zero input and output rows with no explanation. As far as I can tell, this fires whenever a JVM operator above the Comet subtree has already allocated memory by the first pull. Spark's `SortExec` builds its `UnsafeExternalSorter`, which allocates its pointer array, inside `mapPartitionsInternal` before it pulls the child iterator. So a fallback sort above the join would probably disable reduction every time. Could you log this at info or debug level, or add a metric counting skipped admissions, so users can tell "not profitable" from "not admitted"? -- 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]
