sunchao commented on code in PR #7190: URL: https://github.com/apache/datafusion-comet/pull/7190#discussion_r4383440399
########## 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: Removed the scratch reservation, `with_can_spill`, and reservation-refusal path. The bounded scratch now lives only for the synchronous batch, as you suggested. The native planner separately skips reduction if another Spark consumer already holds execution memory, since changing the sort's requests can also affect that consumer. The paired unified-memory controls are in the PR description. I retained the constant fast path because it avoids constructing a hash table and hashing every row in a single-key batch. All-null batches are detected from the null count. It still copies the selected row, so it does not retain the full input buffer. ########## 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: Agreed, there is no per-partition setup to amortize. This gate avoids speculative hashing on short partitions. I clarified the source comment in d92d50e606 and ran an isolated control changing only `PREFIX_ROWS` to zero. Optimized native sort plus left-semi join, Int32 keys, one input batch, medians of 20 rotated-order trials after three warm-up trials: | Input | With prefix bypass | Without prefix bypass | | --- | ---: | ---: | | 65,536 unique rows | 3.111 ms | 3.462 ms | | 8,192 rows, 2,048 distinct keys | 0.258 ms | 0.148 ms | These controls use DataFusion 55.2.0 and Arrow 59.3.1 and exclude Spark and shuffle. They show the trade-off: avoiding a failed hash attempt helps the short unique input, while the gate misses a useful reduction on the short duplicate-heavy input. I retained it as a conservative policy for this version, not as a correctness requirement or an optimal threshold for every workload. Leaving this thread open for your assessment. ########## 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: Updated both the source comment and user guide to say that uniform-size batches go straight to the sort after an unprofitable attempt. A later larger frame can still be tried. The named `short_frames_do_not_disable_larger_batches` test covers that distinction. ########## 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: Changed the unsupported-type arm to `Ok(None)`, so it bypasses this optional reduction instead of failing the query. ########## 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: Removed the stale freshness/ownership wording. Successful reductions copy the selected rows, and bypassed batches are forwarded untouched. ########## 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: The reducer no longer requires shuffle-owned buffers for correctness. I kept this first opt-in version scoped to Comet shuffle inputs because that is where the performance and memory controls were run, including direct and Arrow-stream readers. Bucketed scans could also benefit, but their profitability and memory behavior have not been validated here. I replaced the old ownership explanation with that scope rationale in d92d50e606. Scala establishes Comet-shuffle provenance; the native check verifies the expected direct scan boundary. Leaving this thread open for your assessment. ########## 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: Clarified the responsibilities in the native helper comment and PR description: Spark establishes Comet-shuffle provenance, including AQE/reused exchanges. The native helper is authoritative for inserting the reducer into the executable plan and validates the actual sort, key, ordering, and supported join shape. A serialized request alone does not force insertion. -- 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]
