sunchao commented on code in PR #7190: URL: https://github.com/apache/datafusion-comet/pull/7190#discussion_r4366212677
########## native/operators/src/membership_key_reduction.rs: ########## @@ -0,0 +1,565 @@ +// 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. + +//! Opportunistically remove duplicate right-side rows before a left-semi join's sort. +//! Grouping complete integer tuples keeps the set of matching keys, including nulls. +//! The downstream sort restores ordering and the join preserves duplicate left rows. + +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +use arrow::array::{make_array, ArrayRef, MutableArrayData}; +use arrow::compute::concat_batches; +use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::utils::memory::get_record_batch_memory_size; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use datafusion::physical_plan::execution_plan::{CardinalityEffect, EmissionType}; +use datafusion::physical_plan::metrics::{ + ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, +}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use futures::StreamExt; + +// Skip whole batches ending within the initial prefix. After that, small batches pass +// through unchanged, while an oversized or unprofitable attempt disables later attempts. +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>, + properties: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +impl MembershipKeyReductionExec { + /// The join planner must establish that the input consists of signed integers, and + /// that rows equal in every column yield the same membership key without side effects. + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + // The helper can reorder rows in a partition and does not establish global uniqueness. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(input.schema()), + input.output_partitioning().clone(), + EmissionType::Incremental, + input.boundedness(), + )); + Self { + input, + properties, + 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.properties + } + + 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 stream = futures::stream::try_unfold( + (input, 0usize, false), + move |(mut input, seen, mut bypass)| { + let context = Arc::clone(&context); + let (input_rows, output_rows, bypass_rows, time) = ( + input_rows.clone(), + output_rows.clone(), + bypass_rows.clone(), + time.clone(), + ); + async move { + let Some(batch) = input.next().await.transpose()? else { + return Ok(None); + }; + input_rows.add(batch.num_rows()); + let seen = seen.saturating_add(batch.num_rows()); + let reduced = + if bypass || seen <= PREFIX_ROWS || batch.num_rows() < MIN_BATCH_ROWS { + None + } else { + let reduced = match try_reduce_batch(&batch, &context, &time).await { + Ok(reduced) => reduced, + Err(e) + if matches!( + e.find_root(), + DataFusionError::ResourcesExhausted(_) + ) => + { + None + } + Err(e) => return Err(e), + }; + bypass = reduced.is_none(); + reduced + }; + let batch = match reduced { + Some(reduced) => reduced, + None => { + bypass_rows.add(batch.num_rows()); + let _timer = time.timer(); + compact_batch(batch, &context)? + } + }; + output_rows.add(batch.num_rows()); + Ok(Some((batch, (input, seen, bypass)))) + } + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +fn retained_bytes(batch: &RecordBatch) -> usize { + batch.get_array_memory_size() + + std::mem::size_of::<RecordBatch>() + + batch.num_columns() * std::mem::size_of::<ArrayRef>() +} + +fn oversized_backing(array: &ArrayRef) -> bool { + let values = array.len() * array.data_type().primitive_width().unwrap_or(0); + let validity = if array.nulls().is_some() { + array.len().div_ceil(8) + } else { + 0 + }; + array.get_buffer_memory_size() > 2 * (values + validity) + 128 +} + +// A slice can keep a much larger allocation alive in the parent sort. Only copy such +// columns when the pool has room and the resulting batch retains no more unique bytes. +fn compact_batch(batch: RecordBatch, context: &TaskContext) -> Result<RecordBatch> { + if !batch.columns().iter().any(oversized_backing) { + return Ok(batch); + } + let scratch = MemoryConsumer::new("MembershipKeyReduction compact") + .with_can_spill(true) + .register(context.memory_pool()); + if let Err(error) = scratch.try_grow(retained_bytes(&batch)) { + return if matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)) { + Ok(batch) + } else { + Err(error) + }; + } + let columns = batch + .columns() + .iter() + .map(|array| { + if !oversized_backing(array) { + return Ok(Arc::clone(array)); + } + let data = array.to_data(); + let mut compact = MutableArrayData::new(vec![&data], false, array.len()); + compact.try_extend(0, 0, array.len())?; + Ok(make_array(compact.freeze())) + }) + .collect::<Result<Vec<_>>>()?; + let compact = RecordBatch::try_new(batch.schema(), columns)?; + if get_record_batch_memory_size(&compact) <= get_record_batch_memory_size(&batch) { + Ok(compact) + } else { + Ok(batch) + } +} + +/// Collect one bounded partial aggregate privately so any refusal can return the entire +/// original batch. No helper reservation or unpublished aggregate output crosses a yield. +async fn try_reduce_batch( + original: &RecordBatch, + context: &Arc<TaskContext>, + time: &Time, +) -> Result<Option<RecordBatch>> { + let setup_timer = time.timer(); + if original.num_rows() > MAX_BATCH_ROWS { + return Ok(None); + } + let reservation = MemoryConsumer::new("MembershipKeyReduction batch") + .with_can_spill(true) + .register(context.memory_pool()); + reservation.try_grow(retained_bytes(original))?; + let schema = original.schema(); + let row_width: usize = schema + .fields() + .iter() + .map(|field| field.data_type().primitive_width().unwrap_or(1) + 1) + .sum(); + let scratch = + 2 * (original.num_rows() * (std::mem::size_of::<usize>() + row_width) + 64 * 1024); + // Probe for working space before setting up the aggregate. Its real allocations + // use the same pool, so release the estimate while the aggregate runs. + reservation.try_grow(scratch)?; Review Comment: [P2] `try_grow(scratch)` acquires real Spark execution memory and can force a sibling operator to spill. The immediate `shrink` returns the reservation after the disk write has already happened. Remove this speculative acquire/release pair and let the aggregate’s actual reservations drive the existing replay path. Evidence: On Spark 4.1.3 with `fair_unified`, 64 MiB off heap, 10,747,904 bytes of headroom, and a resident `UnsafeExternalSorter`, a semi join over two 8,192-row `Int32` batches with four keys returned four rows in both arms. The sibling wrote zero disk bytes with reduction off and 2,221,918 with it on. A separate reservation-trace replay isolated the 344,064-byte probe. Removing only that acquire/release pair eliminated the spill. ########## native/operators/src/membership_key_reduction.rs: ########## @@ -0,0 +1,565 @@ +// 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. + +//! Opportunistically remove duplicate right-side rows before a left-semi join's sort. +//! Grouping complete integer tuples keeps the set of matching keys, including nulls. +//! The downstream sort restores ordering and the join preserves duplicate left rows. + +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +use arrow::array::{make_array, ArrayRef, MutableArrayData}; +use arrow::compute::concat_batches; +use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::utils::memory::get_record_batch_memory_size; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use datafusion::physical_plan::execution_plan::{CardinalityEffect, EmissionType}; +use datafusion::physical_plan::metrics::{ + ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, +}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use futures::StreamExt; + +// Skip whole batches ending within the initial prefix. After that, small batches pass +// through unchanged, while an oversized or unprofitable attempt disables later attempts. +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>, + properties: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +impl MembershipKeyReductionExec { + /// The join planner must establish that the input consists of signed integers, and + /// that rows equal in every column yield the same membership key without side effects. + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + // The helper can reorder rows in a partition and does not establish global uniqueness. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(input.schema()), + input.output_partitioning().clone(), + EmissionType::Incremental, + input.boundedness(), + )); + Self { + input, + properties, + 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.properties + } + + 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 stream = futures::stream::try_unfold( + (input, 0usize, false), + move |(mut input, seen, mut bypass)| { + let context = Arc::clone(&context); + let (input_rows, output_rows, bypass_rows, time) = ( + input_rows.clone(), + output_rows.clone(), + bypass_rows.clone(), + time.clone(), + ); + async move { + let Some(batch) = input.next().await.transpose()? else { + return Ok(None); + }; + input_rows.add(batch.num_rows()); + let seen = seen.saturating_add(batch.num_rows()); + let reduced = + if bypass || seen <= PREFIX_ROWS || batch.num_rows() < MIN_BATCH_ROWS { + None + } else { + let reduced = match try_reduce_batch(&batch, &context, &time).await { + Ok(reduced) => reduced, + Err(e) + if matches!( + e.find_root(), + DataFusionError::ResourcesExhausted(_) + ) => + { + None + } + Err(e) => return Err(e), + }; + bypass = reduced.is_none(); + reduced + }; + let batch = match reduced { + Some(reduced) => reduced, + None => { + bypass_rows.add(batch.num_rows()); + let _timer = time.timer(); + compact_batch(batch, &context)? + } + }; + output_rows.add(batch.num_rows()); + Ok(Some((batch, (input, seen, bypass)))) + } + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +fn retained_bytes(batch: &RecordBatch) -> usize { + batch.get_array_memory_size() + + std::mem::size_of::<RecordBatch>() + + batch.num_columns() * std::mem::size_of::<ArrayRef>() +} + +fn oversized_backing(array: &ArrayRef) -> bool { + let values = array.len() * array.data_type().primitive_width().unwrap_or(0); + let validity = if array.nulls().is_some() { + array.len().div_ceil(8) + } else { + 0 + }; + array.get_buffer_memory_size() > 2 * (values + validity) + 128 +} + +// A slice can keep a much larger allocation alive in the parent sort. Only copy such +// columns when the pool has room and the resulting batch retains no more unique bytes. +fn compact_batch(batch: RecordBatch, context: &TaskContext) -> Result<RecordBatch> { + if !batch.columns().iter().any(oversized_backing) { + return Ok(batch); + } + let scratch = MemoryConsumer::new("MembershipKeyReduction compact") + .with_can_spill(true) + .register(context.memory_pool()); + if let Err(error) = scratch.try_grow(retained_bytes(&batch)) { + return if matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)) { + Ok(batch) + } else { + Err(error) + }; + } + let columns = batch + .columns() + .iter() + .map(|array| { + if !oversized_backing(array) { + return Ok(Arc::clone(array)); + } + let data = array.to_data(); + let mut compact = MutableArrayData::new(vec![&data], false, array.len()); + compact.try_extend(0, 0, array.len())?; + Ok(make_array(compact.freeze())) + }) + .collect::<Result<Vec<_>>>()?; + let compact = RecordBatch::try_new(batch.schema(), columns)?; + if get_record_batch_memory_size(&compact) <= get_record_batch_memory_size(&batch) { + Ok(compact) + } else { + Ok(batch) + } +} + +/// Collect one bounded partial aggregate privately so any refusal can return the entire +/// original batch. No helper reservation or unpublished aggregate output crosses a yield. +async fn try_reduce_batch( + original: &RecordBatch, + context: &Arc<TaskContext>, + time: &Time, +) -> Result<Option<RecordBatch>> { + let setup_timer = time.timer(); + if original.num_rows() > MAX_BATCH_ROWS { + return Ok(None); + } + let reservation = MemoryConsumer::new("MembershipKeyReduction batch") + .with_can_spill(true) + .register(context.memory_pool()); + reservation.try_grow(retained_bytes(original))?; + let schema = original.schema(); + let row_width: usize = schema + .fields() + .iter() + .map(|field| field.data_type().primitive_width().unwrap_or(1) + 1) + .sum(); + let scratch = + 2 * (original.num_rows() * (std::mem::size_of::<usize>() + row_width) + 64 * 1024); + // Probe for working space before setting up the aggregate. Its real allocations + // use the same pool, so release the estimate while the aggregate runs. + reservation.try_grow(scratch)?; + reservation.shrink(scratch); + let source = + MemorySourceConfig::try_new_exec(&[vec![original.clone()]], Arc::clone(&schema), None)?; + let keys = schema + .fields() + .iter() + .enumerate() + .map(|(i, field)| { + ( + Arc::new(Column::new(field.name(), i)) as Arc<dyn PhysicalExpr>, + field.name().clone(), + ) + }) + .collect(); + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(keys), + vec![], + vec![], + source, + Arc::clone(&schema), + )?; + let mut stream = aggregate.execute(0, Arc::clone(context))?; + drop(setup_timer); + let mut output = Vec::new(); + while let Some(batch) = stream.next().await.transpose()? { + let _timer = time.timer(); + let batch = compact_batch(batch, context)?; + reservation.try_grow(retained_bytes(&batch))?; + output.push(batch); + } + let _timer = time.timer(); + if let Some(elapsed) = aggregate.metrics().and_then(|m| m.elapsed_compute()) { Review Comment: [P2] `right_key_reduction_time` receives the aggregate’s `elapsed_compute` only after all output batches have been retained. If the output reservation fails, the join successfully replays the original batch while its SQL metric omits the hash work already performed. Transfer the inner compute time on recoverable exits as well. Evidence: Using the same third-reservation refusal as the PR’s operator test, I grouped 65,536 rows with 12 integer columns and 16,384 distinct tuples. All rows were replayed. Across ten refusals, the inner aggregate recorded a median 5.63 ms of completed compute, while the published reduction metric reported 0.33 ms. ########## native/operators/src/membership_key_reduction.rs: ########## @@ -0,0 +1,565 @@ +// 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. + +//! Opportunistically remove duplicate right-side rows before a left-semi join's sort. +//! Grouping complete integer tuples keeps the set of matching keys, including nulls. +//! The downstream sort restores ordering and the join preserves duplicate left rows. + +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +use arrow::array::{make_array, ArrayRef, MutableArrayData}; +use arrow::compute::concat_batches; +use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::utils::memory::get_record_batch_memory_size; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use datafusion::physical_plan::execution_plan::{CardinalityEffect, EmissionType}; +use datafusion::physical_plan::metrics::{ + ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, +}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use futures::StreamExt; + +// Skip whole batches ending within the initial prefix. After that, small batches pass +// through unchanged, while an oversized or unprofitable attempt disables later attempts. +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>, + properties: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +impl MembershipKeyReductionExec { + /// The join planner must establish that the input consists of signed integers, and + /// that rows equal in every column yield the same membership key without side effects. + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + // The helper can reorder rows in a partition and does not establish global uniqueness. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(input.schema()), + input.output_partitioning().clone(), + EmissionType::Incremental, + input.boundedness(), + )); + Self { + input, + properties, + 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.properties + } + + 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 stream = futures::stream::try_unfold( + (input, 0usize, false), + move |(mut input, seen, mut bypass)| { + let context = Arc::clone(&context); + let (input_rows, output_rows, bypass_rows, time) = ( + input_rows.clone(), + output_rows.clone(), + bypass_rows.clone(), + time.clone(), + ); + async move { + let Some(batch) = input.next().await.transpose()? else { + return Ok(None); + }; + input_rows.add(batch.num_rows()); + let seen = seen.saturating_add(batch.num_rows()); + let reduced = + if bypass || seen <= PREFIX_ROWS || batch.num_rows() < MIN_BATCH_ROWS { + None + } else { + let reduced = match try_reduce_batch(&batch, &context, &time).await { + Ok(reduced) => reduced, + Err(e) + if matches!( + e.find_root(), + DataFusionError::ResourcesExhausted(_) + ) => + { + None + } + Err(e) => return Err(e), + }; + bypass = reduced.is_none(); + reduced + }; + let batch = match reduced { + Some(reduced) => reduced, + None => { + bypass_rows.add(batch.num_rows()); + let _timer = time.timer(); + compact_batch(batch, &context)? + } + }; + output_rows.add(batch.num_rows()); + Ok(Some((batch, (input, seen, bypass)))) + } + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +fn retained_bytes(batch: &RecordBatch) -> usize { + batch.get_array_memory_size() + + std::mem::size_of::<RecordBatch>() + + batch.num_columns() * std::mem::size_of::<ArrayRef>() +} + +fn oversized_backing(array: &ArrayRef) -> bool { + let values = array.len() * array.data_type().primitive_width().unwrap_or(0); + let validity = if array.nulls().is_some() { + array.len().div_ceil(8) + } else { + 0 + }; + array.get_buffer_memory_size() > 2 * (values + validity) + 128 +} + +// A slice can keep a much larger allocation alive in the parent sort. Only copy such +// columns when the pool has room and the resulting batch retains no more unique bytes. +fn compact_batch(batch: RecordBatch, context: &TaskContext) -> Result<RecordBatch> { + if !batch.columns().iter().any(oversized_backing) { + return Ok(batch); + } + let scratch = MemoryConsumer::new("MembershipKeyReduction compact") + .with_can_spill(true) + .register(context.memory_pool()); + if let Err(error) = scratch.try_grow(retained_bytes(&batch)) { + return if matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)) { + Ok(batch) + } else { + Err(error) + }; + } + let columns = batch + .columns() + .iter() + .map(|array| { + if !oversized_backing(array) { + return Ok(Arc::clone(array)); + } + let data = array.to_data(); + let mut compact = MutableArrayData::new(vec![&data], false, array.len()); + compact.try_extend(0, 0, array.len())?; + Ok(make_array(compact.freeze())) + }) + .collect::<Result<Vec<_>>>()?; + let compact = RecordBatch::try_new(batch.schema(), columns)?; + if get_record_batch_memory_size(&compact) <= get_record_batch_memory_size(&batch) { + Ok(compact) + } else { + Ok(batch) + } +} + +/// Collect one bounded partial aggregate privately so any refusal can return the entire +/// original batch. No helper reservation or unpublished aggregate output crosses a yield. +async fn try_reduce_batch( + original: &RecordBatch, + context: &Arc<TaskContext>, + time: &Time, +) -> Result<Option<RecordBatch>> { + let setup_timer = time.timer(); + if original.num_rows() > MAX_BATCH_ROWS { + return Ok(None); + } + let reservation = MemoryConsumer::new("MembershipKeyReduction batch") + .with_can_spill(true) + .register(context.memory_pool()); + reservation.try_grow(retained_bytes(original))?; + let schema = original.schema(); + let row_width: usize = schema + .fields() + .iter() + .map(|field| field.data_type().primitive_width().unwrap_or(1) + 1) + .sum(); + let scratch = + 2 * (original.num_rows() * (std::mem::size_of::<usize>() + row_width) + 64 * 1024); + // Probe for working space before setting up the aggregate. Its real allocations + // use the same pool, so release the estimate while the aggregate runs. + reservation.try_grow(scratch)?; + reservation.shrink(scratch); + let source = + MemorySourceConfig::try_new_exec(&[vec![original.clone()]], Arc::clone(&schema), None)?; + let keys = schema + .fields() + .iter() + .enumerate() + .map(|(i, field)| { + ( + Arc::new(Column::new(field.name(), i)) as Arc<dyn PhysicalExpr>, + field.name().clone(), + ) + }) + .collect(); + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(keys), + vec![], + vec![], + source, + Arc::clone(&schema), + )?; + let mut stream = aggregate.execute(0, Arc::clone(context))?; + drop(setup_timer); + let mut output = Vec::new(); + while let Some(batch) = stream.next().await.transpose()? { + let _timer = time.timer(); + let batch = compact_batch(batch, context)?; + reservation.try_grow(retained_bytes(&batch))?; + output.push(batch); + } + let _timer = time.timer(); + if let Some(elapsed) = aggregate.metrics().and_then(|m| m.elapsed_compute()) { + time.add_duration(std::time::Duration::from_nanos(elapsed as u64)); + } + drop(stream); + drop(aggregate); + let rows: usize = output.iter().map(RecordBatch::num_rows).sum(); + if 5 * rows >= 4 * original.num_rows() { Review Comment: [P2] The 80% retention cutoff keeps rebuilding a partial aggregate for every 78%-distinct batch, even when grouping costs more than the sort work it saves. Use a more conservative cutoff or a measured cost signal so the partition can switch to bypass on this workload. Evidence: One partition with 64 batches of 8,192 `Int32` rows and key `(48271 * i) % 6400` returned eight left rows in both arms. The right sort received 411,392 rows with reduction and 524,288 without it. Across 17 alternating optimized native trials, median sort-and-join time rose from 23.31 to 30.38 ms. A control that bypassed after the first 78%-distinct batch measured 23.56 ms enabled versus 23.37 ms disabled. Spark 4.1.3 independently showed median combined right-sort and reduction time rising from 13.1 to 28.6 ms. ########## native/operators/src/membership_key_reduction.rs: ########## @@ -0,0 +1,565 @@ +// 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. + +//! Opportunistically remove duplicate right-side rows before a left-semi join's sort. +//! Grouping complete integer tuples keeps the set of matching keys, including nulls. +//! The downstream sort restores ordering and the join preserves duplicate left rows. + +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +use arrow::array::{make_array, ArrayRef, MutableArrayData}; +use arrow::compute::concat_batches; +use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::utils::memory::get_record_batch_memory_size; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use datafusion::physical_plan::execution_plan::{CardinalityEffect, EmissionType}; +use datafusion::physical_plan::metrics::{ + ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, +}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use futures::StreamExt; + +// Skip whole batches ending within the initial prefix. After that, small batches pass +// through unchanged, while an oversized or unprofitable attempt disables later attempts. +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>, + properties: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +impl MembershipKeyReductionExec { + /// The join planner must establish that the input consists of signed integers, and + /// that rows equal in every column yield the same membership key without side effects. + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + // The helper can reorder rows in a partition and does not establish global uniqueness. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(input.schema()), + input.output_partitioning().clone(), + EmissionType::Incremental, + input.boundedness(), + )); + Self { + input, + properties, + 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.properties + } + + 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 stream = futures::stream::try_unfold( + (input, 0usize, false), + move |(mut input, seen, mut bypass)| { + let context = Arc::clone(&context); + let (input_rows, output_rows, bypass_rows, time) = ( + input_rows.clone(), + output_rows.clone(), + bypass_rows.clone(), + time.clone(), + ); + async move { + let Some(batch) = input.next().await.transpose()? else { + return Ok(None); + }; + input_rows.add(batch.num_rows()); + let seen = seen.saturating_add(batch.num_rows()); + let reduced = + if bypass || seen <= PREFIX_ROWS || batch.num_rows() < MIN_BATCH_ROWS { + None + } else { + let reduced = match try_reduce_batch(&batch, &context, &time).await { + Ok(reduced) => reduced, + Err(e) + if matches!( + e.find_root(), + DataFusionError::ResourcesExhausted(_) + ) => + { + None + } + Err(e) => return Err(e), + }; + bypass = reduced.is_none(); + reduced + }; + let batch = match reduced { + Some(reduced) => reduced, + None => { + bypass_rows.add(batch.num_rows()); + let _timer = time.timer(); + compact_batch(batch, &context)? + } + }; + output_rows.add(batch.num_rows()); + Ok(Some((batch, (input, seen, bypass)))) + } + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +fn retained_bytes(batch: &RecordBatch) -> usize { + batch.get_array_memory_size() + + std::mem::size_of::<RecordBatch>() + + batch.num_columns() * std::mem::size_of::<ArrayRef>() +} + +fn oversized_backing(array: &ArrayRef) -> bool { + let values = array.len() * array.data_type().primitive_width().unwrap_or(0); + let validity = if array.nulls().is_some() { + array.len().div_ceil(8) + } else { + 0 + }; + array.get_buffer_memory_size() > 2 * (values + validity) + 128 Review Comment: [P2] `oversized_backing` compares one column’s logical size with the capacity of the shared IPC body. A fully packed composite-key batch can therefore qualify every column for copying. This continues on every batch after high-cardinality bypass, and the final size comparison happens after allocation. Check the potential saving across the whole batch before copying. Evidence: I passed 128 Comet-decoded batches of 4,096 unique rows and four nullable `Int32` columns through the reducer. All 524,288 rows bypassed reduction, but every output batch had new value buffers. The copies allocated 8,650,752 bytes and retained exactly as many bytes as the inputs. A whole-batch precheck returned identical rows with zero copies and cut median reducer time from 2.59 to 0.50 ms. ########## native/operators/src/membership_key_reduction.rs: ########## @@ -0,0 +1,565 @@ +// 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. + +//! Opportunistically remove duplicate right-side rows before a left-semi join's sort. +//! Grouping complete integer tuples keeps the set of matching keys, including nulls. +//! The downstream sort restores ordering and the join preserves duplicate left rows. + +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +use arrow::array::{make_array, ArrayRef, MutableArrayData}; +use arrow::compute::concat_batches; +use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::utils::memory::get_record_batch_memory_size; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use datafusion::physical_plan::execution_plan::{CardinalityEffect, EmissionType}; +use datafusion::physical_plan::metrics::{ + ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, +}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use futures::StreamExt; + +// Skip whole batches ending within the initial prefix. After that, small batches pass +// through unchanged, while an oversized or unprofitable attempt disables later attempts. +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>, + properties: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +impl MembershipKeyReductionExec { + /// The join planner must establish that the input consists of signed integers, and + /// that rows equal in every column yield the same membership key without side effects. + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + // The helper can reorder rows in a partition and does not establish global uniqueness. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(input.schema()), + input.output_partitioning().clone(), + EmissionType::Incremental, + input.boundedness(), + )); + Self { + input, + properties, + 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.properties + } + + 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 stream = futures::stream::try_unfold( + (input, 0usize, false), + move |(mut input, seen, mut bypass)| { + let context = Arc::clone(&context); + let (input_rows, output_rows, bypass_rows, time) = ( + input_rows.clone(), + output_rows.clone(), + bypass_rows.clone(), + time.clone(), + ); + async move { + let Some(batch) = input.next().await.transpose()? else { + return Ok(None); + }; + input_rows.add(batch.num_rows()); + let seen = seen.saturating_add(batch.num_rows()); + let reduced = + if bypass || seen <= PREFIX_ROWS || batch.num_rows() < MIN_BATCH_ROWS { + None + } else { + let reduced = match try_reduce_batch(&batch, &context, &time).await { + Ok(reduced) => reduced, + Err(e) + if matches!( + e.find_root(), + DataFusionError::ResourcesExhausted(_) + ) => + { + None + } + Err(e) => return Err(e), + }; + bypass = reduced.is_none(); + reduced + }; + let batch = match reduced { + Some(reduced) => reduced, + None => { + bypass_rows.add(batch.num_rows()); + let _timer = time.timer(); + compact_batch(batch, &context)? + } + }; + output_rows.add(batch.num_rows()); + Ok(Some((batch, (input, seen, bypass)))) + } + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +fn retained_bytes(batch: &RecordBatch) -> usize { + batch.get_array_memory_size() Review Comment: [P2] `get_array_memory_size()` charges a shared IPC backing allocation once per column that references it. On a duplicate-heavy composite key, the inflated request can disable reduction for the rest of a partition even though the aggregate fits. Count unique buffers with `get_record_batch_memory_size()` and add array and batch metadata separately. Evidence: Sixteen Comet-decoded batches of 4,096 rows, eight `Int32` columns, and four distinct tuples retained 135,168 unique buffer bytes per batch. The reducer requested 1,082,280 bytes. With a 1 MiB pool, all 65,536 rows passed through. Changing only the reservation calculation reduced the output to 8,248 rows. In a Spark AQE join with 16 keys and an 8 MiB pool, direct shuffle sent all 32,768 rows to the sort and spilled twice. The Arrow-stream control sent 8,204 rows and had zero sort spills. -- 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]
