[AURON #2603] Fix SMJ fallback timing metrics under output backpressure

native-engine/datafusion-ext-plans/src/broadcast_join_build_hash_map_exec.rs

@@ -273,6 +273,164 @@ mod tests {
         )?))
     }
 
+    #[tokio::test]
+    async fn sorted_fallback_metrics_exclude_output_wait() -> Result<()> {
+        MemManager::init(1000000);
+        let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)]));
+        let num_rows = 3 * datafusion_ext_commons::batch_size();
+        let probe_values = (0..num_rows as i32).collect::<Vec<_>>();
+        let probe = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int32Array::from(probe_values.clone()))],
+        )?;
+        // A null table column marks the single build row as sorted for SMJ fallback.
+        let built = RecordBatch::try_new(
+            join_hash_map_schema(&schema),
+            vec![
+                Arc::new(Int32Array::from(vec![-1])),
+                new_null_array(&DataType::Binary, 1),
+            ],
+        )?;
+        for build_left in [false, true] {
+            let (left, right) = if build_left {
+                (memory(built.clone())?, memory(probe.clone())?)
+            } else {
+                (memory(probe.clone())?, memory(built.clone())?)
+            };
+            let output_schema =
+                Arc::new(build_join_schema(&schema, &schema, &JoinType::Inner.try_into()?).0);
+            let join = BroadcastJoinExec::try_new(
+                output_schema,
+                left,
+                right,
+                vec![],
+                JoinType::Inner,
+                if build_left {
+                    JoinSide::Left
+                } else {
+                    JoinSide::Right
+                },
+                true,
+                None,
+                false,
+                None,
+            )?;
+            let mut stream = join.execute(0, SessionContext::new().task_ctx())?;
+            let first = stream
+                .next()
+                .await
+                .transpose()?
+                .expect("first output batch");
+            // Let the bounded output channel fill before draining the remaining batches.
+            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+            let mut output = vec![first];
+            output.extend(collect(stream).await?);
+            let mut actual = vec![];
+            for batch in output {
+                let values = batch
+                    .column(usize::from(build_left))
+                    .as_any()
+                    .downcast_ref::<Int32Array>()
+                    .expect("probe values");
+                let build_values = batch
+                    .column(usize::from(!build_left))
+                    .as_any()
+                    .downcast_ref::<Int32Array>()
+                    .expect("build values");
+                actual.extend(values.values().iter().copied());
+                assert!(build_values.values().iter().all(|&v| v == -1));
+            }
+            actual.sort_unstable();
+            assert_eq!(actual, probe_values);
+            let metrics = join.metrics().expect("join metrics");
+            let fallback_time = metrics
+                .sum_by_name("fallback_sort_merge_join_time")
+                .expect("fallback time")
+                .as_usize() as i64;
+            assert!(
+                fallback_time > 0,
+                "build_left={build_left}: {fallback_time}"
+            );
+            let elapsed_compute = metrics.elapsed_compute().expect("elapsed compute") as i64;
+            assert!(
+                elapsed_compute > 0,
+                "build_left={build_left}: {elapsed_compute}"
+            );
+        }
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn keyed_sorted_fallback_propagates_sort_and_spill_metrics() -> Result<()> {
+        MemManager::init(1_000_000);
+        let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)]));
+        // Exceed the 16 MiB automatic spill check threshold while sorting.
+        let probe = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int32Array::from_iter_values((0..1 << 21).rev()))],
+        )?;
+        let build_values = Int32Array::from(vec![0, 2, 4]);
+        let built = RecordBatch::try_new(
+            join_hash_map_schema(&schema),
+            vec![
+                Arc::new(build_values.clone()),
+                new_null_array(&DataType::Binary, 3), // sorted fallback marker
+            ],
+        )?;
+        for build_left in [false, true] {
+            let (left, right, side) = if build_left {
+                (
+                    memory(built.clone())?,
+                    memory(probe.clone())?,
+                    JoinSide::Left,
+                )
+            } else {
+                (
+                    memory(probe.clone())?,
+                    memory(built.clone())?,
+                    JoinSide::Right,
+                )
+            };
+            let join = BroadcastJoinExec::try_new(
+                Arc::new(build_join_schema(&schema, &schema, &JoinType::Inner.try_into()?).0),
+                left,
+                right,
+                vec![(Arc::new(Column::new("v", 0)), Arc::new(Column::new("v", 0)))],
+                JoinType::Inner,
+                side,
+                true,
+                None,
+                false,
+                None,
+            )?;
+            let output = collect(join.execute(0, SessionContext::new().task_ctx())?).await?;
+            let output = arrow::compute::concat_batches(&join.schema(), &output)?;
+            for column in output.columns() {
+                assert_eq!(
+                    column.as_any().downcast_ref::<Int32Array>(),
+                    Some(&build_values)
+                );
+            }
+            let metrics = join.metrics().expect("join metrics");
+            let metric_value = |name: &str| metrics.sum_by_name(name).expect(name).as_usize();
+            // The probe sort's compute metric is appended after the join baseline.
+            let compute = metrics
+                .iter()
+                .filter(|m| m.value().name() == "elapsed_compute")
+                .map(|m| m.value().as_usize())
+                .collect::<Vec<_>>();
+            assert_eq!(compute.len(), 2, "{metrics}");
+            assert!(
+                compute[1] > 0 && metric_value("fallback_sort_merge_join_time") > compute[1],
+                "{metrics}"
+            );
+            for name in ["disk_spill_size", "disk_spill_iotime"] {
+                assert!(metric_value(name) > 0, "{metrics}");
+            }
+        }
+        Ok(())
+    }
+
     #[tokio::test]
     async fn nested_loop_sorted_fallback_preserves_condition_and_projection() -> Result<()> {
         MemManager::init(1000000);

native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs

@@ -437,6 +437,9 @@ async fn execute_join_with_smj_fallback(
     exec_ctx: Arc<ExecutionContext>,
     sender: Arc<WrappedRecordBatchSender>,
 ) -> Result<()> {
+    let elapsed_compute = exec_ctx.baseline_metrics().elapsed_compute().clone();
+    let _timer = elapsed_compute.timer();
+
     // remove the table data column from the built stream
     let built_schema = built.schema();
     let built_sorted: Arc<dyn ExecutionPlan> = {
@@ -460,6 +463,8 @@ async fn execute_join_with_smj_fallback(
 
     let left_keys = smj_fallback_keys(&join_params.left_keys);
     let right_keys = smj_fallback_keys(&join_params.right_keys);
+    // Keep sort timing separate from the join's output-wait deductions.
+    let sort_metrics = ExecutionPlanMetricsSet::new();
 
     // create sorted streams, build side is already sorted
     let (left_exec, right_exec) = match broadcast_side {
@@ -468,15 +473,15 @@ async fn execute_join_with_smj_fallback(
             create_default_ascending_sort_exec(
                 probed_plan,
                 &right_keys,
-                Some(exec_ctx.execution_plan_metrics().clone()),
+                Some(sort_metrics.clone()),
                 false, // do not record output metric
             ),
         ),
         JoinSide::Right => (
             create_default_ascending_sort_exec(
                 probed_plan,
                 &left_keys,
-                Some(exec_ctx.execution_plan_metrics().clone()),
+                Some(sort_metrics.clone()),
                 false, // do not record output metric
             ),
             built_sorted,
@@ -518,34 +523,33 @@ async fn execute_join_with_smj_fallback(
         smj_exec.execute_projected(exec_ctx.partition_id(), exec_ctx.task_ctx(), &projection)?;
 
     // send all outputs
-    while let Some(batch) = join_output.next().await.transpose()? {
+    while let Some(batch) = elapsed_compute
+        .exclude_timer_async(join_output.next())
+        .await
+        .transpose()?
+    {
         sender.send(batch).await;
     }
 
     // elapsed_compute = sort time + merge time
     let smj_time = exec_ctx.register_timer_metric("fallback_sort_merge_join_time");
+    let sort_metrics = sort_metrics.clone_inner();
     smj_time.add_duration(Duration::from_nanos(
-        left_exec
-            .metrics()
-            .and_then(|m| m.elapsed_compute())
-            .unwrap_or(0) as u64,
+        sort_metrics.elapsed_compute().unwrap_or(0) as u64,
     ));
-    smj_time.add_duration(Duration::from_nanos(
-        right_exec
-            .metrics()
-            .and_then(|m| m.elapsed_compute())
-            .unwrap_or(0) as u64,
-    ));
-    smj_time.add_duration(Duration::from_nanos(
+    let merge_time = Duration::from_nanos(
         smj_exec
             .metrics()
             .and_then(|m| m.elapsed_compute())
             .unwrap_or(0) as u64,
-    ));
-    exec_ctx
-        .baseline_metrics()
-        .elapsed_compute()
-        .add_duration(smj_time.duration());
+    );
+    smj_time.add_duration(merge_time);
+    elapsed_compute.add_duration(merge_time);
+    // Publish all sort metrics, including elapsed_compute, without counting it
+    // twice.
+    for metric in sort_metrics.iter() {
+        exec_ctx.execution_plan_metrics().register(metric.clone());
+    }
     Ok(())
 }