diff --git a/native-engine/datafusion-ext-plans/src/broadcast_join_build_hash_map_exec.rs b/native-engine/datafusion-ext-plans/src/broadcast_join_build_hash_map_exec.rs index 666248501..428ae2325 100644 --- a/native-engine/datafusion-ext-plans/src/broadcast_join_build_hash_map_exec.rs +++ b/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::>(); + 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::() + .expect("probe values"); + let build_values = batch + .column(usize::from(!build_left)) + .as_any() + .downcast_ref::() + .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::(), + 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::>(); + 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); diff --git a/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs b/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs index 485f420a9..33527f326 100644 --- a/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs +++ b/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs @@ -437,6 +437,9 @@ async fn execute_join_with_smj_fallback( exec_ctx: Arc, sender: Arc, ) -> 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 = { @@ -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,7 +473,7 @@ 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 ), ), @@ -476,7 +481,7 @@ async fn execute_join_with_smj_fallback( 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(()) }