Skip to main content

datafusion_expr/
window_state.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Structures used to hold window function state (for implementing WindowUDFs)
19
20use std::{collections::VecDeque, ops::Range, sync::Arc};
21
22use crate::{WindowFrame, WindowFrameBound, WindowFrameUnits};
23
24use arrow::{
25    array::ArrayRef,
26    compute::{SortOptions, concat, concat_batches},
27    datatypes::{DataType, SchemaRef},
28    record_batch::RecordBatch,
29};
30use datafusion_common::{
31    Result, ScalarValue, internal_datafusion_err, internal_err,
32    utils::{compare_rows, get_row_at_idx, search_in_slice},
33};
34
35/// Holds the state of evaluating a window function
36#[derive(Debug, Clone)]
37pub struct WindowAggState {
38    /// The range that we calculate the window function
39    pub window_frame_range: Range<usize>,
40    pub window_frame_ctx: Option<WindowFrameContext>,
41    /// The index of the last row that its result is calculated inside the partition record batch buffer.
42    pub last_calculated_index: usize,
43    /// The offset of the deleted row number
44    pub offset_pruned_rows: usize,
45    /// Stores the results calculated by window frame
46    pub out_col: ArrayRef,
47    /// Keeps track of how many rows should be generated to be in sync with input record_batch.
48    // (For each row in the input record batch we need to generate a window result).
49    pub n_row_result_missing: usize,
50    /// Flag indicating whether we have received all data for this partition
51    pub is_end: bool,
52}
53
54impl WindowAggState {
55    pub fn prune_state(&mut self, n_prune: usize) {
56        self.window_frame_range = Range {
57            start: self.window_frame_range.start - n_prune,
58            end: self.window_frame_range.end - n_prune,
59        };
60        self.last_calculated_index -= n_prune;
61        self.offset_pruned_rows += n_prune;
62
63        match self.window_frame_ctx.as_mut() {
64            // Rows have no state do nothing
65            Some(WindowFrameContext::Rows(_)) => {}
66            Some(WindowFrameContext::Range { .. }) => {}
67            Some(WindowFrameContext::Groups { state, .. }) => {
68                let mut n_group_to_del = 0;
69                for (_, end_idx) in &state.group_end_indices {
70                    if n_prune < *end_idx {
71                        break;
72                    }
73                    n_group_to_del += 1;
74                }
75                state.group_end_indices.drain(0..n_group_to_del);
76                state
77                    .group_end_indices
78                    .iter_mut()
79                    .for_each(|(_, start_idx)| *start_idx -= n_prune);
80                state.current_group_idx -= n_group_to_del;
81            }
82            None => {}
83        };
84    }
85
86    pub fn update(
87        &mut self,
88        out_col: &ArrayRef,
89        partition_batch_state: &PartitionBatchState,
90    ) -> Result<()> {
91        self.last_calculated_index += out_col.len();
92        // no need to use concat if the current `out_col` is empty
93        if self.out_col.is_empty() {
94            self.out_col = Arc::clone(out_col);
95        } else {
96            self.out_col = concat(&[&self.out_col, &out_col])?;
97        }
98        self.n_row_result_missing =
99            partition_batch_state.record_batch.num_rows() - self.last_calculated_index;
100        self.is_end = partition_batch_state.is_end;
101        Ok(())
102    }
103
104    /// Returns true when this state is fully up to date with the partition's
105    /// buffered batch, meaning another evaluation pass over the partition could
106    /// not produce any new results or change any state:
107    ///
108    /// - `last_calculated_index` has reached the end of the partition's
109    ///   buffered batch, so every row of this partition that has arrived so
110    ///   far already has a result.
111    /// - When a partition ends, a final evaluation pass is needed to bring
112    ///   derived state up to date.
113    #[inline]
114    pub fn is_up_to_date_with(
115        &self,
116        partition_batch_state: &PartitionBatchState,
117    ) -> bool {
118        let all_rows_have_results =
119            self.last_calculated_index == partition_batch_state.record_batch.num_rows();
120        if all_rows_have_results {
121            debug_assert_eq!(self.n_row_result_missing, 0);
122        }
123
124        // `self.is_end` holds the flag as of the previous evaluation pass.
125        let partition_just_ended = !self.is_end && partition_batch_state.is_end;
126        all_rows_have_results && !partition_just_ended
127    }
128
129    pub fn new(out_type: &DataType) -> Result<Self> {
130        let empty_out_col = ScalarValue::try_from(out_type)?.to_array_of_size(0)?;
131        Ok(Self {
132            window_frame_range: Range { start: 0, end: 0 },
133            window_frame_ctx: None,
134            last_calculated_index: 0,
135            offset_pruned_rows: 0,
136            out_col: empty_out_col,
137            n_row_result_missing: 0,
138            is_end: false,
139        })
140    }
141}
142
143/// This object stores the window frame state for use in incremental calculations.
144#[derive(Debug, Clone)]
145pub enum WindowFrameContext {
146    /// ROWS frames are inherently stateless.
147    Rows(Arc<WindowFrame>),
148    /// RANGE frames are stateful, they store indices specifying where the
149    /// previous search left off. This amortizes the overall cost to O(n)
150    /// where n denotes the row count.
151    Range {
152        window_frame: Arc<WindowFrame>,
153        state: WindowFrameStateRange,
154    },
155    /// GROUPS frames are stateful, they store group boundaries and indices
156    /// specifying where the previous search left off. This amortizes the
157    /// overall cost to O(n) where n denotes the row count.
158    Groups {
159        window_frame: Arc<WindowFrame>,
160        state: WindowFrameStateGroups,
161    },
162}
163
164impl WindowFrameContext {
165    /// Create a new state object for the given window frame.
166    pub fn new(window_frame: Arc<WindowFrame>, sort_options: Vec<SortOptions>) -> Self {
167        match window_frame.units {
168            WindowFrameUnits::Rows => WindowFrameContext::Rows(window_frame),
169            WindowFrameUnits::Range => WindowFrameContext::Range {
170                window_frame,
171                state: WindowFrameStateRange::new(sort_options),
172            },
173            WindowFrameUnits::Groups => WindowFrameContext::Groups {
174                window_frame,
175                state: WindowFrameStateGroups::default(),
176            },
177        }
178    }
179
180    /// This function calculates beginning/ending indices for the frame of the current row.
181    pub fn calculate_range(
182        &mut self,
183        range_columns: &[ArrayRef],
184        last_range: &Range<usize>,
185        length: usize,
186        idx: usize,
187    ) -> Result<Range<usize>> {
188        match self {
189            WindowFrameContext::Rows(window_frame) => {
190                Self::calculate_range_rows(window_frame, length, idx)
191            }
192            // Sort options is used in RANGE mode calculations because the
193            // ordering or position of NULLs impact range calculations and
194            // comparison of rows.
195            WindowFrameContext::Range {
196                window_frame,
197                state,
198            } => state.calculate_range(
199                window_frame,
200                last_range,
201                range_columns,
202                length,
203                idx,
204            ),
205            // Sort options is not used in GROUPS mode calculations as the
206            // inequality of two rows indicates a group change, and ordering
207            // or position of NULLs do not impact inequality.
208            WindowFrameContext::Groups {
209                window_frame,
210                state,
211            } => state.calculate_range(window_frame, range_columns, length, idx),
212        }
213    }
214
215    /// This function calculates beginning/ending indices for the frame of the current row.
216    fn calculate_range_rows(
217        window_frame: &Arc<WindowFrame>,
218        length: usize,
219        idx: usize,
220    ) -> Result<Range<usize>> {
221        let start = match window_frame.start_bound {
222            // UNBOUNDED PRECEDING
223            WindowFrameBound::Preceding(ScalarValue::UInt64(None)) => 0,
224            WindowFrameBound::Preceding(ScalarValue::UInt64(Some(n))) => {
225                idx.saturating_sub(n as usize)
226            }
227            WindowFrameBound::CurrentRow => idx,
228            // UNBOUNDED FOLLOWING
229            WindowFrameBound::Following(ScalarValue::UInt64(None)) => {
230                return internal_err!(
231                    "Frame start cannot be UNBOUNDED FOLLOWING '{window_frame:?}'"
232                );
233            }
234            WindowFrameBound::Following(ScalarValue::UInt64(Some(n))) => {
235                std::cmp::min(idx + n as usize, length)
236            }
237            // ERRONEOUS FRAMES
238            WindowFrameBound::Preceding(_) | WindowFrameBound::Following(_) => {
239                return internal_err!("Rows should be UInt64");
240            }
241        };
242        let end = match window_frame.end_bound {
243            // UNBOUNDED PRECEDING
244            WindowFrameBound::Preceding(ScalarValue::UInt64(None)) => {
245                return internal_err!(
246                    "Frame end cannot be UNBOUNDED PRECEDING '{window_frame:?}'"
247                );
248            }
249            WindowFrameBound::Preceding(ScalarValue::UInt64(Some(n))) => {
250                if idx >= n as usize {
251                    idx - n as usize + 1
252                } else {
253                    0
254                }
255            }
256            WindowFrameBound::CurrentRow => idx + 1,
257            // UNBOUNDED FOLLOWING
258            WindowFrameBound::Following(ScalarValue::UInt64(None)) => length,
259            WindowFrameBound::Following(ScalarValue::UInt64(Some(n))) => {
260                std::cmp::min(idx + n as usize + 1, length)
261            }
262            // ERRONEOUS FRAMES
263            WindowFrameBound::Preceding(_) | WindowFrameBound::Following(_) => {
264                return internal_err!("Rows should be UInt64");
265            }
266        };
267        Ok(Range { start, end })
268    }
269}
270
271/// State for each unique partition determined according to PARTITION BY column(s)
272#[derive(Debug, Clone, PartialEq)]
273pub struct PartitionBatchState {
274    /// The record batch belonging to current partition
275    pub record_batch: RecordBatch,
276    /// Flag indicating whether we have received all data for this partition
277    pub is_end: bool,
278    /// Number of rows emitted for this partition since the last pruning pass
279    pub n_out_row: usize,
280}
281
282impl PartitionBatchState {
283    pub fn new(schema: SchemaRef) -> Self {
284        Self {
285            record_batch: RecordBatch::new_empty(schema),
286            is_end: false,
287            n_out_row: 0,
288        }
289    }
290
291    pub fn new_with_batch(batch: RecordBatch) -> Self {
292        Self {
293            record_batch: batch,
294            is_end: false,
295            n_out_row: 0,
296        }
297    }
298
299    pub fn extend(&mut self, batch: &RecordBatch) -> Result<()> {
300        self.record_batch =
301            concat_batches(&self.record_batch.schema(), [&self.record_batch, batch])?;
302        Ok(())
303    }
304}
305
306/// This structure encapsulates all the state information we require as we scan
307/// ranges of data while processing RANGE frames.
308/// Attribute `sort_options` stores the column ordering specified by the ORDER
309/// BY clause. This information is used to calculate the range.
310#[derive(Debug, Default, Clone)]
311pub struct WindowFrameStateRange {
312    sort_options: Vec<SortOptions>,
313}
314
315impl WindowFrameStateRange {
316    /// Create a new object to store the search state.
317    fn new(sort_options: Vec<SortOptions>) -> Self {
318        Self { sort_options }
319    }
320
321    /// This function calculates beginning/ending indices for the frame of the current row.
322    // Argument `last_range` stores the resulting indices from the previous search. Since the indices only
323    // advance forward, we start from `last_range` subsequently. Thus, the overall
324    // time complexity of linear search amortizes to O(n) where n denotes the total
325    // row count.
326    fn calculate_range(
327        &mut self,
328        window_frame: &Arc<WindowFrame>,
329        last_range: &Range<usize>,
330        range_columns: &[ArrayRef],
331        length: usize,
332        idx: usize,
333    ) -> Result<Range<usize>> {
334        let start = match window_frame.start_bound {
335            WindowFrameBound::Preceding(ref n) => {
336                if n.is_null() {
337                    // UNBOUNDED PRECEDING
338                    0
339                } else {
340                    self.calculate_index_of_row::<true, true>(
341                        range_columns,
342                        last_range,
343                        idx,
344                        Some(n),
345                        length,
346                    )?
347                }
348            }
349            WindowFrameBound::CurrentRow => self.calculate_index_of_row::<true, true>(
350                range_columns,
351                last_range,
352                idx,
353                None,
354                length,
355            )?,
356            WindowFrameBound::Following(ref n) => self
357                .calculate_index_of_row::<true, false>(
358                    range_columns,
359                    last_range,
360                    idx,
361                    Some(n),
362                    length,
363                )?,
364        };
365        let end = match window_frame.end_bound {
366            WindowFrameBound::Preceding(ref n) => self
367                .calculate_index_of_row::<false, true>(
368                    range_columns,
369                    last_range,
370                    idx,
371                    Some(n),
372                    length,
373                )?,
374            WindowFrameBound::CurrentRow => self.calculate_index_of_row::<false, false>(
375                range_columns,
376                last_range,
377                idx,
378                None,
379                length,
380            )?,
381            WindowFrameBound::Following(ref n) => {
382                if n.is_null() {
383                    // UNBOUNDED FOLLOWING
384                    length
385                } else {
386                    self.calculate_index_of_row::<false, false>(
387                        range_columns,
388                        last_range,
389                        idx,
390                        Some(n),
391                        length,
392                    )?
393                }
394            }
395        };
396        Ok(Range { start, end })
397    }
398
399    /// This function does the heavy lifting when finding range boundaries. It is meant to be
400    /// called twice, in succession, to get window frame start and end indices (with `SIDE`
401    /// supplied as true and false, respectively).
402    fn calculate_index_of_row<const SIDE: bool, const SEARCH_SIDE: bool>(
403        &mut self,
404        range_columns: &[ArrayRef],
405        last_range: &Range<usize>,
406        idx: usize,
407        delta: Option<&ScalarValue>,
408        length: usize,
409    ) -> Result<usize> {
410        let current_row_values = get_row_at_idx(range_columns, idx)?;
411        let search_start = if SIDE {
412            last_range.start
413        } else {
414            last_range.end
415        };
416        let end_range = if let Some(delta) = delta {
417            let is_descending: bool = self
418                .sort_options
419                .first()
420                .ok_or_else(|| {
421                    internal_datafusion_err!(
422                        "Sort options unexpectedly absent in a window frame"
423                    )
424                })?
425                .descending;
426
427            // On overflow the boundary exceeds the type's range and is
428            // effectively unbounded within the partition. Collapse to the
429            // partition edge rather than feeding `search_in_slice` a
430            // wrapped-around target: PRECEDING searches reach `search_start`,
431            // FOLLOWING searches reach `length`.
432            let unbounded_edge = if SEARCH_SIDE { search_start } else { length };
433            let mut targets = Vec::with_capacity(current_row_values.len());
434            for value in &current_row_values {
435                if value.is_null() {
436                    targets.push(value.clone());
437                    continue;
438                }
439                let target = if SEARCH_SIDE == is_descending {
440                    match value.add_checked(delta) {
441                        Ok(v) => v,
442                        Err(_) => return Ok(unbounded_edge),
443                    }
444                } else if value.is_unsigned() && value < delta {
445                    // NOTE: This gets a polymorphic zero without having long coercion code for ScalarValue.
446                    //       If we decide to implement a "default" construction mechanism for ScalarValue,
447                    //       change the following statement to use that.
448                    value.sub(value)?
449                } else {
450                    match value.sub_checked(delta) {
451                        Ok(v) => v,
452                        Err(_) => return Ok(unbounded_edge),
453                    }
454                };
455                targets.push(target);
456            }
457            targets
458        } else {
459            current_row_values
460        };
461        let compare_fn = |current: &[ScalarValue], target: &[ScalarValue]| {
462            let cmp = compare_rows(current, target, &self.sort_options)?;
463            Ok(if SIDE { cmp.is_lt() } else { cmp.is_le() })
464        };
465        search_in_slice(range_columns, &end_range, compare_fn, search_start, length)
466    }
467}
468
469// In GROUPS mode, rows with duplicate sorting values are grouped together.
470// Therefore, there must be an ORDER BY clause in the window definition to use GROUPS mode.
471// The syntax is as follows:
472//     GROUPS frame_start [ frame_exclusion ]
473//     GROUPS BETWEEN frame_start AND frame_end [ frame_exclusion ]
474// The optional frame_exclusion specifier is not yet supported.
475// The frame_start and frame_end parameters allow us to specify which rows the window
476// frame starts and ends with. They accept the following values:
477//    - UNBOUNDED PRECEDING: Start with the first row of the partition. Possible only in frame_start.
478//    - offset PRECEDING: When used in frame_start, it refers to the first row of the group
479//                        that comes "offset" groups before the current group (i.e. the group
480//                        containing the current row). When used in frame_end, it refers to the
481//                        last row of the group that comes "offset" groups before the current group.
482//    - CURRENT ROW: When used in frame_start, it refers to the first row of the group containing
483//                   the current row. When used in frame_end, it refers to the last row of the group
484//                   containing the current row.
485//    - offset FOLLOWING: When used in frame_start, it refers to the first row of the group
486//                        that comes "offset" groups after the current group (i.e. the group
487//                        containing the current row). When used in frame_end, it refers to the
488//                        last row of the group that comes "offset" groups after the current group.
489//    - UNBOUNDED FOLLOWING: End with the last row of the partition. Possible only in frame_end.
490
491/// This structure encapsulates all the state information we require as we
492/// scan groups of data while processing window frames.
493#[derive(Debug, Default, Clone)]
494pub struct WindowFrameStateGroups {
495    /// A tuple containing group values and the row index where the group ends.
496    /// Example: [[1, 1], [1, 1], [2, 1], [2, 1], ...] would correspond to
497    ///          [([1, 1], 2), ([2, 1], 4), ...].
498    pub group_end_indices: VecDeque<(Vec<ScalarValue>, usize)>,
499    /// The group index to which the row index belongs.
500    pub current_group_idx: usize,
501}
502
503impl WindowFrameStateGroups {
504    fn calculate_range(
505        &mut self,
506        window_frame: &Arc<WindowFrame>,
507        range_columns: &[ArrayRef],
508        length: usize,
509        idx: usize,
510    ) -> Result<Range<usize>> {
511        let start = match window_frame.start_bound {
512            WindowFrameBound::Preceding(ref n) => {
513                if n.is_null() {
514                    // UNBOUNDED PRECEDING
515                    0
516                } else {
517                    self.calculate_index_of_row::<true, true>(
518                        range_columns,
519                        idx,
520                        Some(n),
521                        length,
522                    )?
523                }
524            }
525            WindowFrameBound::CurrentRow => self.calculate_index_of_row::<true, true>(
526                range_columns,
527                idx,
528                None,
529                length,
530            )?,
531            WindowFrameBound::Following(ref n) => self
532                .calculate_index_of_row::<true, false>(
533                    range_columns,
534                    idx,
535                    Some(n),
536                    length,
537                )?,
538        };
539        let end = match window_frame.end_bound {
540            WindowFrameBound::Preceding(ref n) => self
541                .calculate_index_of_row::<false, true>(
542                    range_columns,
543                    idx,
544                    Some(n),
545                    length,
546                )?,
547            WindowFrameBound::CurrentRow => self.calculate_index_of_row::<false, false>(
548                range_columns,
549                idx,
550                None,
551                length,
552            )?,
553            WindowFrameBound::Following(ref n) => {
554                if n.is_null() {
555                    // UNBOUNDED FOLLOWING
556                    length
557                } else {
558                    self.calculate_index_of_row::<false, false>(
559                        range_columns,
560                        idx,
561                        Some(n),
562                        length,
563                    )?
564                }
565            }
566        };
567        Ok(Range { start, end })
568    }
569
570    /// This function does the heavy lifting when finding range boundaries. It is meant to be
571    /// called twice, in succession, to get window frame start and end indices (with `SIDE`
572    /// supplied as true and false, respectively). Generic argument `SEARCH_SIDE` determines
573    /// the sign of `delta` (where true/false represents negative/positive respectively).
574    fn calculate_index_of_row<const SIDE: bool, const SEARCH_SIDE: bool>(
575        &mut self,
576        range_columns: &[ArrayRef],
577        idx: usize,
578        delta: Option<&ScalarValue>,
579        length: usize,
580    ) -> Result<usize> {
581        let delta = if let Some(delta) = delta {
582            if let ScalarValue::UInt64(Some(value)) = delta {
583                *value as usize
584            } else {
585                return internal_err!(
586                    "Unexpectedly got a non-UInt64 value in a GROUPS mode window frame"
587                );
588            }
589        } else {
590            0
591        };
592        let mut group_start = 0;
593        let last_group = self.group_end_indices.back_mut();
594        if let Some((group_row, group_end)) = last_group {
595            if *group_end < length {
596                let new_group_row = get_row_at_idx(range_columns, *group_end)?;
597                // If last/current group keys are the same, we extend the last group:
598                if new_group_row.eq(group_row) {
599                    // Update the end boundary of the group (search right boundary):
600                    *group_end = search_in_slice(
601                        range_columns,
602                        group_row,
603                        check_equality,
604                        *group_end,
605                        length,
606                    )?;
607                }
608            }
609            // Start searching from the last group boundary:
610            group_start = *group_end;
611        }
612
613        // Advance groups until `idx` is inside a group:
614        while idx >= group_start {
615            let group_row = get_row_at_idx(range_columns, group_start)?;
616            // Find end boundary of the group (search right boundary):
617            let group_end = search_in_slice(
618                range_columns,
619                &group_row,
620                check_equality,
621                group_start,
622                length,
623            )?;
624            self.group_end_indices.push_back((group_row, group_end));
625            group_start = group_end;
626        }
627
628        // Update the group index `idx` belongs to:
629        while self.current_group_idx < self.group_end_indices.len()
630            && idx >= self.group_end_indices[self.current_group_idx].1
631        {
632            self.current_group_idx += 1;
633        }
634
635        // Find the group index of the frame boundary:
636        let group_idx = if SEARCH_SIDE {
637            self.current_group_idx.saturating_sub(delta)
638        } else {
639            self.current_group_idx + delta
640        };
641
642        // Extend `group_start_indices` until it includes at least `group_idx`:
643        while self.group_end_indices.len() <= group_idx && group_start < length {
644            let group_row = get_row_at_idx(range_columns, group_start)?;
645            // Find end boundary of the group (search right boundary):
646            let group_end = search_in_slice(
647                range_columns,
648                &group_row,
649                check_equality,
650                group_start,
651                length,
652            )?;
653            self.group_end_indices.push_back((group_row, group_end));
654            group_start = group_end;
655        }
656
657        // Calculate index of the group boundary:
658        Ok(match (SIDE, SEARCH_SIDE) {
659            // Window frame start:
660            (true, _) => {
661                let group_idx = std::cmp::min(group_idx, self.group_end_indices.len());
662                if group_idx > 0 {
663                    // Normally, start at the boundary of the previous group.
664                    self.group_end_indices[group_idx - 1].1
665                } else {
666                    // If previous group is out of the table, start at zero.
667                    0
668                }
669            }
670            // Window frame end, PRECEDING n
671            (false, true) => {
672                if self.current_group_idx >= delta {
673                    let group_idx = self.current_group_idx - delta;
674                    self.group_end_indices[group_idx].1
675                } else {
676                    // Group is out of the table, therefore end at zero.
677                    0
678                }
679            }
680            // Window frame end, FOLLOWING n
681            (false, false) => {
682                let group_idx = std::cmp::min(
683                    self.current_group_idx + delta,
684                    self.group_end_indices.len() - 1,
685                );
686                self.group_end_indices[group_idx].1
687            }
688        })
689    }
690}
691
692fn check_equality(current: &[ScalarValue], target: &[ScalarValue]) -> Result<bool> {
693    Ok(current == target)
694}
695
696#[cfg(test)]
697mod tests {
698    use super::*;
699
700    use arrow::array::Float64Array;
701
702    fn get_test_data() -> (Vec<ArrayRef>, Vec<SortOptions>) {
703        let range_columns: Vec<ArrayRef> = vec![Arc::new(Float64Array::from(vec![
704            5.0, 7.0, 8.0, 8.0, 9., 10., 10., 10., 11.,
705        ]))];
706        let sort_options = vec![SortOptions {
707            descending: false,
708            nulls_first: false,
709        }];
710
711        (range_columns, sort_options)
712    }
713
714    fn assert_group_ranges(
715        window_frame: &Arc<WindowFrame>,
716        expected_results: Vec<(Range<usize>, usize)>,
717    ) -> Result<()> {
718        let mut window_frame_groups = WindowFrameStateGroups::default();
719        let (range_columns, _) = get_test_data();
720        let n_row = range_columns[0].len();
721        for (idx, (expected_range, expected_group_idx)) in
722            expected_results.into_iter().enumerate()
723        {
724            let range = window_frame_groups.calculate_range(
725                window_frame,
726                &range_columns,
727                n_row,
728                idx,
729            )?;
730            assert_eq!(range, expected_range);
731            assert_eq!(window_frame_groups.current_group_idx, expected_group_idx);
732        }
733        Ok(())
734    }
735
736    fn assert_frame_ranges(
737        window_frame: &Arc<WindowFrame>,
738        expected_results: Vec<Range<usize>>,
739    ) -> Result<()> {
740        let mut window_frame_context =
741            WindowFrameContext::new(Arc::clone(window_frame), vec![]);
742        let (range_columns, _) = get_test_data();
743        let n_row = range_columns[0].len();
744        let mut last_range = Range { start: 0, end: 0 };
745        for (idx, expected_range) in expected_results.into_iter().enumerate() {
746            let range = window_frame_context.calculate_range(
747                &range_columns,
748                &last_range,
749                n_row,
750                idx,
751            )?;
752            assert_eq!(range, expected_range);
753            last_range = range;
754        }
755        Ok(())
756    }
757
758    #[test]
759    fn test_default_window_frame_group_boundaries() -> Result<()> {
760        let window_frame = Arc::new(WindowFrame::new(None));
761        assert_group_ranges(
762            &window_frame,
763            vec![
764                (Range { start: 0, end: 9 }, 0),
765                (Range { start: 0, end: 9 }, 0),
766                (Range { start: 0, end: 9 }, 0),
767                (Range { start: 0, end: 9 }, 0),
768                (Range { start: 0, end: 9 }, 0),
769                (Range { start: 0, end: 9 }, 0),
770                (Range { start: 0, end: 9 }, 0),
771                (Range { start: 0, end: 9 }, 0),
772                (Range { start: 0, end: 9 }, 0),
773            ],
774        )?;
775
776        assert_frame_ranges(
777            &window_frame,
778            vec![
779                Range { start: 0, end: 9 },
780                Range { start: 0, end: 9 },
781                Range { start: 0, end: 9 },
782                Range { start: 0, end: 9 },
783                Range { start: 0, end: 9 },
784                Range { start: 0, end: 9 },
785                Range { start: 0, end: 9 },
786                Range { start: 0, end: 9 },
787                Range { start: 0, end: 9 },
788            ],
789        )?;
790
791        Ok(())
792    }
793
794    #[test]
795    fn test_unordered_window_frame_group_boundaries() -> Result<()> {
796        let window_frame = Arc::new(WindowFrame::new(Some(false)));
797        assert_group_ranges(
798            &window_frame,
799            vec![
800                (Range { start: 0, end: 1 }, 0),
801                (Range { start: 0, end: 2 }, 1),
802                (Range { start: 0, end: 4 }, 2),
803                (Range { start: 0, end: 4 }, 2),
804                (Range { start: 0, end: 5 }, 3),
805                (Range { start: 0, end: 8 }, 4),
806                (Range { start: 0, end: 8 }, 4),
807                (Range { start: 0, end: 8 }, 4),
808                (Range { start: 0, end: 9 }, 5),
809            ],
810        )?;
811
812        assert_frame_ranges(
813            &window_frame,
814            vec![
815                Range { start: 0, end: 9 },
816                Range { start: 0, end: 9 },
817                Range { start: 0, end: 9 },
818                Range { start: 0, end: 9 },
819                Range { start: 0, end: 9 },
820                Range { start: 0, end: 9 },
821                Range { start: 0, end: 9 },
822                Range { start: 0, end: 9 },
823                Range { start: 0, end: 9 },
824            ],
825        )?;
826
827        Ok(())
828    }
829
830    #[test]
831    fn test_ordered_window_frame_group_boundaries() -> Result<()> {
832        let window_frame = Arc::new(WindowFrame::new(Some(true)));
833        assert_group_ranges(
834            &window_frame,
835            vec![
836                (Range { start: 0, end: 1 }, 0),
837                (Range { start: 0, end: 2 }, 1),
838                (Range { start: 0, end: 4 }, 2),
839                (Range { start: 0, end: 4 }, 2),
840                (Range { start: 0, end: 5 }, 3),
841                (Range { start: 0, end: 8 }, 4),
842                (Range { start: 0, end: 8 }, 4),
843                (Range { start: 0, end: 8 }, 4),
844                (Range { start: 0, end: 9 }, 5),
845            ],
846        )?;
847
848        assert_frame_ranges(
849            &window_frame,
850            vec![
851                Range { start: 0, end: 1 },
852                Range { start: 0, end: 2 },
853                Range { start: 0, end: 3 },
854                Range { start: 0, end: 4 },
855                Range { start: 0, end: 5 },
856                Range { start: 0, end: 6 },
857                Range { start: 0, end: 7 },
858                Range { start: 0, end: 8 },
859                Range { start: 0, end: 9 },
860            ],
861        )?;
862
863        Ok(())
864    }
865
866    #[test]
867    fn test_window_frame_group_boundaries() -> Result<()> {
868        let window_frame = Arc::new(WindowFrame::new_bounds(
869            WindowFrameUnits::Groups,
870            WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))),
871            WindowFrameBound::Following(ScalarValue::UInt64(Some(1))),
872        ));
873        assert_group_ranges(
874            &window_frame,
875            vec![
876                (Range { start: 0, end: 2 }, 0),
877                (Range { start: 0, end: 4 }, 1),
878                (Range { start: 1, end: 5 }, 2),
879                (Range { start: 1, end: 5 }, 2),
880                (Range { start: 2, end: 8 }, 3),
881                (Range { start: 4, end: 9 }, 4),
882                (Range { start: 4, end: 9 }, 4),
883                (Range { start: 4, end: 9 }, 4),
884                (Range { start: 5, end: 9 }, 5),
885            ],
886        )
887    }
888
889    #[test]
890    fn test_window_frame_group_boundaries_both_following() -> Result<()> {
891        let window_frame = Arc::new(WindowFrame::new_bounds(
892            WindowFrameUnits::Groups,
893            WindowFrameBound::Following(ScalarValue::UInt64(Some(1))),
894            WindowFrameBound::Following(ScalarValue::UInt64(Some(2))),
895        ));
896        assert_group_ranges(
897            &window_frame,
898            vec![
899                (Range::<usize> { start: 1, end: 4 }, 0),
900                (Range::<usize> { start: 2, end: 5 }, 1),
901                (Range::<usize> { start: 4, end: 8 }, 2),
902                (Range::<usize> { start: 4, end: 8 }, 2),
903                (Range::<usize> { start: 5, end: 9 }, 3),
904                (Range::<usize> { start: 8, end: 9 }, 4),
905                (Range::<usize> { start: 8, end: 9 }, 4),
906                (Range::<usize> { start: 8, end: 9 }, 4),
907                (Range::<usize> { start: 9, end: 9 }, 5),
908            ],
909        )
910    }
911
912    #[test]
913    fn test_window_frame_group_boundaries_both_preceding() -> Result<()> {
914        let window_frame = Arc::new(WindowFrame::new_bounds(
915            WindowFrameUnits::Groups,
916            WindowFrameBound::Preceding(ScalarValue::UInt64(Some(2))),
917            WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))),
918        ));
919        assert_group_ranges(
920            &window_frame,
921            vec![
922                (Range::<usize> { start: 0, end: 0 }, 0),
923                (Range::<usize> { start: 0, end: 1 }, 1),
924                (Range::<usize> { start: 0, end: 2 }, 2),
925                (Range::<usize> { start: 0, end: 2 }, 2),
926                (Range::<usize> { start: 1, end: 4 }, 3),
927                (Range::<usize> { start: 2, end: 5 }, 4),
928                (Range::<usize> { start: 2, end: 5 }, 4),
929                (Range::<usize> { start: 2, end: 5 }, 4),
930                (Range::<usize> { start: 4, end: 8 }, 5),
931            ],
932        )
933    }
934}