viirya commented on code in PR #7190: URL: https://github.com/apache/datafusion-comet/pull/7190#discussion_r4382079230
########## 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. + let attempt = seen >= PREFIX_ROWS && (min_rows..=MAX_BATCH_ROWS).contains(&rows); + seen = seen.saturating_add(rows); + if attempt { + match try_reduce_batch(&batch, &context, &time) { + Ok(Some(reduced)) => { + output_rows.add(reduced.num_rows()); + return Ok(reduced); + } + // A short shuffle frame may have too few repeats while a later full + // batch is worth reducing. Retry only at the next larger batch size. + Ok(None) => min_rows = (rows + 1).next_power_of_two(), + Err(e) if matches!(e.find_root(), DataFusionError::ResourcesExhausted(_)) => { + min_rows = MAX_BATCH_ROWS + 1; + } + Err(e) => return Err(e), + } + } + bypass_rows.add(rows); + output_rows.add(rows); + Ok(batch) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +fn try_reduce_batch( + batch: &RecordBatch, + context: &TaskContext, + time: &Time, +) -> Result<Option<RecordBatch>> { + let _timer = time.timer(); + let column = batch.column(0).as_ref(); + downcast_integer_array!(column => reduce_integer(column, batch, context), + data_type => datafusion::common::internal_err!("Expected an integer membership key, got {data_type}")) +} + +fn reduce_integer<T: ArrowPrimitiveType>( + column: &PrimitiveArray<T>, + batch: &RecordBatch, + context: &TaskContext, +) -> Result<Option<RecordBatch>> +where + T::Native: Eq + Hash, +{ + let rows = column.len(); + // A low-cardinality shuffle can put just one key in a partition. Its copied + // one-row output needs no full-sized table or index vector. + let constant = column.null_count() == rows + || (column.null_count() == 0 && column.values()[1..] == column.values()[..rows - 1]); + let limit = if constant { 1 } else { rows.div_ceil(3) }; + let reservation = MemoryConsumer::new("MembershipKeyReduction batch") + .with_can_spill(true) + .register(context.memory_pool()); + // The first key past the limit is inserted before it triggers bypass. Budget the + // Arrow value and validity buffers, scratch indices, and array/buffer headers. + let table_bytes = estimate_memory_size::<Option<T::Native>>(limit + 1, 1024)?; + reservation.try_grow( Review Comment: This scratch space is bounded by `MAX_BATCH_ROWS` (roughly 1 MB at worst for i64) and freed before the function returns. That makes it comparable to the temporary allocations in `filter` or `take`, which Comet and DataFusion don't reserve for. Reserving it costs a register, try_grow and release per attempted batch. Under the unified pool, each attempt also goes through JNI to `acquireExecutionMemory`, which can make a JVM consumer in the same task spill. That was the effect behind the first finding in the earlier review, and it still applies, just with a smaller amount. Without the reservation, the `ResourcesExhausted` arm, `min_rows = MAX_BATCH_ROWS + 1`, the constant fast path, the `with_can_spill(true)` question and the pool test could all go. Is there a case where accounting for this memory changes the outcome? ########## 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: What setup cost does this avoid? The only per-attempt cost is the `HashSet` for that batch, and nothing is set up per partition. With the default 8,192-row batch size, this means the first full batch of every partition is never reduced. That covers most of a partition after AQE coalescing leaves it with one or two batches. Do you have numbers showing the gate helps? If not, I'd drop it along with its doc paragraph. ########## spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala: ########## @@ -1680,6 +1683,125 @@ class CometJoinSuite extends CometTestBase { } } + test("right key reduction requires an unfiltered left-semi join sorted over a Comet shuffle") { + val config = CometConf.COMET_EXEC_SORT_MERGE_JOIN_RIGHT_KEY_REDUCTION_ENABLED + assert(!config.get(new SQLConf)) Review Comment: These two asserts check a config default and a proto3 bool default. I don't think they add coverage and they could be removed. ########## native/core/src/execution/planner/join.rs: ########## @@ -201,6 +204,45 @@ impl OperatorBuilder for HashJoinBuilder { /// Builder for `SortMergeJoin`. pub struct SortMergeJoinBuilder; +/// Spark enables reduction only for a single integer key read from a Comet shuffle. Review Comment: The eligibility checks are split between the two layers. - Scala decides: join type, condition, Attribute key, integral type, shuffle source. - Native re-checks: join type, filter, single field, type. It alone checks the sort expression, direction and null ordering. Could you note which side is authoritative and which checks are defensive? That would make it less likely that a future change updates only one side. ########## 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. + let attempt = seen >= PREFIX_ROWS && (min_rows..=MAX_BATCH_ROWS).contains(&rows); + seen = seen.saturating_add(rows); + if attempt { + match try_reduce_batch(&batch, &context, &time) { + Ok(Some(reduced)) => { + output_rows.add(reduced.num_rows()); + return Ok(reduced); + } + // A short shuffle frame may have too few repeats while a later full + // batch is worth reducing. Retry only at the next larger batch size. + Ok(None) => min_rows = (rows + 1).next_power_of_two(), + Err(e) if matches!(e.find_root(), DataFusionError::ResourcesExhausted(_)) => { + min_rows = MAX_BATCH_ROWS + 1; + } + Err(e) => return Err(e), + } + } + bypass_rows.add(rows); + output_rows.add(rows); + Ok(batch) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +fn try_reduce_batch( + batch: &RecordBatch, + context: &TaskContext, + time: &Time, +) -> Result<Option<RecordBatch>> { + let _timer = time.timer(); + let column = batch.column(0).as_ref(); + downcast_integer_array!(column => reduce_integer(column, batch, context), + data_type => datafusion::common::internal_err!("Expected an integer membership key, got {data_type}")) +} + +fn reduce_integer<T: ArrowPrimitiveType>( + column: &PrimitiveArray<T>, + batch: &RecordBatch, + context: &TaskContext, +) -> Result<Option<RecordBatch>> +where + T::Native: Eq + Hash, +{ + let rows = column.len(); + // A low-cardinality shuffle can put just one key in a partition. Its copied + // one-row output needs no full-sized table or index vector. + let constant = column.null_count() == rows + || (column.null_count() == 0 && column.values()[1..] == column.values()[..rows - 1]); + let limit = if constant { 1 } else { rows.div_ceil(3) }; + let reservation = MemoryConsumer::new("MembershipKeyReduction batch") + .with_can_spill(true) + .register(context.memory_pool()); + // The first key past the limit is inserted before it triggers bypass. Budget the + // Arrow value and validity buffers, scratch indices, and array/buffer headers. + let table_bytes = estimate_memory_size::<Option<T::Native>>(limit + 1, 1024)?; + reservation.try_grow( + table_bytes + limit * (size_of::<u32>() + size_of::<T::Native>()) + limit.div_ceil(8), + )?; + + let indices = if constant { + vec![0] + } else { + let mut seen = HashSet::with_capacity(limit + 1); + let mut indices = Vec::with_capacity(limit); + for (row, value) in column.iter().enumerate() { + if seen.insert(value) { + if indices.len() == limit { + return Ok(None); + } + indices.push(row as u32); + } + } + indices + }; + Ok(Some(take_record_batch(batch, &UInt32Array::from(indices))?)) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, Int32Array}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + use datafusion::physical_plan::common::collect; + use datafusion::prelude::{SessionConfig, SessionContext}; + + fn batch(values: impl Iterator<Item = i32>) -> RecordBatch { + RecordBatch::try_from_iter([( + "k", + Arc::new(Int32Array::from_iter_values(values)) as ArrayRef, + )]) + .unwrap() + } + + #[test] + fn keeps_distinct_nullable_keys_in_input_order() { + let keys = [None, Some(i32::MIN), Some(0), Some(-1)]; + let input = RecordBatch::try_from_iter([( + "k", + Arc::new(Int32Array::from_iter( + (0..1025).map(|i| keys[i % keys.len()]), + )) as ArrayRef, + )]) + .unwrap() + .slice(1, 1024); + let reduced = try_reduce_batch(&input, &SessionContext::new().task_ctx(), &Time::new()) + .unwrap() + .unwrap(); + assert_eq!(reduced, input.slice(0, keys.len())); + } + + #[test] + fn keeps_a_different_key_at_the_end_of_an_otherwise_constant_batch() { + for last in [0, 1] { + // A constant batch only needs room for its one-row output. + let limit = if last == 0 { 2048 } else { usize::MAX }; + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(limit))) + .build_arc() + .unwrap(); + let context = + SessionContext::new_with_config_rt(SessionConfig::new(), runtime).task_ctx(); + let input = batch((0..1024).map(|i| if i == 1023 { last } else { 0 })); + let output = try_reduce_batch(&input, &context, &Time::new()) + .unwrap() + .unwrap(); + assert_eq!(output, batch(0..=last)); + } + } + + #[tokio::test] + async fn short_frames_do_not_disable_larger_batches() { + let prefix = batch(0..PREFIX_ROWS as i32); Review Comment: Two scenarios share one loop over `(limit, attempted)`, and the expected output depends on an `if limit == 1` branch. It took me a while to work out what each case asserts. Two named tests would be easier to follow: "a refused reservation bypasses the rest of the partition" and "a short high-cardinality frame does not disable a later full batch". ########## 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. Review Comment: "fresh shuffle batch" no longer describes a requirement of this operator. Any batch whose first column is a signed integer key works. ########## 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. + let attempt = seen >= PREFIX_ROWS && (min_rows..=MAX_BATCH_ROWS).contains(&rows); + seen = seen.saturating_add(rows); + if attempt { + match try_reduce_batch(&batch, &context, &time) { + Ok(Some(reduced)) => { + output_rows.add(reduced.num_rows()); + return Ok(reduced); + } + // A short shuffle frame may have too few repeats while a later full + // batch is worth reducing. Retry only at the next larger batch size. + Ok(None) => min_rows = (rows + 1).next_power_of_two(), Review Comment: With uniform 8,192-row shuffle batches, this sets `min_rows` to 16,384, so the partition never tries again. The behavior is reasonable, but the comment and the user doc describe it as retrying at the next larger tier. In practice it's usually "stop for this partition". Could the wording say that? ########## 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. + let attempt = seen >= PREFIX_ROWS && (min_rows..=MAX_BATCH_ROWS).contains(&rows); + seen = seen.saturating_add(rows); + if attempt { + match try_reduce_batch(&batch, &context, &time) { + Ok(Some(reduced)) => { + output_rows.add(reduced.num_rows()); + return Ok(reduced); + } + // A short shuffle frame may have too few repeats while a later full + // batch is worth reducing. Retry only at the next larger batch size. + Ok(None) => min_rows = (rows + 1).next_power_of_two(), + Err(e) if matches!(e.find_root(), DataFusionError::ResourcesExhausted(_)) => { + min_rows = MAX_BATCH_ROWS + 1; + } + Err(e) => return Err(e), + } + } + bypass_rows.add(rows); + output_rows.add(rows); + Ok(batch) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +fn try_reduce_batch( + batch: &RecordBatch, + context: &TaskContext, + time: &Time, +) -> Result<Option<RecordBatch>> { + let _timer = time.timer(); + let column = batch.column(0).as_ref(); + downcast_integer_array!(column => reduce_integer(column, batch, context), + data_type => datafusion::common::internal_err!("Expected an integer membership key, got {data_type}")) Review Comment: The planner makes this arm unreachable today. Still, since this is an optimization, bypassing (returning `Ok(None)`) seems safer than failing the query if a future input type ever reaches here. ########## 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: This comment gives the reason from the shared-IPC-buffer version. Since the reducer now copies only kept rows and passes bypassed batches through unchanged, is there still a correctness or profitability reason to require a Comet shuffle here? ########## spark/src/test/scala/org/apache/spark/sql/benchmark/CometSortMergeJoinBenchmark.scala: ########## @@ -117,6 +131,66 @@ object CometSortMergeJoinBenchmark extends CometBenchmarkBase { cometConfigs) } + runBenchmark("SortMergeJoin - right key reduction") { + val settings = cometConfigs ++ Map( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") + withSQLConf(settings.toSeq: _*) { + val flag = CometConf.COMET_EXEC_SORT_MERGE_JOIN_RIGHT_KEY_REDUCTION_ENABLED.key + for ((table, distinctKeys, batchSize, shouldReduce) <- Seq( + ("build_repeated", 128, 1024, true), + ("build_balanced", 2048, 8192, true), + ("build_skewed", buildRows / 25 * 12 + 7, 1024, false), + ("build", buildRows, 8192, false))) { + val benchmark = new Benchmark( + s"left semi, $distinctKeys distinct right keys in $buildRows rows", + buildRows.toLong, + output = output) + val query = s"""SELECT /*+ MERGE(b) */ count(*) + |FROM (SELECT k FROM probe WHERE k < $buildRows) p + |LEFT SEMI JOIN $table b ON p.k = b.k""".stripMargin + // Check the result and the actual reducer work outside the timed section. Review Comment: Most of the roughly 60 new lines here assert results and metrics, which `CometJoinSuite` already covers, and the benchmark doesn't run in CI. Could the benchmark keep just the timed cases and leave verification to the suite? It would be shorter and easier to change. ########## docs/source/user-guide/latest/tuning/operators.md: ########## @@ -39,6 +39,20 @@ set, else `spark.sql.shuffle.partitions`). When broadcasts are disabled with a n threshold of 10 MB is used instead. Set `spark.comet.exec.forceShuffledHashJoin.maxBuildSize` to a size in bytes to use a fixed limit, or to a non-positive value to convert every eligible join regardless of size. +### Left-semi sort-merge joins + +A left-semi join returns each left row if any right row matches. For an unfiltered batch join +on a single integer equality key (`=`), set `spark.comet.exec.sortMergeJoin.rightKeyReduction.enabled=true` +to keep one copy of each repeated key per batch before the native sort. The option is off by +default. The right sort must read the key column directly from a Comet shuffle. + +The reducer forwards whole batches until it has seen at least 8,192 rows in the partition. Review Comment: These thresholds (8,192, 1,024 to 65,536, a third) are implementation details that will probably be retuned. I'd keep the user guide to the eligibility conditions, what the metrics mean, and the fact that unprofitable batches are bypassed. -- 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]
