You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/23 09:15:40 UTC
(datafusion-comet) 01/01: fix: report output rows and time of the native window operator (#7140)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-7140-653c7954a485666a4b7c5487616b8755bbcf643c
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 333bae814b676e2c03fb49fd55d9a9c1a3f43de7
Author: Oleks V <co...@users.noreply.github.com>
AuthorDate: Fri Oct 23 08:55:16 2026 +0000
fix: report output rows and time of the native window operator (#7140)
* fix: report output rows and time of the native window operator
`CometWindowExec` declared only `dataSize` and `numPartitions`, which
nothing sets, so the time spent in the native window never reached the
Spark UI or the event log. Report the baseline metrics instead. When the
planner casts the window's results in a projection, add the window as an
additional native plan so that its time is reported with the
projection's.
* test: overwrite the temp directory the window metrics test writes to
* test: check the native window time is added to the cast projection's
Also update the window debugging example for the new SparkPlan construction.
* test: share the scan-driving loop between the planner tests
* test: leave the cast projection window to the native test
---
docs/source/contributor-guide/debugging.md | 18 ++-
native/core/src/execution/planner.rs | 175 ++++++++++++++++-----
native/core/src/execution/planner/window.rs | 23 ++-
.../apache/spark/sql/comet/CometWindowExec.scala | 7 +-
.../apache/comet/exec/CometWindowExecSuite.scala | 22 +++
5 files changed, 187 insertions(+), 58 deletions(-)
diff --git a/docs/source/contributor-guide/debugging.md b/docs/source/contributor-guide/debugging.md
index ee9ddbb8ea..8795854193 100644
--- a/docs/source/contributor-guide/debugging.md
+++ b/docs/source/contributor-guide/debugging.md
@@ -262,9 +262,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 {
@@ -288,14 +289,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:
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index dfcc0072a6..a47c19d1a2 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -3114,6 +3114,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;
@@ -3129,7 +3130,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::{
@@ -4304,27 +4308,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);
@@ -4344,6 +4330,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 {
@@ -4816,30 +4898,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() {
@@ -4888,6 +4958,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,
diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs
index 915ff02955..85499fac4b 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<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)))
}
}
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 7a798a67cf..69fa6e371a 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 98fd1ea1e2..76b39c34a5 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"))
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org