diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index b4d3d09069b14..fd424467ab72d 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -38,16 +38,23 @@ pub struct WindowAggState { /// The range that we calculate the window function pub window_frame_range: Range, pub window_frame_ctx: Option, - /// The index of the last row that its result is calculated inside the partition record batch buffer. + /// The index of the first row in the partition's buffered record batch + /// whose result has not been calculated yet; equivalently, the number + /// of buffered rows with calculated results. The next evaluation pass + /// resumes from this index. pub last_calculated_index: usize, /// The offset of the deleted row number pub offset_pruned_rows: usize, /// Stores the results calculated by window frame pub out_col: ArrayRef, - /// Keeps track of how many rows should be generated to be in sync with input record_batch. + /// Keeps track of how many rows should be generated to be in sync with input record_batch // (For each row in the input record batch we need to generate a window result). pub n_row_result_missing: usize, - /// Flag indicating whether we have received all data for this partition + /// Snapshot of [`PartitionBatchState::n_rows_received`] as of this + /// partition's last evaluation pass + pub n_rows_received: usize, + /// Snapshot of [`PartitionBatchState::is_end`] as of this partition's + /// last evaluation pass pub is_end: bool, } @@ -97,6 +104,7 @@ impl WindowAggState { } self.n_row_result_missing = partition_batch_state.record_batch.num_rows() - self.last_calculated_index; + self.n_rows_received = partition_batch_state.n_rows_received; self.is_end = partition_batch_state.is_end; Ok(()) } @@ -126,6 +134,24 @@ impl WindowAggState { all_rows_have_results && !partition_just_ended } + /// Returns true when the partition's input is unchanged since the last + /// evaluation pass: no new rows have arrived and the end flag has not + /// changed. + /// + /// Standard (non-aggregate) window functions read nothing outside their + /// partition, so we can skip re-evaluating them when their input hasn't + /// changed. Aggregate window functions cannot use this test to skip + /// evaluation, because they also consult the ORDER BY values of the most + /// recent input row across *all* partitions. + #[inline] + pub fn is_input_unchanged( + &self, + partition_batch_state: &PartitionBatchState, + ) -> bool { + self.n_rows_received == partition_batch_state.n_rows_received + && self.is_end == partition_batch_state.is_end + } + pub fn new(out_type: &DataType) -> Result { let empty_out_col = ScalarValue::try_from(out_type)?.to_array_of_size(0)?; Ok(Self { @@ -135,6 +161,7 @@ impl WindowAggState { offset_pruned_rows: 0, out_col: empty_out_col, n_row_result_missing: 0, + n_rows_received: 0, is_end: false, }) } @@ -273,6 +300,8 @@ impl WindowFrameContext { pub struct PartitionBatchState { /// The record batch belonging to current partition pub record_batch: RecordBatch, + /// Total number of rows this partition has ever received + pub n_rows_received: usize, /// Flag indicating whether we have received all data for this partition pub is_end: bool, /// Number of rows emitted for this partition since the last pruning pass @@ -283,14 +312,17 @@ impl PartitionBatchState { pub fn new(schema: SchemaRef) -> Self { Self { record_batch: RecordBatch::new_empty(schema), + n_rows_received: 0, is_end: false, n_out_row: 0, } } pub fn new_with_batch(batch: RecordBatch) -> Self { + let n_rows_received = batch.num_rows(); Self { record_batch: batch, + n_rows_received, is_end: false, n_out_row: 0, } @@ -299,6 +331,7 @@ impl PartitionBatchState { pub fn extend(&mut self, batch: &RecordBatch) -> Result<()> { self.record_batch = concat_batches(&self.record_batch.schema(), [&self.record_batch, batch])?; + self.n_rows_received += batch.num_rows(); Ok(()) } } @@ -697,7 +730,8 @@ fn check_equality(current: &[ScalarValue], target: &[ScalarValue]) -> Result (Vec, Vec) { let range_columns: Vec = vec![Arc::new(Float64Array::from(vec![ @@ -711,6 +745,62 @@ mod tests { (range_columns, sort_options) } + #[test] + fn test_is_input_unchanged() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("ts", DataType::UInt64, false)])); + let batch = |values: &[u64]| -> Result { + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt64Array::from(values.to_vec()))], + )?) + }; + let results = |n: usize| ScalarValue::UInt64(Some(0)).to_array_of_size(n); + + let mut partition = PartitionBatchState::new_with_batch(batch(&[1, 2, 3])?); + assert_eq!(partition.n_rows_received, 3); + + // A freshly created state has seen none of the partition's rows, so + // it must not report the input as unchanged. + let mut state = WindowAggState::new(&DataType::UInt64)?; + assert!(!state.is_input_unchanged(&partition)); + + // After an evaluation pass (here producing results for two of the + // three rows), the state has seen the partition's current input. + state.update(&results(2)?, &partition)?; + assert!(state.is_input_unchanged(&partition)); + + // New rows must be detected... + partition.extend(&batch(&[4])?)?; + assert_eq!(partition.n_rows_received, 4); + assert!(!state.is_input_unchanged(&partition)); + // ...until the next evaluation pass catches up again. + state.update(&results(2)?, &partition)?; + assert!(state.is_input_unchanged(&partition)); + + // Pruning rows whose results are already calculated shrinks the + // buffer but not the received-row count, so it must not make the + // input appear changed. + let n_prune = 2; + let batch_ref = &partition.record_batch; + partition.record_batch = batch_ref.slice(n_prune, batch_ref.num_rows() - n_prune); + // An evaluator would have advanced the frame at least past the + // pruned rows; prune_state requires this. + state.window_frame_range = n_prune..n_prune; + state.prune_state(n_prune); + assert_eq!(partition.n_rows_received, 4); + assert!(state.is_input_unchanged(&partition)); + + // The end-of-partition transition must trigger one final pass, even + // though no rows arrived. + partition.is_end = true; + assert!(!state.is_input_unchanged(&partition)); + state.update(&results(0)?, &partition)?; + assert!(state.is_input_unchanged(&partition)); + + Ok(()) + } + fn assert_group_ranges( window_frame: &Arc, expected_results: Vec<(Range, usize)>, diff --git a/datafusion/physical-expr/src/window/standard.rs b/datafusion/physical-expr/src/window/standard.rs index b70ad38e95a0a..2a1239793c80a 100644 --- a/datafusion/physical-expr/src/window/standard.rs +++ b/datafusion/physical-expr/src/window/standard.rs @@ -179,6 +179,11 @@ impl WindowExpr for StandardWindowExpr { published: false, }) }; + // Skip partitions whose input is unchanged since the last + // evaluation pass. + if window_state.state.is_input_unchanged(partition_batch_state) { + continue; + } let WindowFn::Builtin(evaluator) = &mut window_state.window_fn else { unreachable!() }; diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 11d0f677600ea..b52be42c5e356 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -1642,8 +1642,10 @@ mod tests { }; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::sum::sum_udaf; + use datafusion_functions_window::lead_lag::lead_udwf; use datafusion_functions_window::nth_value::last_value_udwf; use datafusion_functions_window::nth_value::nth_value_udwf; + use datafusion_functions_window::row_number::row_number_udwf; use datafusion_physical_expr::expressions::{Column, Literal, col}; use datafusion_physical_expr::window::{PartitionKey, StandardWindowExpr}; use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; @@ -2183,6 +2185,115 @@ mod tests { Ok(()) } + // In `Linear` mode, a partition may receive no new rows for several + // input batches while other partitions keep growing. The evaluation + // sweep skips a partition whose input is unchanged, so this test + // drives a partition through quiet batches and then resumes it: the + // results after the gap must continue from the retained evaluator + // state. ROW_NUMBER is causal, so the quiet partition is fully + // calculated while it waits; LEAD is not, so its result for the + // partition's last buffered row stays pending across the quiet + // batches and must materialize once the partition receives another + // row. + #[tokio::test] + async fn bounded_window_linear_quiet_partition_resume_standard() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::UInt64, false), + Field::new("ts", DataType::UInt64, false), + ])); + let make_batch = |rows: &[(u64, u64)]| -> Result { + let mut pk = UInt64Builder::with_capacity(rows.len()); + let mut ts = UInt64Builder::with_capacity(rows.len()); + for (p, t) in rows { + pk.append_value(*p); + ts.append_value(*t); + } + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(pk.finish()), Arc::new(ts.finish())], + )?) + }; + // `ts` ascends globally; partition 0 is absent from the middle batches. + let batches = vec![ + make_batch(&[(0, 0), (0, 1), (1, 2)])?, + make_batch(&[(1, 3), (1, 4)])?, + make_batch(&[(1, 5)])?, + make_batch(&[(0, 6), (1, 7)])?, + ]; + let memory_exec = + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + + let partition_by = vec![col("pk", &schema)?]; + let order_by = [PhysicalSortExpr { + expr: col("ts", &schema)?, + options: SortOptions::default(), + }]; + // Both functions use the default frame of a window with an ORDER BY + // clause (RANGE UNBOUNDED PRECEDING..CURRENT ROW). + let row_number_expr = create_window_expr( + &WindowFunctionDefinition::WindowUDF(row_number_udwf()), + "row_number".to_string(), + &[], + &partition_by, + &order_by, + Arc::new(WindowFrame::new(Some(false))), + Arc::clone(&schema), + false, + false, + None, + )?; + let lead_expr = create_window_expr( + &WindowFunctionDefinition::WindowUDF(lead_udwf()), + "lead".to_string(), + &[col("ts", &schema)?], + &partition_by, + &order_by, + Arc::new(WindowFrame::new(Some(false))), + Arc::clone(&schema), + false, + false, + None, + )?; + let physical_plan = BoundedWindowAggExec::try_new( + vec![row_number_expr, lead_expr], + memory_exec, + InputOrderMode::Linear, + true, + ) + .map(|e| Arc::new(e) as Arc)?; + + let batches = collect(physical_plan.execute(0, task_context())?).await?; + + // The skip must not delay results that are ready to be finalized; + // they stream out as soon as every window expression has produced + // them. LEAD holds back only the last buffered row of a partition, + // so one row unblocks right after the first input batch, five more + // when partition 0 resumes in the last input batch, and the final + // two (whose LEAD results need the end of the input) in the flush + // after the input is exhausted. + assert_eq!( + batches.iter().map(|b| b.num_rows()).collect::>(), + vec![1, 5, 2], + "expected results to stream as they become final" + ); + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+------------+------+ + | pk | ts | row_number | lead | + +----+----+------------+------+ + | 0 | 0 | 1 | 1 | + | 0 | 1 | 2 | 6 | + | 1 | 2 | 1 | 3 | + | 1 | 3 | 2 | 4 | + | 1 | 4 | 3 | 5 | + | 1 | 5 | 4 | 7 | + | 0 | 6 | 3 | | + | 1 | 7 | 5 | | + +----+----+------------+------+ + "); + Ok(()) + } + // This test, tests whether most recent row guarantee by the input batch of the `BoundedWindowAggExec` // helps `BoundedWindowAggExec` to generate low latency result in the `Linear` mode. // Input data generated at the source is