diff --git a/docs/source/contributor-guide/debugging.md b/docs/source/contributor-guide/debugging.md index b33a6849cd4..ba8f48ee31a 100644 --- a/docs/source/contributor-guide/debugging.md +++ b/docs/source/contributor-guide/debugging.md @@ -257,9 +257,10 @@ so `spark.comet.debug.enabled` and `COMET_LOG_LEVEL` do not affect it. In `WindowBuilder::build` in `native/core/src/execution/planner/window.rs`, the window operator builds `window_agg`, either a `BoundedWindowAggExec` or a `WindowAggExec` depending on whether every window expression runs with bounded -memory. It then optionally wraps `window_agg` in a casting `ProjectionExec` as -`final_plan` and hands `final_plan` to `SparkPlan::new`. To see the raw window -output, wrap `window_agg` right after it is built, before the cast projection: +memory. It then optionally wraps `window_agg` in a casting `ProjectionExec` and +builds the `SparkPlan` from either the projection or `window_agg`. To see the raw +window output, wrap `window_agg` right after it is built, before the cast +projection: ```rust let window_agg: Arc = if all_bounded { @@ -283,14 +284,15 @@ let window_agg: Arc = { Arc::new(DebugExecutionDataStream::new("window-output", window_agg)) }; -// ... `final_plan` is built from `window_agg`, and -// `SparkPlan::new(spark_plan.plan_id, final_plan, vec![child])` wraps it. +// ... the casting `ProjectionExec` is built from `window_agg`, and +// `SparkPlan::new_with_additional(spark_plan.plan_id, projection, vec![child], +// vec![window_agg])` wraps it. Without a cast, `SparkPlan::new` wraps `window_agg`. ``` Wrapping `window_agg` leaves the rest of the builder unchanged. To see the -output after the cast, wrap `final_plan` the same way just before the -`SparkPlan::new` call. The `let` rebinding is only compiled in when -`debug_assertions` is on, so release builds use the original binding. +output after the cast, wrap `projection` the same way just before the +`SparkPlan::new_with_additional` call. The `let` rebinding is only compiled in +when `debug_assertions` is on, so release builds use the original binding. Rebuild with `make core` and run a window test, for example: diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 69e4738dde3..b81ab5f76ae 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3051,6 +3051,7 @@ mod tests { FileGroup, FileScanConfigBuilder, FileSource, ParquetSource, }; use datafusion::error::DataFusionError; + use datafusion::execution::SendableRecordBatchStream; use datafusion::logical_expr::ScalarUDF; use datafusion::logical_expr::{AggregateUDF, EmitTo}; use datafusion::physical_expr::expressions::Column; @@ -3066,7 +3067,10 @@ mod tests { use tempfile::TempDir; use tokio::sync::mpsc; - use crate::execution::{operators::InputBatch, planner::PhysicalPlanner}; + use crate::execution::{ + operators::{InputBatch, ScanExec}, + planner::PhysicalPlanner, + }; use crate::execution::operators::ExecutionError; use crate::execution::planner::{ @@ -4241,27 +4245,9 @@ mod tests { .collect::>() .into_iter() }); - let mut stream = planned.native_plan.execute(0, session.task_ctx()).unwrap(); - let mut output_rows = 0; - while let Some(batch) = futures::future::poll_fn(|cx| { - let result = stream.poll_next_unpin(cx); - if result.is_pending() { - for (scan, input) in scans.iter_mut().zip(inputs.iter_mut()) { - if scan.batch.try_lock().unwrap().is_none() { - if let Some(batch) = input.next() { - scan.set_input_batch(batch); - cx.waker().wake_by_ref(); - } - } - } - } - result - }) - .await - { - output_rows += batch.unwrap().num_rows(); - } - assert_eq!(output_rows, 4); + let stream = planned.native_plan.execute(0, session.task_ctx()).unwrap(); + let batches = collect_with_inputs(stream, &mut scans, &mut inputs).await; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 4); let metrics = to_native_metric_node(&planned).unwrap(); assert_eq!(metrics.children.len(), 2); assert_eq!(metrics.metrics["output_rows"], 4); @@ -4281,6 +4267,102 @@ mod tests { } } + #[tokio::test] + async fn window_cast_projection_reports_window_time() { + use crate::execution::metrics::utils::to_native_metric_node; + use spark_operator::{ + lower_window_frame_bound::LowerFrameBoundStruct, + upper_window_frame_bound::UpperFrameBoundStruct, CurrentRow, LowerWindowFrameBound, + UnboundedPreceding, UpperWindowFrameBound, WindowExpr, WindowFrame, WindowFrameType, + WindowSpecDefinition, + }; + + let session = Arc::new(SessionContext::new()); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + // `row_number` returns UInt64 natively and Spark expects Int32, so the planner adds a + // casting projection above the window. + let row_number = WindowExpr { + built_in_window_function: Some(Expr { + expr_struct: Some(ScalarFunc(spark_expression::ScalarFunc { + func: "row_number".to_string(), + ..Default::default() + })), + query_context: None, + expr_id: None, + }), + spec: Some(WindowSpecDefinition { + frame_specification: Some(WindowFrame { + frame_type: WindowFrameType::Rows as i32, + lower_bound: Some(LowerWindowFrameBound { + lower_frame_bound_struct: Some(LowerFrameBoundStruct::UnboundedPreceding( + UnboundedPreceding {}, + )), + }), + upper_bound: Some(UpperWindowFrameBound { + upper_frame_bound_struct: Some(UpperFrameBoundStruct::CurrentRow( + CurrentRow {}, + )), + }), + }), + ..Default::default() + }), + result_type: Some(create_proto_datatype()), + ..Default::default() + }; + let operator = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window { + window_expr: vec![row_number], + order_by_list: vec![Expr { + expr_struct: Some(SortOrder(Box::new(spark_expression::SortOrder { + child: Some(Box::new(create_bound_reference(0))), + direction: 0, + null_ordering: 0, + }))), + query_context: None, + expr_id: None, + }], + ..Default::default() + }))), + ..Default::default() + }; + let (mut scans, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + assert_eq!(planned.additional_native_plans.len(), 1); + + let rows = 100; + let mut inputs = [vec![ + InputBatch::Batch( + vec![Arc::new(Int32Array::from((0..rows).collect::>())) as ArrayRef], + rows as usize, + ), + InputBatch::EOF, + ] + .into_iter()]; + let stream = planned.native_plan.execute(0, session.task_ctx()).unwrap(); + let batches = collect_with_inputs(stream, &mut scans, &mut inputs).await; + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + rows as usize + ); + + let elapsed = |plan: &Arc| { + plan.metrics() + .and_then(|m| m.elapsed_compute()) + .unwrap_or(0) as i64 + }; + let window_time = elapsed(&planned.additional_native_plans[0]); + let projection_time = elapsed(&planned.native_plan); + assert!(window_time > 0); + + let metrics = to_native_metric_node(&planned).unwrap(); + // Output rows come from the projection only, and the time covers both operators. + assert_eq!(metrics.metrics["output_rows"], rows as i64); + assert_eq!( + metrics.metrics["elapsed_compute"], + window_time + projection_time + ); + } + #[test] fn projection_rolls_up_same_plan_id_child_metrics() { let op_scan = Operator { @@ -4753,30 +4835,18 @@ mod tests { Arc::new(Int32Array::from(vec![20, 21, 22])) as ArrayRef, Arc::new(Int32Array::from(vec![30, 31, 32])) as ArrayRef, ]; - let mut input = vec![ + let mut inputs = [vec![ InputBatch::Batch(columns.clone(), 3), InputBatch::Batch(columns, 3), InputBatch::EOF, ] - .into_iter(); - let mut stream = planned + .into_iter()]; + let stream = planned .native_plan .execute(0, SessionContext::new().task_ctx()) .unwrap(); let mut rows = 0; - while let Some(batch) = futures::future::poll_fn(|cx| { - let result = stream.poll_next_unpin(cx); - if result.is_pending() && scans[0].batch.try_lock().unwrap().is_none() { - if let Some(batch) = input.next() { - scans[0].set_input_batch(batch); - cx.waker().wake_by_ref(); - } - } - result - }) - .await - { - let batch = batch.unwrap(); + for batch in collect_with_inputs(stream, &mut scans, &mut inputs).await { assert_eq!(batch.num_columns(), indices.len()); for row in 0..batch.num_rows() { for (column, &source) in indices.iter().enumerate() { @@ -4825,6 +4895,35 @@ mod tests { } } + /// Drives `stream` to completion, feeding each scan its next input batch whenever the stream + /// is pending and that scan has none, and returns the output batches. + async fn collect_with_inputs( + mut stream: SendableRecordBatchStream, + scans: &mut [ScanExec], + inputs: &mut [std::vec::IntoIter], + ) -> Vec { + let mut batches = vec![]; + while let Some(batch) = futures::future::poll_fn(|cx| { + let result = stream.poll_next_unpin(cx); + if result.is_pending() { + for (scan, input) in scans.iter_mut().zip(inputs.iter_mut()) { + if scan.batch.try_lock().unwrap().is_none() { + if let Some(batch) = input.next() { + scan.set_input_batch(batch); + cx.waker().wake_by_ref(); + } + } + } + } + result + }) + .await + { + batches.push(batch.unwrap()); + } + batches + } + fn create_scan() -> Operator { Operator { plan_id: 0, diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs index 915ff029553..85499fac4bd 100644 --- a/native/core/src/execution/planner/window.rs +++ b/native/core/src/execution/planner/window.rs @@ -152,7 +152,7 @@ impl OperatorBuilder for WindowBuilder { // Spark expects Int32). If any window expression carries a // `result_type` that differs from the actual output type, wrap the // aggregate in a projection that casts the mismatched columns. - let final_plan: Arc = { + let plan = { let agg_schema = window_agg.schema(); let input_field_count = input_schema.fields().len(); let mut needs_cast = false; @@ -180,17 +180,24 @@ impl OperatorBuilder for WindowBuilder { proj_exprs.push((expr, field.name().to_string())); } if needs_cast { - Arc::new(ProjectionExec::try_new(proj_exprs, window_agg)?) + let projection: Arc = Arc::new(ProjectionExec::try_new( + proj_exprs, + Arc::clone(&window_agg), + )?); + // The window is an additional plan so that its time is reported with the + // projection's rather than left out. + SparkPlan::new_with_additional( + spark_plan.plan_id, + projection, + vec![child], + vec![window_agg], + ) } else { - window_agg + SparkPlan::new(spark_plan.plan_id, window_agg, vec![child]) } }; - Ok(( - scans, - shuffle_scans, - Arc::new(SparkPlan::new(spark_plan.plan_id, final_plan, vec![child])), - )) + Ok((scans, shuffle_scans, Arc::new(plan))) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala index aefa252615f..bc8e097c5f4 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, Complete, Count, First, Last, Max, Min, Sum} import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, DataType, DateType, Decimal, DecimalType, DoubleType, LongType, MapType, NumericType, StructType} @@ -768,9 +768,8 @@ case class CometWindowExec( override def nodeName: String = "CometWindowExec" override def producedAttributes: AttributeSet = outputSet ++ AttributeSet(windowExpression) - override lazy val metrics: Map[String, SQLMetric] = Map( - "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"), - "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of partitions")) + override lazy val metrics: Map[String, SQLMetric] = + CometMetricNode.baselineMetrics(sparkContext) override def outputOrdering: Seq[SortOrder] = child.outputOrdering diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 63fe78298c2..836783673be 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -521,6 +521,28 @@ class CometWindowExecSuite extends CometTestBase { } } + test("Native window operator reports its output rows and time") { + withTempDir { dir => + (0 until 1000) + .map(i => (i % 10, i)) + .toDF("k", "v") + .write + .mode("overwrite") + .parquet(dir.toString) + withTempView("window_metrics") { + spark.read.parquet(dir.toString).createOrReplaceTempView("window_metrics") + // `max` needs no cast projection. The planner test in `planner.rs` covers the window + // that does, such as `row_number`. + val (_, plan) = checkSparkAnswerAndOperator( + sql("SELECT max(v) OVER (PARTITION BY k ORDER BY v) FROM window_metrics")) + val windows = collect(plan) { case w: CometWindowExec => w } + assert(windows.length == 1) + assert(windows.head.metrics("output_rows").value == 1000) + assert(windows.head.metrics("elapsed_compute").value > 0) + } + } + } + test("Window range frame with long boundary should not fail") { val df = Seq((1L, "1"), (1L, "1"), (2147483650L, "1"), (3L, "2"), (2L, "1"), (2147483650L, "2"))