perf(functions-aggregate): optimize sliding window MIN/MAX using monotonic deques (#23826)#23827
perf(functions-aggregate): optimize sliding window MIN/MAX using monotonic deques (#23826)#23827pavan51 wants to merge 8 commits into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #23827 +/- ##
==========================================
- Coverage 80.72% 80.71% -0.01%
==========================================
Files 1089 1089
Lines 368911 368893 -18
Branches 368911 368893 -18
==========================================
- Hits 297785 297755 -30
- Misses 53374 53381 +7
- Partials 17752 17757 +5 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
neilconway
left a comment
neilconway
left a comment
There was a problem hiding this comment.
Thanks for the contribution @pavan51 ! I think this is directionally a good idea and a clear improvement over the previous implementation.
The PR fails some linter checks; can you fix those, please?
| // Licensed to the Apache Software Foundation (ASF) under one | ||
| // or more contributor license agreements. See the NOTICE file | ||
| // distributed with this work for additional information | ||
| // regarding copyright ownership. The ASF licenses this file | ||
| // to you under the Apache License, Version 2.0 (the | ||
| // "License"); you may not use this file except in compliance | ||
| // with the License. You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, | ||
| // software distributed under the License is distributed on an | ||
| // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| // KIND, either express or implied. See the License for the | ||
| // specific language governing permissions and limitations | ||
| // under the License. | ||
|
|
||
| use datafusion_common::ScalarValue; | ||
| use rand::Rng; | ||
| use std::collections::VecDeque; | ||
| use std::time::{Duration, Instant}; | ||
|
|
||
| // === Old Two-Stack Queue Implementation === | ||
| #[derive(Debug)] | ||
| pub struct TwoStackMax<T> { | ||
| push_stack: Vec<(T, T)>, | ||
| pop_stack: Vec<(T, T)>, | ||
| } | ||
|
|
||
| impl<T: Clone + PartialOrd> TwoStackMax<T> { | ||
| pub fn new() -> Self { | ||
| Self { | ||
| push_stack: Vec::new(), | ||
| pop_stack: Vec::new(), | ||
| } | ||
| } | ||
|
|
||
| pub fn push(&mut self, val: T) { | ||
| self.push_stack.push(match self.push_stack.last() { | ||
| Some((_, max)) => { | ||
| if val < *max { | ||
| (val, max.clone()) | ||
| } else { | ||
| (val.clone(), val) | ||
| } | ||
| } | ||
| None => (val.clone(), val), | ||
| }); | ||
| } | ||
|
|
||
| pub fn pop(&mut self) -> Option<T> { | ||
| if self.pop_stack.is_empty() { | ||
| match self.push_stack.pop() { | ||
| Some((val, _)) => { | ||
| let mut last = (val.clone(), val); | ||
| self.pop_stack.push(last.clone()); | ||
| while let Some((val, _)) = self.push_stack.pop() { | ||
| let max = if last.1 > val { | ||
| last.1.clone() | ||
| } else { | ||
| val.clone() | ||
| }; | ||
| last = (val.clone(), max); | ||
| self.pop_stack.push(last.clone()); | ||
| } | ||
| } | ||
| None => return None, | ||
| } | ||
| } | ||
| self.pop_stack.pop().map(|(val, _)| val) | ||
| } | ||
| } | ||
|
|
||
| // === Production Monotonic Deque Implementation === | ||
| #[derive(Debug)] | ||
| pub struct MonotonicMax<T> { | ||
| fifo: VecDeque<T>, | ||
| deque: VecDeque<T>, | ||
| } | ||
|
|
||
| impl<T: Clone + PartialOrd> MonotonicMax<T> { | ||
| pub fn new() -> Self { | ||
| Self { | ||
| fifo: VecDeque::new(), | ||
| deque: VecDeque::new(), | ||
| } | ||
| } | ||
|
|
||
| pub fn push(&mut self, val: T) { | ||
| self.fifo.push_back(val.clone()); | ||
| while self.deque.back().map_or(false, |back_val| *back_val < val) { | ||
| self.deque.pop_back(); | ||
| } | ||
| self.deque.push_back(val); | ||
| } | ||
|
|
||
| pub fn pop(&mut self) -> Option<T> { | ||
| if let Some(popped) = self.fifo.pop_front() { | ||
| if self | ||
| .deque | ||
| .front() | ||
| .map_or(false, |front_val| *front_val == popped) | ||
| { | ||
| self.deque.pop_front(); | ||
| } | ||
| Some(popped) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // === Generator Utilities === | ||
| fn generate_random_i64(size: usize) -> Vec<i64> { | ||
| let mut rng = rand::thread_rng(); | ||
| (0..size).map(|_| rng.gen_range(0..1000000)).collect() | ||
| } | ||
|
|
||
| fn generate_random_strings(size: usize) -> Vec<String> { | ||
| let mut rng = rand::thread_rng(); | ||
| (0..size) | ||
| .map(|_| { | ||
| let len = rng.gen_range(10..40); | ||
| (0..len) | ||
| .map(|_| rng.gen_range(b'a'..=b'z') as char) | ||
| .collect() | ||
| }) | ||
| .collect() | ||
| } | ||
|
|
||
| fn main() { | ||
| let dataset_size = 5000000; // 5M elements | ||
| println!("Generating raw inputs of size {}...", dataset_size); | ||
| let i64_raw = generate_random_i64(dataset_size); | ||
| let str_raw = generate_random_strings(dataset_size); | ||
|
|
||
| let scalar_int_data: Vec<ScalarValue> = i64_raw | ||
| .iter() | ||
| .map(|&val| ScalarValue::Int64(Some(val))) | ||
| .collect(); | ||
| let scalar_utf8_data: Vec<ScalarValue> = str_raw | ||
| .iter() | ||
| .map(|val| ScalarValue::Utf8(Some(val.clone()))) | ||
| .collect(); | ||
|
|
||
| let datasets = vec![ | ||
| ("scalar_int64", scalar_int_data), | ||
| ("scalar_utf8_string", scalar_utf8_data), | ||
| ]; | ||
|
|
||
| println!("Starting end-to-end latency measurements..."); | ||
| for (label, data) in datasets { | ||
| println!("\n=================================================="); | ||
| println!("DATATYPE: {}", label); | ||
| println!("=================================================="); | ||
|
|
||
| for window_size in [1000, 10000, 100000] { | ||
| println!(" --- Window Size: {} ---", window_size); | ||
|
|
||
| // Two-Stack Max Latency | ||
| { | ||
| let mut q = TwoStackMax::new(); | ||
| let mut max_push_time = Duration::ZERO; | ||
| let mut max_pop_time = Duration::ZERO; | ||
|
|
||
| let start_total = Instant::now(); | ||
| for (i, val) in data.iter().enumerate() { | ||
| let t0 = Instant::now(); | ||
| q.push(val.clone()); | ||
| let t_push = t0.elapsed(); | ||
| if t_push > max_push_time { | ||
| max_push_time = t_push; | ||
| } | ||
|
|
||
| if i >= window_size { | ||
| let t0 = Instant::now(); | ||
| q.pop(); | ||
| let t_pop = t0.elapsed(); | ||
| if t_pop > max_pop_time { | ||
| max_pop_time = t_pop; | ||
| } | ||
| } | ||
| } | ||
| let total_duration = start_total.elapsed(); | ||
| println!( | ||
| " Two-Stack Queue: Total Time = {:?}, Max Push = {:?}, Max Pop = {:?}", | ||
| total_duration, max_push_time, max_pop_time | ||
| ); | ||
| } | ||
|
|
||
| // Monotonic Deque Max Latency | ||
| { | ||
| let mut q = MonotonicMax::new(); | ||
| let mut max_push_time = Duration::ZERO; | ||
| let mut max_pop_time = Duration::ZERO; | ||
|
|
||
| let start_total = Instant::now(); | ||
| for (i, val) in data.iter().enumerate() { | ||
| let t0 = Instant::now(); | ||
| q.push(val.clone()); | ||
| let t_push = t0.elapsed(); | ||
| if t_push > max_push_time { | ||
| max_push_time = t_push; | ||
| } | ||
|
|
||
| if i >= window_size { | ||
| let t0 = Instant::now(); | ||
| q.pop(); | ||
| let t_pop = t0.elapsed(); | ||
| if t_pop > max_pop_time { | ||
| max_pop_time = t_pop; | ||
| } | ||
| } | ||
| } | ||
| let total_duration = start_total.elapsed(); | ||
| println!( | ||
| " Monotonic Deque: Total Time = {:?}, Max Push = {:?}, Max Pop = {:?}", | ||
| total_duration, max_push_time, max_pop_time | ||
| ); | ||
| } | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
It is useful to report these statistics in the PR discussion, but I think we should omit the file from the PR itself.
There was a problem hiding this comment.
Done. I've removed max_latency.rs from the PR and will post the benchmark results in the main PR discussion thread.
| // === Old Two-Stack Queue Implementation === | ||
| #[derive(Debug)] | ||
| pub struct TwoStackMax<T> { | ||
| push_stack: Vec<(T, T)>, | ||
| pop_stack: Vec<(T, T)>, | ||
| } |
There was a problem hiding this comment.
Generally, we like benchmarks to:
- Directly test the production implementation, not a copy of it. If necessary we can refactor the production code to make this possible.
- Only test the current version of the code, not prior versions or alternative implementations.
We'd then establish the performance delta of the PR by comparing the benchmark results with and without the optimization applied, rather than having both old and new algorithms in the benchmark code that is committed to main.
There was a problem hiding this comment.
Updated. The sliding_min_max.rs benchmark now only tests the production MovingMax implementation directly. To establish the performance delta, I ran this benchmark against both the main
branch and this branch individually(results posted in the main discussion).
There was a problem hiding this comment.
This comment needs updating.
There was a problem hiding this comment.
Updated. The documentation now accurately describes the new monotonic deque approach and sequence numbers instead of the old two-stack algorithm.
|
|
||
| // === Generator Utilities === | ||
| fn generate_random_i64(size: usize) -> Vec<i64> { | ||
| let mut rng = rand::thread_rng(); |
There was a problem hiding this comment.
We should see the RNG to ensure determinism; see how the other benchmarks work as an example. Also rand::rngs::StdRng not thread_rng.
There was a problem hiding this comment.
Fixed. I've updated the benchmark to use StdRng::seed_from_u64(42) for deterministic data generation across all types.
| Self { | ||
| push_stack: Vec::with_capacity(capacity), | ||
| pop_stack: Vec::with_capacity(capacity), | ||
| fifo: std::collections::VecDeque::with_capacity(capacity), |
There was a problem hiding this comment.
use std::collections::VecDeque would be a bit more readable.
There was a problem hiding this comment.
Done. Added the top-level import and updated all usages for better readability.
| /// Removes and returns the last value of the sliding window. | ||
| #[inline] | ||
| pub fn pop(&mut self) -> Option<T> { |
There was a problem hiding this comment.
I notice that we don't actually use the return value of pop in the sliding window accumulators. If we stop having pop return a value, I wonder if we can simplify the implementation -- e.g., get rid of fifo and just use a single deque that stores candidates in the window and their positions / sequence numbers. Something like
structure MovingMax<T>:
candidates : deque of (seq, value) # front = oldest; values strictly decreasing
push_seq : integer = 0 # total elements ever pushed
pop_seq : integer = 0 # total elements ever popped
# window = elements with seq in [pop_seq, push_seq)
On push(), we evict all the candidates that are dominated by the newly pushed value, if any.
On pop(), if the popped element is the front of candidates (based on sequence number, not value), remove it.
There was a problem hiding this comment.
I have refactored the implementation to use this sequence number approach. The fifo queue has been removed entirely, and pop() now returns (), which significantly simplified the code.
82168be to
31cf348
Compare
Fixed. I've resolved all the clippy::unnecessary_map_or warnings. |
|
@neilconway Thanks for the detailed review and excellent feedback. I've addressed all the comments. The sliding_min_max.rs benchmark now only tests the production MovingMax implementation directly. To establish the performance delta, I ran the benchmark against both the main branch (old two-stack implementation) and this new branch (monotonic deque sequence number). Here are the results: Key Metrics (Old vs New)1. 64-bit Integers (
|
3 participants
Which issue does this PR close?
Rationale for this change
Summary
This PR replaces the two-stack queue-based sliding window minimum (
MovingMin) and maximum (MovingMax) implementations with a double-ended queue-based Monotonic Deque implementation.This optimization provides several key benefits:
pop()latency ofpop_stackwas empty, because it had to drain and reverse thepush_stack. The new monotonic deque design ensures thatpop()andevaluate()(min()/max()) are strictlyTisScalarValuecontaining heap-allocated elements like strings, list vectors, or decimals).Implementation Details
fifo: A standardVecDeque<T>that preserves the exact FIFO ordering of the active window elements to support thepop()API.deque: A monotonicVecDeque<T>that retains only the active candidates (strictly increasing forMovingMin, strictly decreasing forMovingMax).MovingMin<T>,MovingMax<T>) and methods, preserving backwards compatibility and passing all existing unit and doc tests without changes.Benchmark Comparison (Monotonic Deque vs Two-Stack Queue)
MovingMax/MovingMin)push()Time Complexitypop()Time Complexitymax()/min()Time Complexity.clone()calls during stack re-shufflesBenchmark Results (50,000 Elements)
Running the criterion benchmark suite on random datasets (
sliding_min_max) comparingTwoStackMax<ScalarValue>againstMonotonicMax<ScalarValue>.Because DataFusion executes sliding window aggregate states using the
ScalarValueenum wrapper (not raw primitive Rust types), these benchmarks reflect real-world performance:1.
ScalarValue::Int64(Standard Integer Columns)1.661 ms1.220 ms(1.36x speedup)1.740 ms1.288 ms(1.35x speedup)1.664 ms1.208 ms(1.38x speedup)2.
ScalarValue::Float64(Standard Float Columns)1.740 ms1.290 ms(1.35x speedup)1.820 ms1.267 ms(1.44x speedup)1.742 ms1.294 ms(1.35x speedup)3.
ScalarValue::TimestampNanosecond(Time-Series / Timestamp Columns)1.847 ms1.291 ms(1.43x speedup)1.897 ms1.276 ms(1.49x speedup)1.801 ms1.291 ms(1.40x speedup)4.
ScalarValue::Decimal128(Financial / Precise Decimal Columns)1.832 ms1.290 ms(1.42x speedup)1.908 ms1.362 ms(1.40x speedup)1.916 ms1.315 ms(1.46x speedup)5.
ScalarValue::Utf8(String / Text Columns)10.987 ms4.756 ms(2.31x speedup)10.453 ms4.788 ms(2.18x speedup)10.691 ms5.009 ms(2.13x speedup)Scaled Benchmark Results (5,000,000 Elements)
To verify how performance behaves under extreme scale, we ran an end-to-end total execution time and worst-case latency benchmark on a dataset of 5,000,000 elements:
1. Total Execution Time Comparison (5M Elements)
As expected, the overall throughput speedup multiplier remains highly consistent as the dataset size scales to 5 million records, showing no performance fading:
ScalarValue::Int64(Standard Integer Columns):675.46 ms| Monotonic Deque:466.98 ms(1.44x speedup)545.11 ms| Monotonic Deque:465.74 ms(1.17x speedup)565.46 ms| Monotonic Deque:514.02 ms(1.10x speedup)ScalarValue::Utf8(String / Text Columns):1.428 s| Monotonic Deque:0.736 s(1.94x speedup)1.336 s| Monotonic Deque:0.732 s(1.82x speedup)1.327 s| Monotonic Deque:0.830 s(1.60x speedup)2. Worst-Case Single-Row Pop Latency Comparison (At$W = 100,000$ )
At a massive window size of 100,000, the worst-case single-row pop latency (peak thread pause/jitter) of the Two-Stack Queue balloons to 6.4 milliseconds for integers and 20.2 milliseconds for strings due to linear stack-reversal overhead. The Monotonic Deque minimizes this overhead:
ScalarValue::Int646.486 ms0.131 ms(130.9 µs)ScalarValue::Utf8(String)20.240 ms2.654 msAnalysis of freezes (jitter) vs total execution time:
What changes are included in this PR?
MovingMin<T>andMovingMax<T>indatafusion/functions-aggregate/src/min_max.rsto use a Monotonic Deque (via twoVecDeques: one for the FIFO window, one for the monotonic candidate values).sliding_min_max.rs) for 5 majorScalarValuetypes (Int64,Float64,TimestampNanosecond,Decimal128,Utf8).max_latency.rs) to prove the elimination ofAre these changes tested?
Yes.
datafusion-functions-aggregatepass perfectly, ensuring 100% correctness and parity with the previous implementation.datafusion/functions-aggregate/benches/and registered inCargo.toml.Are there any user-facing changes?
No