sunchao commented on code in PR #7190: URL: https://github.com/apache/datafusion-comet/pull/7190#discussion_r4415880189
########## native/operators/src/membership_key_reduction.rs: ########## @@ -0,0 +1,317 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Keep the first occurrence of each integer key in a duplicate-heavy shuffle batch. +//! A left-semi join only needs to know whether a right-side key exists. + +use std::fmt::Formatter; +use std::hash::Hash; +use std::sync::Arc; + +use arrow::array::{downcast_integer_array, Array, PrimitiveArray, UInt32Array}; +use arrow::compute::{take_record_batch, BatchCoalescer}; +use arrow::datatypes::ArrowPrimitiveType; +use arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, HashSet, Result}; +use datafusion::execution::{async_try_stream, TaskContext}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; +use futures::{StreamExt, TryStreamExt}; + +const PREFIX_ROWS: usize = 8192; +const MIN_BATCH_ROWS: usize = 1024; +const MAX_BATCH_ROWS: usize = 64 * 1024; +// A 128-row target amortizes retained batch metadata while keeping packing buffers small. +const MIN_OUTPUT_ROWS: usize = 128; + +#[derive(Debug)] +pub struct MembershipKeyReductionExec { + input: Arc<dyn ExecutionPlan>, + metrics: ExecutionPlanMetricsSet, +} + +impl MembershipKeyReductionExec { + /// The join planner supplies a single signed integer key column. + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + Self { + input, + metrics: ExecutionPlanMetricsSet::new(), + } + } +} + +impl DisplayAs for MembershipKeyReductionExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CometMembershipKeyReductionExec") + } +} + +impl ExecutionPlan for MembershipKeyReductionExec { + fn name(&self) -> &str { + "CometMembershipKeyReductionExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + self.input.properties() + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::LowerEqual + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + let [input]: [_; 1] = children.try_into().map_err(|_| { + DataFusionError::Internal("Membership key reduction requires one child".into()) + })?; + Ok(Arc::new(Self::new(input))) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let input = self.input.execute(partition, context)?; + let input_rows = + MetricBuilder::new(&self.metrics).counter("right_key_reduction_input_rows", partition); + let output_rows = + MetricBuilder::new(&self.metrics).counter("right_key_reduction_output_rows", partition); + let bypass_rows = + MetricBuilder::new(&self.metrics).counter("right_key_reduction_bypass_rows", partition); + let time = + MetricBuilder::new(&self.metrics).subset_time("right_key_reduction_time", partition); + let mut seen = 0usize; + let mut min_rows = MIN_BATCH_ROWS; + let mut batches = input.map(move |batch| -> Result<_> { + let batch = batch?; + let rows = batch.num_rows(); + input_rows.add(rows); + // Avoid speculative hashing on short partitions, including a single large batch. + let attempt = seen >= PREFIX_ROWS && (min_rows..=MAX_BATCH_ROWS).contains(&rows); Review Comment: Added a many-map control in `4ee20523f5`: identical 4,194,304 right rows / 8,192 repeated keys, 200 reducers, and 4 versus 200 map tasks. It verifies the actual map/reducer counts, the native join flag, the result, and reducer input/output/bypass counters. Run `CometSortMergeJoinBenchmark` with `--many-map-only`. Your core concern is valid: the current reducer bypasses every batch below 1,024 rows, even after the partition has passed its 8,192-row prefix. `N / (M * R)` is a useful balanced-data average, though skew and writer flushes mean 41 million rows is not a hard threshold for every block. The new Scala benchmark compiles with fatal warnings and passes formatting. A diagnostic run with the earlier optimized CI native library from `8650e826` returned 8,192 matches in both shapes. With reduction enabled, the 4-map control sorted 2,153,287 of 4,194,304 input rows; the 200-map case bypassed and sorted all 4,194,304. These are explicitly diagnostic results from that older native build, not new-head performance measurements. The current build is still queued, and the previous current-source native artifact expired; no current-head timing claim is added. I documented the small-frame limitation and tracked cross-batch reduction in #7318. Keeping a key set across input pulls changes the current synchronous per-batch scratch lifetime, so it needs separate memory and performance controls rather than being added to this PR. This addresses the benchmark/scope gap without expanding the reducer's state lifetime. -- 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]
