fix: report output rows and time of the native window operator

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<dyn ExecutionPlan> = if all_bounded {
@@ -283,14 +284,15 @@ let window_agg: Arc<dyn ExecutionPlan> = {
     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:
 

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::<Vec<_>>()
                         .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::<usize>(), 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::<Vec<_>>())) 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::<usize>(),
+            rows as usize
+        );
+
+        let elapsed = |plan: &Arc<dyn ExecutionPlan>| {
+            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<InputBatch>],
+    ) -> Vec<RecordBatch> {
+        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,

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<dyn ExecutionPlan> = {
+        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<dyn ExecutionPlan> = 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)))
     }
 }
 

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
 

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"))