sunchao commented on code in PR #7190:
URL: https://github.com/apache/datafusion-comet/pull/7190#discussion_r4391164853
##########
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:
Updated in 5ba986ac5c. Added the mechanism to the native planner comment:
Comet's fair pool checks its own limits before calling Spark. A smaller sort
reservation after reduction can pass that check and reach
`acquireExecutionMemory`, which can spill a resident JVM consumer. The
unreduced request could instead have been refused locally, causing the native
sort to spill.
The comment and user guide now describe admission as a best-effort check at
first execution. It does not cover consumers that allocate later, and different
tasks can make different admission decisions. The new committed Spark test
exercises both outcomes through the native planner.
##########
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:
Updated in 5ba986ac5c. Added a debug message on rejected admission,
including the right sort's plan ID and the reason: another Spark consumer holds
execution memory in the task. I verified the message during the new
resident-memory test. This distinguishes skipped admission from an admitted
reducer deciding that individual batches are unprofitable without adding
another metric.
##########
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:
Updated in 5ba986ac5c. Added `right key reduction skips a task with resident
JVM execution memory` to `CometJoinSuite`. It holds 4 KiB through a same-task
off-heap `MemoryConsumer` before the first native pull, then verifies correct
duplicate-preserving results, zero reducer rows, and all 32,768 right rows
reaching the native sort.
A no-consumer control proves that reduction actually runs on the same input.
The test executes the native-to-Scala callback rather than calling the helper
directly, so it also covers the JNI method lookup. The reservation is released
in `finally`. It passed with Spark 3.5.9 and the rebuilt native library.
##########
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:
Updated in 5ba986ac5c. Changed both the per-input drain and the EOF drain to
`while let Some(batch) = small.next_completed_batch()`. Packed output still
precedes each forwarded batch. The existing packing test covers full packed
batches, bypass ordering, original-buffer identity for bypassed input, and the
EOF tail.
##########
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:
Updated in 5ba986ac5c. Added a comment explaining the 128-row target:
amortize retained batch metadata while keeping the packing buffer small. It is
a conservative packing target, not a correctness boundary or a claim of an
optimal value. Flushes before bypassed batches and at EOF can still produce
shorter batches.
--
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]
