viirya commented on code in PR #7190:
URL: https://github.com/apache/datafusion-comet/pull/7190#discussion_r4392062797
##########
native/core/src/execution/planner/join.rs:
##########
@@ -201,6 +206,64 @@ impl OperatorBuilder for HashJoinBuilder {
/// Builder for `SortMergeJoin`.
pub struct SortMergeJoinBuilder;
+/// Spark identifies the shuffle input. Check the actual sort and its key
before replacing its child.
+fn reduce_before_sort(
+ planner: &PhysicalPlanner,
+ params: &mut JoinParameters,
+ sort_options: &[SortOptions],
+) -> Result<Option<Arc<dyn ExecutionPlan>>, ExecutionError> {
+ if params.join_type != DFJoinType::LeftSemi ||
params.join_filter.is_some() {
+ return Ok(None);
+ }
+ let Some(sort) = params.right.native_plan.downcast_ref::<SortExec>() else {
+ return Ok(None);
+ };
+ let input = sort.input();
+ let ([(_, key)], [options]) = (params.join_on.as_slice(), sort_options)
else {
+ return Ok(None);
+ };
+ let schema = input.schema();
+ let order = sort.expr().first();
+ if !(input.is::<ScanExec>() || input.is::<ShuffleScanExec>())
+ || schema.fields().len() != 1
+ || !schema.field(0).data_type().is_signed_integer()
+ || !key.is::<Column>()
+ || sort.expr().len() != 1
+ || order.options != *options
+ || !order.expr.eq(key)
+ {
+ return Ok(None);
+ }
+ // Check at first execution, after Spark has constructed the input
iterators.
Review Comment:
Thanks, the comment now explains why the guard exists and what it doesn't
cover. I left one follow-up on pool scope on `hasOtherExecutionMemory`.
Resolving this one.
##########
spark/src/main/scala/org/apache/comet/CometExecIterator.scala:
##########
@@ -449,6 +449,17 @@ object CometExecIterator extends Logging {
}
}
+ /** Called by the native planner before reducing keys for an eligible semi
join. */
+ def hasOtherExecutionMemory(context: TaskContext): Boolean = {
+ if (!CometSparkSessionExtensions.isOffHeapEnabled(SparkEnv.get.conf)) {
Review Comment:
The mechanism in the planner comment depends on `CometFairMemoryPool`. Its
`try_grow` checks the fair limit locally, and only calls `spark.try_acquire` if
that check passes. The greedy unified pool has no local limit. Its `try_grow`
sends every request straight to `spark.try_acquire`, so the unreduced sort
already reaches Spark on every request. Reduction only makes those requests
smaller.
Does `greedy_unified` need this guard too? As it stands, a fallback JVM sort
or aggregate above the join turns reduction off under the greedy pool as well,
for a reason the comment doesn't give there. If greedy doesn't need it, could
the check apply only to `fair_unified`? If it does, could the comment add the
greedy reason?
Not blocking; a comment-only answer is fine too.
##########
native/core/src/execution/planner/join.rs:
##########
@@ -201,6 +206,64 @@ impl OperatorBuilder for HashJoinBuilder {
/// Builder for `SortMergeJoin`.
pub struct SortMergeJoinBuilder;
+/// Spark identifies the shuffle input. Check the actual sort and its key
before replacing its child.
+fn reduce_before_sort(
+ planner: &PhysicalPlanner,
+ params: &mut JoinParameters,
+ sort_options: &[SortOptions],
+) -> Result<Option<Arc<dyn ExecutionPlan>>, ExecutionError> {
+ if params.join_type != DFJoinType::LeftSemi ||
params.join_filter.is_some() {
+ return Ok(None);
+ }
+ let Some(sort) = params.right.native_plan.downcast_ref::<SortExec>() else {
+ return Ok(None);
+ };
+ let input = sort.input();
+ let ([(_, key)], [options]) = (params.join_on.as_slice(), sort_options)
else {
+ return Ok(None);
+ };
+ let schema = input.schema();
+ let order = sort.expr().first();
+ if !(input.is::<ScanExec>() || input.is::<ShuffleScanExec>())
+ || schema.fields().len() != 1
+ || !schema.field(0).data_type().is_signed_integer()
+ || !key.is::<Column>()
+ || sort.expr().len() != 1
+ || order.options != *options
+ || !order.expr.eq(key)
+ {
+ return Ok(None);
+ }
+ // Check at first execution, after Spark has constructed the input
iterators.
+ // Reducing the sort's request must not make an already-resident JVM
consumer spill.
+ if let Some(context) = &planner.task_context {
+ let occupied = JVMClasses::with_env(|env| -> Result<bool, CometError> {
+ let result = env.call_static_method(
+ jni::jni_str!("org/apache/comet/CometExecIterator"),
+ jni::jni_str!("hasOtherExecutionMemory"),
+ jni::jni_sig!("(Lorg/apache/spark/TaskContext;)Z"),
+ &[JValue::Object(context.as_obj())],
+ );
+ if let Some(error) = check_exception(env)? {
+ return Err(error);
+ }
+ Ok(result?.z()?)
+ })?;
+ if occupied {
Review Comment:
A debug message is enough to tell "not admitted" from "not profitable".
Resolving.
##########
spark/src/main/scala/org/apache/comet/CometExecIterator.scala:
##########
@@ -449,6 +449,17 @@ object CometExecIterator extends Logging {
}
}
+ /** Called by the native planner before reducing keys for an eligible semi
join. */
+ def hasOtherExecutionMemory(context: TaskContext): Boolean = {
Review Comment:
Thanks, the test does what I asked: the no-consumer control proves reduction
runs on the same input, and it goes through the real native-to-Scala callback.
Resolving.
##########
native/operators/src/membership_key_reduction.rs:
##########
@@ -0,0 +1,316 @@
+// 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;
+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);
+ seen = seen.saturating_add(rows);
+ if attempt {
+ let _timer = time.timer();
+ if let Some(reduced) = try_reduce_batch(&batch)? {
+ return Ok((reduced.num_rows() < MIN_OUTPUT_ROWS, reduced));
+ }
+ // A larger shuffle frame may still be worth reducing. With
uniform
+ // frame sizes the rest of the partition goes straight to the
sort.
+ min_rows = (rows + 1).next_power_of_two();
+ }
+ bypass_rows.add(rows);
+ Ok((false, batch))
+ });
+ // The sort accounts for Arrow buffers, but each small RecordBatch
also has
+ // metadata. Pack tiny reductions so this overhead stays small on long
partitions.
+ let mut small = BatchCoalescer::new(self.schema(), MIN_OUTPUT_ROWS);
+ let stream = async_try_stream(|mut out| async move {
+ while let Some((pack, batch)) = batches.try_next().await? {
+ let forward = if pack {
+ small.push_batch(batch)?;
+ None
+ } else {
+ small.finish_buffered_batch()?;
+ Some(batch)
+ };
+ for batch in [small.next_completed_batch(), forward]
Review Comment:
Looks good. Resolving.
##########
native/operators/src/membership_key_reduction.rs:
##########
@@ -0,0 +1,316 @@
+// 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;
+const MIN_OUTPUT_ROWS: usize = 128;
Review Comment:
Thanks. Resolving.
--
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]
