comphead commented on code in PR #7190:
URL: https://github.com/apache/datafusion-comet/pull/7190#discussion_r4414769440
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -3646,6 +3667,9 @@ object CometSortMergeJoinExec extends
CometOperatorSerde[SortMergeJoinExec] {
val joinBuilder = OperatorOuterClass.SortMergeJoin
.newBuilder()
.setJoinType(joinType)
+ .setRightKeyReductionEnabled(
+
CometConf.COMET_EXEC_SORT_MERGE_JOIN_RIGHT_KEY_REDUCTION_ENABLED.get(join.conf)
&&
+ join.joinType == LeftSemi && join.condition.isEmpty &&
canReduceRightKey(join))
Review Comment:
Spark applies the same idea more widely in
`ShuffledHashJoinExec.ignoreDuplicatedKey`, which covers every `LeftExistence`
join type and conditions that reference only streamed attributes and build
keys, from 3.4.3 through 4.1.3. Since this PR already requires the right sort
to output only the key column, rows sharing a key are identical, so a residual
condition gives the same answer for every copy, and an unfiltered `LeftAnti`
join sees the same right key set after reduction. DataFusion 55.2 runs both
through `BitwiseSortMergeJoinStream`. Were anti joins and residual conditions
left out on purpose? If they are a follow-up, a sentence in the description
would help.
##########
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:
Reduction is only tried for batches of 1,024 rows or more, but Comet's
shuffle readers emit one batch per shuffle block and do not regroup them
(`ShuffleScanExec` and `NativeBatchDecoderIterator` both decode block by
block). A block holds one map task's rows for one reducer, so a right side of N
rows from M map tasks into R partitions arrives in blocks of about N / (M * R)
rows. With 200 map tasks and 200 partitions, N has to pass about 41 million
rows before any block reaches 1,024 rows, and below that every batch is
bypassed. The end-to-end run in the description used four map tasks and four
partitions, and `CometSortMergeJoinBenchmark` uses one partition, so neither
covers that shape. Could the benchmark add a many-map case? If the bypass shows
there, a key set kept for the whole partition and reset when full, capped at
the `MAX_BATCH_ROWS.div_ceil(3)` entries the per-batch scratch already allows,
would filter small blocks too. This is arithmetic from the code and the
descripti
on, not a measured run.
--
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]
