sunchao commented on code in PR #7190:
URL: https://github.com/apache/datafusion-comet/pull/7190#discussion_r4383441683
##########
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:
Removed the fixed thresholds from the user guide. It now describes
eligibility, bypass behavior, and the meaning of the row and attempt-time
metrics.
##########
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:
Removed both default-value assertions. The suite retains checks for actual
planner eligibility, native execution, matching results, and reduction metrics.
##########
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:
The reservation-refusal case is gone with the reservation. The remaining
behavior is now a single named test,
`short_frames_do_not_disable_larger_batches`, with direct output assertions and
no `(limit, attempted)` loop or conditional expectation.
##########
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:
Removed the benchmark's duplicate result/metric assertions. It now contains
the timed enabled/disabled cases. `CometJoinSuite` retains the result and
metric checks, including the BIGINT and final-batch match cases.
--
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]
