Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/windows/mod.rs
Line
Count
Source (jump to first uncovered line)
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
//! Physical expressions for window functions
19
20
use std::borrow::Borrow;
21
use std::sync::Arc;
22
23
use crate::{
24
    expressions::{
25
        cume_dist, dense_rank, lag, lead, percent_rank, rank, Literal, NthValue, Ntile,
26
        PhysicalSortExpr,
27
    },
28
    ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PhysicalExpr,
29
};
30
31
use arrow::datatypes::Schema;
32
use arrow_schema::{DataType, Field, SchemaRef};
33
use datafusion_common::{
34
    exec_datafusion_err, exec_err, DataFusionError, Result, ScalarValue,
35
};
36
use datafusion_expr::{
37
    BuiltInWindowFunction, PartitionEvaluator, ReversedUDWF, WindowFrame,
38
    WindowFunctionDefinition, WindowUDF,
39
};
40
use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
41
use datafusion_physical_expr::equivalence::collapse_lex_req;
42
use datafusion_physical_expr::{
43
    reverse_order_bys,
44
    window::{BuiltInWindowFunctionExpr, SlidingAggregateWindowExpr},
45
    ConstExpr, EquivalenceProperties, LexOrdering, PhysicalSortRequirement,
46
};
47
use itertools::Itertools;
48
49
mod bounded_window_agg_exec;
50
mod utils;
51
mod window_agg_exec;
52
53
pub use bounded_window_agg_exec::BoundedWindowAggExec;
54
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
55
use datafusion_physical_expr::expressions::Column;
56
pub use datafusion_physical_expr::window::{
57
    BuiltInWindowExpr, PlainAggregateWindowExpr, WindowExpr,
58
};
59
use datafusion_physical_expr_common::sort_expr::LexRequirement;
60
pub use window_agg_exec::WindowAggExec;
61
62
/// Build field from window function and add it into schema
63
0
pub fn schema_add_window_field(
64
0
    args: &[Arc<dyn PhysicalExpr>],
65
0
    schema: &Schema,
66
0
    window_fn: &WindowFunctionDefinition,
67
0
    fn_name: &str,
68
0
) -> Result<Arc<Schema>> {
69
0
    let data_types = args
70
0
        .iter()
71
0
        .map(|e| Arc::clone(e).as_ref().data_type(schema))
72
0
        .collect::<Result<Vec<_>>>()?;
73
0
    let nullability = args
74
0
        .iter()
75
0
        .map(|e| Arc::clone(e).as_ref().nullable(schema))
76
0
        .collect::<Result<Vec<_>>>()?;
77
0
    let window_expr_return_type =
78
0
        window_fn.return_type(&data_types, &nullability, fn_name)?;
79
0
    let mut window_fields = schema
80
0
        .fields()
81
0
        .iter()
82
0
        .map(|f| f.as_ref().clone())
83
0
        .collect_vec();
84
0
    // Skip extending schema for UDAF
85
0
    if let WindowFunctionDefinition::AggregateUDF(_) = window_fn {
86
0
        Ok(Arc::new(Schema::new(window_fields)))
87
    } else {
88
0
        window_fields.extend_from_slice(&[Field::new(
89
0
            fn_name,
90
0
            window_expr_return_type,
91
0
            false,
92
0
        )]);
93
0
        Ok(Arc::new(Schema::new(window_fields)))
94
    }
95
0
}
96
97
/// Create a physical expression for window function
98
#[allow(clippy::too_many_arguments)]
99
2
pub fn create_window_expr(
100
2
    fun: &WindowFunctionDefinition,
101
2
    name: String,
102
2
    args: &[Arc<dyn PhysicalExpr>],
103
2
    partition_by: &[Arc<dyn PhysicalExpr>],
104
2
    order_by: &[PhysicalSortExpr],
105
2
    window_frame: Arc<WindowFrame>,
106
2
    input_schema: &Schema,
107
2
    ignore_nulls: bool,
108
2
) -> Result<Arc<dyn WindowExpr>> {
109
2
    Ok(match fun {
110
0
        WindowFunctionDefinition::BuiltInWindowFunction(fun) => {
111
0
            Arc::new(BuiltInWindowExpr::new(
112
0
                create_built_in_window_expr(fun, args, input_schema, name, ignore_nulls)?,
113
0
                partition_by,
114
0
                order_by,
115
0
                window_frame,
116
            ))
117
        }
118
2
        WindowFunctionDefinition::AggregateUDF(fun) => {
119
2
            let aggregate = AggregateExprBuilder::new(Arc::clone(fun), args.to_vec())
120
2
                .schema(Arc::new(input_schema.clone()))
121
2
                .alias(name)
122
2
                .with_ignore_nulls(ignore_nulls)
123
2
                .build()
?0
;
124
2
            window_expr_from_aggregate_expr(
125
2
                partition_by,
126
2
                order_by,
127
2
                window_frame,
128
2
                aggregate,
129
2
            )
130
        }
131
        // TODO: Ordering not supported for Window UDFs yet
132
0
        WindowFunctionDefinition::WindowUDF(fun) => Arc::new(BuiltInWindowExpr::new(
133
0
            create_udwf_window_expr(fun, args, input_schema, name, ignore_nulls)?,
134
0
            partition_by,
135
0
            order_by,
136
0
            window_frame,
137
        )),
138
    })
139
2
}
140
141
/// Creates an appropriate [`WindowExpr`] based on the window frame and
142
2
fn window_expr_from_aggregate_expr(
143
2
    partition_by: &[Arc<dyn PhysicalExpr>],
144
2
    order_by: &[PhysicalSortExpr],
145
2
    window_frame: Arc<WindowFrame>,
146
2
    aggregate: AggregateFunctionExpr,
147
2
) -> Arc<dyn WindowExpr> {
148
2
    // Is there a potentially unlimited sized window frame?
149
2
    let unbounded_window = window_frame.start_bound.is_unbounded();
150
2
151
2
    if !unbounded_window {
152
1
        Arc::new(SlidingAggregateWindowExpr::new(
153
1
            aggregate,
154
1
            partition_by,
155
1
            order_by,
156
1
            window_frame,
157
1
        ))
158
    } else {
159
1
        Arc::new(PlainAggregateWindowExpr::new(
160
1
            aggregate,
161
1
            partition_by,
162
1
            order_by,
163
1
            window_frame,
164
1
        ))
165
    }
166
2
}
167
168
0
fn get_scalar_value_from_args(
169
0
    args: &[Arc<dyn PhysicalExpr>],
170
0
    index: usize,
171
0
) -> Result<Option<ScalarValue>> {
172
0
    Ok(if let Some(field) = args.get(index) {
173
0
        let tmp = field
174
0
            .as_any()
175
0
            .downcast_ref::<Literal>()
176
0
            .ok_or_else(|| DataFusionError::NotImplemented(
177
0
                format!("There is only support Literal types for field at idx: {index} in Window Function"),
178
0
            ))?
179
0
            .value()
180
0
            .clone();
181
0
        Some(tmp)
182
    } else {
183
0
        None
184
    })
185
0
}
186
187
0
fn get_signed_integer(value: ScalarValue) -> Result<i64> {
188
0
    if value.is_null() {
189
0
        return Ok(0);
190
0
    }
191
0
192
0
    if !value.data_type().is_integer() {
193
0
        return exec_err!("Expected an integer value");
194
0
    }
195
0
196
0
    value.cast_to(&DataType::Int64)?.try_into()
197
0
}
198
199
0
fn get_unsigned_integer(value: ScalarValue) -> Result<u64> {
200
0
    if value.is_null() {
201
0
        return Ok(0);
202
0
    }
203
0
204
0
    if !value.data_type().is_integer() {
205
0
        return exec_err!("Expected an integer value");
206
0
    }
207
0
208
0
    value.cast_to(&DataType::UInt64)?.try_into()
209
0
}
210
211
0
fn get_casted_value(
212
0
    default_value: Option<ScalarValue>,
213
0
    dtype: &DataType,
214
0
) -> Result<ScalarValue> {
215
0
    match default_value {
216
0
        Some(v) if !v.data_type().is_null() => v.cast_to(dtype),
217
        // If None or Null datatype
218
0
        _ => ScalarValue::try_from(dtype),
219
    }
220
0
}
221
222
0
fn create_built_in_window_expr(
223
0
    fun: &BuiltInWindowFunction,
224
0
    args: &[Arc<dyn PhysicalExpr>],
225
0
    input_schema: &Schema,
226
0
    name: String,
227
0
    ignore_nulls: bool,
228
0
) -> Result<Arc<dyn BuiltInWindowFunctionExpr>> {
229
    // derive the output datatype from incoming schema
230
0
    let out_data_type: &DataType = input_schema.field_with_name(&name)?.data_type();
231
0
232
0
    Ok(match fun {
233
0
        BuiltInWindowFunction::Rank => Arc::new(rank(name, out_data_type)),
234
0
        BuiltInWindowFunction::DenseRank => Arc::new(dense_rank(name, out_data_type)),
235
0
        BuiltInWindowFunction::PercentRank => Arc::new(percent_rank(name, out_data_type)),
236
0
        BuiltInWindowFunction::CumeDist => Arc::new(cume_dist(name, out_data_type)),
237
        BuiltInWindowFunction::Ntile => {
238
0
            let n = get_scalar_value_from_args(args, 0)?.ok_or_else(|| {
239
0
                DataFusionError::Execution(
240
0
                    "NTILE requires a positive integer".to_string(),
241
0
                )
242
0
            })?;
243
244
0
            if n.is_null() {
245
0
                return exec_err!("NTILE requires a positive integer, but finds NULL");
246
0
            }
247
0
248
0
            if n.is_unsigned() {
249
0
                let n = get_unsigned_integer(n)?;
250
0
                Arc::new(Ntile::new(name, n, out_data_type))
251
            } else {
252
0
                let n: i64 = get_signed_integer(n)?;
253
0
                if n <= 0 {
254
0
                    return exec_err!("NTILE requires a positive integer");
255
0
                }
256
0
                Arc::new(Ntile::new(name, n as u64, out_data_type))
257
            }
258
        }
259
        BuiltInWindowFunction::Lag => {
260
0
            let arg = Arc::clone(&args[0]);
261
0
            let shift_offset = get_scalar_value_from_args(args, 1)?
262
0
                .map(get_signed_integer)
263
0
                .map_or(Ok(None), |v| v.map(Some))?;
264
0
            let default_value =
265
0
                get_casted_value(get_scalar_value_from_args(args, 2)?, out_data_type)?;
266
0
            Arc::new(lag(
267
0
                name,
268
0
                out_data_type.clone(),
269
0
                arg,
270
0
                shift_offset,
271
0
                default_value,
272
0
                ignore_nulls,
273
0
            ))
274
        }
275
        BuiltInWindowFunction::Lead => {
276
0
            let arg = Arc::clone(&args[0]);
277
0
            let shift_offset = get_scalar_value_from_args(args, 1)?
278
0
                .map(get_signed_integer)
279
0
                .map_or(Ok(None), |v| v.map(Some))?;
280
0
            let default_value =
281
0
                get_casted_value(get_scalar_value_from_args(args, 2)?, out_data_type)?;
282
0
            Arc::new(lead(
283
0
                name,
284
0
                out_data_type.clone(),
285
0
                arg,
286
0
                shift_offset,
287
0
                default_value,
288
0
                ignore_nulls,
289
0
            ))
290
        }
291
        BuiltInWindowFunction::NthValue => {
292
0
            let arg = Arc::clone(&args[0]);
293
0
            let n = get_signed_integer(
294
0
                args[1]
295
0
                    .as_any()
296
0
                    .downcast_ref::<Literal>()
297
0
                    .ok_or_else(|| {
298
0
                        exec_datafusion_err!("Expected a signed integer literal for the second argument of nth_value, got {}", args[1])
299
0
                    })?
300
0
                    .value()
301
0
                    .clone(),
302
0
            )?;
303
0
            Arc::new(NthValue::nth(
304
0
                name,
305
0
                arg,
306
0
                out_data_type.clone(),
307
0
                n,
308
0
                ignore_nulls,
309
0
            )?)
310
        }
311
        BuiltInWindowFunction::FirstValue => {
312
0
            let arg = Arc::clone(&args[0]);
313
0
            Arc::new(NthValue::first(
314
0
                name,
315
0
                arg,
316
0
                out_data_type.clone(),
317
0
                ignore_nulls,
318
0
            ))
319
        }
320
        BuiltInWindowFunction::LastValue => {
321
0
            let arg = Arc::clone(&args[0]);
322
0
            Arc::new(NthValue::last(
323
0
                name,
324
0
                arg,
325
0
                out_data_type.clone(),
326
0
                ignore_nulls,
327
0
            ))
328
        }
329
    })
330
0
}
331
332
/// Creates a `BuiltInWindowFunctionExpr` suitable for a user defined window function
333
0
fn create_udwf_window_expr(
334
0
    fun: &Arc<WindowUDF>,
335
0
    args: &[Arc<dyn PhysicalExpr>],
336
0
    input_schema: &Schema,
337
0
    name: String,
338
0
    ignore_nulls: bool,
339
0
) -> Result<Arc<dyn BuiltInWindowFunctionExpr>> {
340
    // need to get the types into an owned vec for some reason
341
0
    let input_types: Vec<_> = args
342
0
        .iter()
343
0
        .map(|arg| arg.data_type(input_schema))
344
0
        .collect::<Result<_>>()?;
345
346
0
    Ok(Arc::new(WindowUDFExpr {
347
0
        fun: Arc::clone(fun),
348
0
        args: args.to_vec(),
349
0
        input_types,
350
0
        name,
351
0
        is_reversed: false,
352
0
        ignore_nulls,
353
0
    }))
354
0
}
355
356
/// Implements [`BuiltInWindowFunctionExpr`] for [`WindowUDF`]
357
#[derive(Clone, Debug)]
358
struct WindowUDFExpr {
359
    fun: Arc<WindowUDF>,
360
    args: Vec<Arc<dyn PhysicalExpr>>,
361
    /// Display name
362
    name: String,
363
    /// Types of input expressions
364
    input_types: Vec<DataType>,
365
    /// This is set to `true` only if the user-defined window function
366
    /// expression supports evaluation in reverse order, and the
367
    /// evaluation order is reversed.
368
    is_reversed: bool,
369
    /// Set to `true` if `IGNORE NULLS` is defined, `false` otherwise.
370
    ignore_nulls: bool,
371
}
372
373
impl BuiltInWindowFunctionExpr for WindowUDFExpr {
374
0
    fn as_any(&self) -> &dyn std::any::Any {
375
0
        self
376
0
    }
377
378
0
    fn field(&self) -> Result<Field> {
379
0
        self.fun
380
0
            .field(WindowUDFFieldArgs::new(&self.input_types, &self.name))
381
0
    }
382
383
0
    fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
384
0
        self.args.clone()
385
0
    }
386
387
0
    fn create_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
388
0
        self.fun.partition_evaluator_factory()
389
0
    }
390
391
0
    fn name(&self) -> &str {
392
0
        &self.name
393
0
    }
394
395
0
    fn reverse_expr(&self) -> Option<Arc<dyn BuiltInWindowFunctionExpr>> {
396
0
        match self.fun.reverse_expr() {
397
0
            ReversedUDWF::Identical => Some(Arc::new(self.clone())),
398
0
            ReversedUDWF::NotSupported => None,
399
0
            ReversedUDWF::Reversed(fun) => Some(Arc::new(WindowUDFExpr {
400
0
                fun,
401
0
                args: self.args.clone(),
402
0
                name: self.name.clone(),
403
0
                input_types: self.input_types.clone(),
404
0
                is_reversed: !self.is_reversed,
405
0
                ignore_nulls: self.ignore_nulls,
406
0
            })),
407
        }
408
0
    }
409
410
0
    fn get_result_ordering(&self, schema: &SchemaRef) -> Option<PhysicalSortExpr> {
411
0
        self.fun
412
0
            .sort_options()
413
0
            .zip(schema.column_with_name(self.name()))
414
0
            .map(|(options, (idx, field))| {
415
0
                let expr = Arc::new(Column::new(field.name(), idx));
416
0
                PhysicalSortExpr { expr, options }
417
0
            })
418
0
    }
419
}
420
421
4
pub(crate) fn calc_requirements<
422
4
    T: Borrow<Arc<dyn PhysicalExpr>>,
423
4
    S: Borrow<PhysicalSortExpr>,
424
4
>(
425
4
    partition_by_exprs: impl IntoIterator<Item = T>,
426
4
    orderby_sort_exprs: impl IntoIterator<Item = S>,
427
4
) -> Option<LexRequirement> {
428
4
    let mut sort_reqs = LexRequirement::new(
429
4
        partition_by_exprs
430
4
            .into_iter()
431
5
            .map(|partition_by| {
432
5
                PhysicalSortRequirement::new(Arc::clone(partition_by.borrow()), None)
433
5
            })
434
4
            .collect::<Vec<_>>(),
435
4
    );
436
6
    for element in 
orderby_sort_exprs.into_iter()4
{
437
6
        let PhysicalSortExpr { expr, options } = element.borrow();
438
9
        if !sort_reqs.iter().any(|e| e.expr.eq(expr)
)6
{
439
4
            sort_reqs.push(PhysicalSortRequirement::new(
440
4
                Arc::clone(expr),
441
4
                Some(*options),
442
4
            ));
443
4
        }
2
444
    }
445
    // Convert empty result to None. Otherwise wrap result inside Some()
446
4
    (!sort_reqs.is_empty()).then_some(sort_reqs)
447
4
}
448
449
/// This function calculates the indices such that when partition by expressions reordered with the indices
450
/// resulting expressions define a preset for existing ordering.
451
/// For instance, if input is ordered by a, b, c and PARTITION BY b, a is used,
452
/// this vector will be [1, 0]. It means that when we iterate b, a columns with the order [1, 0]
453
/// resulting vector (a, b) is a preset of the existing ordering (a, b, c).
454
52
pub fn get_ordered_partition_by_indices(
455
52
    partition_by_exprs: &[Arc<dyn PhysicalExpr>],
456
52
    input: &Arc<dyn ExecutionPlan>,
457
52
) -> Vec<usize> {
458
52
    let (_, indices) = input
459
52
        .equivalence_properties()
460
52
        .find_longest_permutation(partition_by_exprs);
461
52
    indices
462
52
}
463
464
3
pub(crate) fn get_partition_by_sort_exprs(
465
3
    input: &Arc<dyn ExecutionPlan>,
466
3
    partition_by_exprs: &[Arc<dyn PhysicalExpr>],
467
3
    ordered_partition_by_indices: &[usize],
468
3
) -> Result<LexOrdering> {
469
3
    let ordered_partition_exprs = ordered_partition_by_indices
470
3
        .iter()
471
3
        .map(|idx| 
Arc::clone(&partition_by_exprs[*idx])0
)
472
3
        .collect::<Vec<_>>();
473
3
    // Make sure ordered section doesn't move over the partition by expression
474
3
    assert!(ordered_partition_by_indices.len() <= partition_by_exprs.len());
475
3
    let (ordering, _) = input
476
3
        .equivalence_properties()
477
3
        .find_longest_permutation(&ordered_partition_exprs);
478
3
    if ordering.len() == ordered_partition_exprs.len() {
479
3
        Ok(ordering)
480
    } else {
481
0
        exec_err!("Expects PARTITION BY expression to be ordered")
482
    }
483
3
}
484
485
3
pub(crate) fn window_equivalence_properties(
486
3
    schema: &SchemaRef,
487
3
    input: &Arc<dyn ExecutionPlan>,
488
3
    window_expr: &[Arc<dyn WindowExpr>],
489
3
) -> EquivalenceProperties {
490
3
    // We need to update the schema, so we can not directly use
491
3
    // `input.equivalence_properties()`.
492
3
    let mut window_eq_properties = EquivalenceProperties::new(Arc::clone(schema))
493
3
        .extend(input.equivalence_properties().clone());
494
495
8
    for 
expr5
in window_expr {
496
3
        if let Some(builtin_window_expr) =
497
5
            expr.as_any().downcast_ref::<BuiltInWindowExpr>()
498
3
        {
499
3
            builtin_window_expr.add_equal_orderings(&mut window_eq_properties);
500
3
        }
2
501
    }
502
3
    window_eq_properties
503
3
}
504
505
/// Constructs the best-fitting windowing operator (a `WindowAggExec` or a
506
/// `BoundedWindowExec`) for the given `input` according to the specifications
507
/// of `window_exprs` and `physical_partition_keys`. Here, best-fitting means
508
/// not requiring additional sorting and/or partitioning for the given input.
509
/// - A return value of `None` represents that there is no way to construct a
510
///   windowing operator that doesn't need additional sorting/partitioning for
511
///   the given input. Existing ordering should be changed to run the given
512
///   windowing operation.
513
/// - A `Some(window exec)` value contains the optimal windowing operator (a
514
///   `WindowAggExec` or a `BoundedWindowExec`) for the given input.
515
0
pub fn get_best_fitting_window(
516
0
    window_exprs: &[Arc<dyn WindowExpr>],
517
0
    input: &Arc<dyn ExecutionPlan>,
518
0
    // These are the partition keys used during repartitioning.
519
0
    // They are either the same with `window_expr`'s PARTITION BY columns,
520
0
    // or it is empty if partitioning is not desirable for this windowing operator.
521
0
    physical_partition_keys: &[Arc<dyn PhysicalExpr>],
522
0
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
523
0
    // Contains at least one window expr and all of the partition by and order by sections
524
0
    // of the window_exprs are same.
525
0
    let partitionby_exprs = window_exprs[0].partition_by();
526
0
    let orderby_keys = window_exprs[0].order_by();
527
0
    let (should_reverse, input_order_mode) =
528
0
        if let Some((should_reverse, input_order_mode)) =
529
0
            get_window_mode(partitionby_exprs, orderby_keys, input)
530
        {
531
0
            (should_reverse, input_order_mode)
532
        } else {
533
0
            return Ok(None);
534
        };
535
0
    let is_unbounded = input.execution_mode().is_unbounded();
536
0
    if !is_unbounded && input_order_mode != InputOrderMode::Sorted {
537
        // Executor has bounded input and `input_order_mode` is not `InputOrderMode::Sorted`
538
        // in this case removing the sort is not helpful, return:
539
0
        return Ok(None);
540
0
    };
541
542
0
    let window_expr = if should_reverse {
543
0
        if let Some(reversed_window_expr) = window_exprs
544
0
            .iter()
545
0
            .map(|e| e.get_reverse_expr())
546
0
            .collect::<Option<Vec<_>>>()
547
        {
548
0
            reversed_window_expr
549
        } else {
550
            // Cannot take reverse of any of the window expr
551
            // In this case, with existing ordering window cannot be run
552
0
            return Ok(None);
553
        }
554
    } else {
555
0
        window_exprs.to_vec()
556
    };
557
558
    // If all window expressions can run with bounded memory, choose the
559
    // bounded window variant:
560
0
    if window_expr.iter().all(|e| e.uses_bounded_memory()) {
561
0
        Ok(Some(Arc::new(BoundedWindowAggExec::try_new(
562
0
            window_expr,
563
0
            Arc::clone(input),
564
0
            physical_partition_keys.to_vec(),
565
0
            input_order_mode,
566
0
        )?) as _))
567
0
    } else if input_order_mode != InputOrderMode::Sorted {
568
        // For `WindowAggExec` to work correctly PARTITION BY columns should be sorted.
569
        // Hence, if `input_order_mode` is not `Sorted` we should convert
570
        // input ordering such that it can work with `Sorted` (add `SortExec`).
571
        // Effectively `WindowAggExec` works only in `Sorted` mode.
572
0
        Ok(None)
573
    } else {
574
0
        Ok(Some(Arc::new(WindowAggExec::try_new(
575
0
            window_expr,
576
0
            Arc::clone(input),
577
0
            physical_partition_keys.to_vec(),
578
0
        )?) as _))
579
    }
580
0
}
581
582
/// Compares physical ordering (output ordering of the `input` operator) with
583
/// `partitionby_exprs` and `orderby_keys` to decide whether existing ordering
584
/// is sufficient to run the current window operator.
585
/// - A `None` return value indicates that we can not remove the sort in question
586
///   (input ordering is not sufficient to run current window executor).
587
/// - A `Some((bool, InputOrderMode))` value indicates that the window operator
588
///   can run with existing input ordering, so we can remove `SortExec` before it.
589
///
590
/// The `bool` field in the return value represents whether we should reverse window
591
/// operator to remove `SortExec` before it. The `InputOrderMode` field represents
592
/// the mode this window operator should work in to accommodate the existing ordering.
593
75
pub fn get_window_mode(
594
75
    partitionby_exprs: &[Arc<dyn PhysicalExpr>],
595
75
    orderby_keys: &[PhysicalSortExpr],
596
75
    input: &Arc<dyn ExecutionPlan>,
597
75
) -> Option<(bool, InputOrderMode)> {
598
75
    let input_eqs = input.equivalence_properties().clone();
599
75
    let mut partition_by_reqs: LexRequirement = LexRequirement::new(vec![]);
600
75
    let (_, indices) = input_eqs.find_longest_permutation(partitionby_exprs);
601
76
    vec![].extend(indices.iter().map(|&idx| PhysicalSortRequirement {
602
76
        expr: Arc::clone(&partitionby_exprs[idx]),
603
76
        options: None,
604
76
    }));
605
75
    partition_by_reqs
606
75
        .inner
607
76
        .extend(indices.iter().map(|&idx| PhysicalSortRequirement {
608
76
            expr: Arc::clone(&partitionby_exprs[idx]),
609
76
            options: None,
610
76
        }));
611
75
    // Treat partition by exprs as constant. During analysis of requirements are satisfied.
612
75
    let const_exprs = partitionby_exprs.iter().map(ConstExpr::from);
613
75
    let partition_by_eqs = input_eqs.with_constants(const_exprs);
614
75
    let order_by_reqs = PhysicalSortRequirement::from_sort_exprs(orderby_keys);
615
75
    let reverse_order_by_reqs =
616
75
        PhysicalSortRequirement::from_sort_exprs(&reverse_order_bys(orderby_keys));
617
89
    for (should_swap, order_by_reqs) in
618
75
        [(false, order_by_reqs), (true, reverse_order_by_reqs)]
619
    {
620
89
        let req = LexRequirement::new(
621
89
            [partition_by_reqs.inner.clone(), order_by_reqs.inner].concat(),
622
89
        );
623
89
        let req = collapse_lex_req(req);
624
89
        if partition_by_eqs.ordering_satisfy_requirement(&req) {
625
            // Window can be run with existing ordering
626
63
            let mode = if indices.len() == partitionby_exprs.len() {
627
30
                InputOrderMode::Sorted
628
33
            } else if indices.is_empty() {
629
20
                InputOrderMode::Linear
630
            } else {
631
13
                InputOrderMode::PartiallySorted(indices)
632
            };
633
63
            return Some((should_swap, mode));
634
26
        }
635
    }
636
12
    None
637
75
}
638
639
#[cfg(test)]
640
mod tests {
641
    use super::*;
642
    use crate::collect;
643
    use crate::expressions::col;
644
    use crate::streaming::StreamingTableExec;
645
    use crate::test::assert_is_pending;
646
    use crate::test::exec::{assert_strong_count_converges_to_zero, BlockingExec};
647
648
    use arrow::compute::SortOptions;
649
    use datafusion_execution::TaskContext;
650
651
    use datafusion_functions_aggregate::count::count_udaf;
652
    use futures::FutureExt;
653
    use InputOrderMode::{Linear, PartiallySorted, Sorted};
654
655
2
    fn create_test_schema() -> Result<SchemaRef> {
656
2
        let nullable_column = Field::new("nullable_col", DataType::Int32, true);
657
2
        let non_nullable_column = Field::new("non_nullable_col", DataType::Int32, false);
658
2
        let schema = Arc::new(Schema::new(vec![nullable_column, non_nullable_column]));
659
2
660
2
        Ok(schema)
661
2
    }
662
663
1
    fn create_test_schema2() -> Result<SchemaRef> {
664
1
        let a = Field::new("a", DataType::Int32, true);
665
1
        let b = Field::new("b", DataType::Int32, true);
666
1
        let c = Field::new("c", DataType::Int32, true);
667
1
        let d = Field::new("d", DataType::Int32, true);
668
1
        let e = Field::new("e", DataType::Int32, true);
669
1
        let schema = Arc::new(Schema::new(vec![a, b, c, d, e]));
670
1
        Ok(schema)
671
1
    }
672
673
    // Generate a schema which consists of 5 columns (a, b, c, d, e)
674
2
    fn create_test_schema3() -> Result<SchemaRef> {
675
2
        let a = Field::new("a", DataType::Int32, true);
676
2
        let b = Field::new("b", DataType::Int32, false);
677
2
        let c = Field::new("c", DataType::Int32, true);
678
2
        let d = Field::new("d", DataType::Int32, false);
679
2
        let e = Field::new("e", DataType::Int32, false);
680
2
        let schema = Arc::new(Schema::new(vec![a, b, c, d, e]));
681
2
        Ok(schema)
682
2
    }
683
684
    /// make PhysicalSortExpr with default options
685
8
    pub fn sort_expr(name: &str, schema: &Schema) -> PhysicalSortExpr {
686
8
        sort_expr_options(name, schema, SortOptions::default())
687
8
    }
688
689
    /// PhysicalSortExpr with specified options
690
8
    pub fn sort_expr_options(
691
8
        name: &str,
692
8
        schema: &Schema,
693
8
        options: SortOptions,
694
8
    ) -> PhysicalSortExpr {
695
8
        PhysicalSortExpr {
696
8
            expr: col(name, schema).unwrap(),
697
8
            options,
698
8
        }
699
8
    }
700
701
    /// Created a sorted Streaming Table exec
702
2
    pub fn streaming_table_exec(
703
2
        schema: &SchemaRef,
704
2
        sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
705
2
        infinite_source: bool,
706
2
    ) -> Result<Arc<dyn ExecutionPlan>> {
707
2
        let sort_exprs = sort_exprs.into_iter().collect();
708
2
709
2
        Ok(Arc::new(StreamingTableExec::try_new(
710
2
            Arc::clone(schema),
711
2
            vec![],
712
2
            None,
713
2
            Some(sort_exprs),
714
2
            infinite_source,
715
2
            None,
716
2
        )
?0
))
717
2
    }
718
719
    #[tokio::test]
720
1
    async fn test_calc_requirements() -> Result<()> {
721
1
        let schema = create_test_schema2()
?0
;
722
1
        let test_data = vec![
723
1
            // PARTITION BY a, ORDER BY b ASC NULLS FIRST
724
1
            (
725
1
                vec!["a"],
726
1
                vec![("b", true, true)],
727
1
                vec![("a", None), ("b", Some((true, true)))],
728
1
            ),
729
1
            // PARTITION BY a, ORDER BY a ASC NULLS FIRST
730
1
            (vec!["a"], vec![("a", true, true)], vec![("a", None)]),
731
1
            // PARTITION BY a, ORDER BY b ASC NULLS FIRST, c DESC NULLS LAST
732
1
            (
733
1
                vec!["a"],
734
1
                vec![("b", true, true), ("c", false, false)],
735
1
                vec![
736
1
                    ("a", None),
737
1
                    ("b", Some((true, true))),
738
1
                    ("c", Some((false, false))),
739
1
                ],
740
1
            ),
741
1
            // PARTITION BY a, c, ORDER BY b ASC NULLS FIRST, c DESC NULLS LAST
742
1
            (
743
1
                vec!["a", "c"],
744
1
                vec![("b", true, true), ("c", false, false)],
745
1
                vec![("a", None), ("c", None), ("b", Some((true, true)))],
746
1
            ),
747
1
        ];
748
5
        for (
pb_params, ob_params, expected_params4
) in test_data {
749
4
            let mut partitionbys = vec![];
750
9
            for 
col_name5
in pb_params {
751
5
                partitionbys.push(col(col_name, &schema)
?0
);
752
1
            }
753
1
754
4
            let mut orderbys = vec![];
755
10
            for (
col_name, descending, nulls_first6
) in ob_params {
756
6
                let expr = col(col_name, &schema)
?0
;
757
6
                let options = SortOptions {
758
6
                    descending,
759
6
                    nulls_first,
760
6
                };
761
6
                orderbys.push(PhysicalSortExpr { expr, options });
762
1
            }
763
1
764
4
            let mut expected: Option<LexRequirement> = None;
765
13
            for (
col_name, reqs9
) in expected_params {
766
9
                let options = reqs.map(|(descending, nulls_first)| SortOptions {
767
4
                    descending,
768
4
                    nulls_first,
769
9
                });
770
9
                let expr = col(col_name, &schema)
?0
;
771
9
                let res = PhysicalSortRequirement::new(expr, options);
772
9
                if let Some(
expected5
) = &mut expected {
773
5
                    expected.push(res);
774
5
                } else {
775
4
                    expected = Some(LexRequirement::new(vec![res]));
776
4
                }
777
1
            }
778
4
            assert_eq!(calc_requirements(partitionbys, orderbys), expected);
779
1
        }
780
1
        Ok(())
781
1
    }
782
783
    #[tokio::test]
784
1
    async fn test_drop_cancel() -> Result<()> {
785
1
        let task_ctx = Arc::new(TaskContext::default());
786
1
        let schema =
787
1
            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)]));
788
1
789
1
        let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1));
790
1
        let refs = blocking_exec.refs();
791
1
        let window_agg_exec = Arc::new(WindowAggExec::try_new(
792
1
            vec![create_window_expr(
793
1
                &WindowFunctionDefinition::AggregateUDF(count_udaf()),
794
1
                "count".to_owned(),
795
1
                &[col("a", &schema)
?0
],
796
1
                &[],
797
1
                &[],
798
1
                Arc::new(WindowFrame::new(None)),
799
1
                schema.as_ref(),
800
1
                false,
801
1
            )
?0
],
802
1
            blocking_exec,
803
1
            vec![],
804
1
        )
?0
);
805
1
806
1
        let fut = collect(window_agg_exec, task_ctx);
807
1
        let mut fut = fut.boxed();
808
1
809
1
        assert_is_pending(&mut fut);
810
1
        drop(fut);
811
1
        assert_strong_count_converges_to_zero(refs).
await0
;
812
1
813
1
        Ok(())
814
1
    }
815
816
    #[tokio::test]
817
1
    async fn test_satisfy_nullable() -> Result<()> {
818
1
        let schema = create_test_schema()
?0
;
819
1
        let params = vec![
820
1
            ((true, true), (false, false), false),
821
1
            ((true, true), (false, true), false),
822
1
            ((true, true), (true, false), false),
823
1
            ((true, false), (false, true), false),
824
1
            ((true, false), (false, false), false),
825
1
            ((true, false), (true, true), false),
826
1
            ((true, false), (true, false), true),
827
1
        ];
828
1
        for (
829
7
            (physical_desc, physical_nulls_first),
830
7
            (req_desc, req_nulls_first),
831
7
            expected,
832
8
        ) in params
833
1
        {
834
7
            let physical_ordering = PhysicalSortExpr {
835
7
                expr: col("nullable_col", &schema)
?0
,
836
7
                options: SortOptions {
837
7
                    descending: physical_desc,
838
7
                    nulls_first: physical_nulls_first,
839
7
                },
840
1
            };
841
7
            let required_ordering = PhysicalSortExpr {
842
7
                expr: col("nullable_col", &schema)
?0
,
843
7
                options: SortOptions {
844
7
                    descending: req_desc,
845
7
                    nulls_first: req_nulls_first,
846
7
                },
847
7
            };
848
7
            let res = physical_ordering.satisfy(&required_ordering.into(), &schema);
849
7
            assert_eq!(res, expected);
850
1
        }
851
1
852
1
        Ok(())
853
1
    }
854
855
    #[tokio::test]
856
1
    async fn test_satisfy_non_nullable() -> Result<()> {
857
1
        let schema = create_test_schema()
?0
;
858
1
859
1
        let params = vec![
860
1
            ((true, true), (false, false), false),
861
1
            ((true, true), (false, true), false),
862
1
            ((true, true), (true, false), true),
863
1
            ((true, false), (false, true), false),
864
1
            ((true, false), (false, false), false),
865
1
            ((true, false), (true, true), true),
866
1
            ((true, false), (true, false), true),
867
1
        ];
868
1
        for (
869
7
            (physical_desc, physical_nulls_first),
870
7
            (req_desc, req_nulls_first),
871
7
            expected,
872
8
        ) in params
873
1
        {
874
7
            let physical_ordering = PhysicalSortExpr {
875
7
                expr: col("non_nullable_col", &schema)
?0
,
876
7
                options: SortOptions {
877
7
                    descending: physical_desc,
878
7
                    nulls_first: physical_nulls_first,
879
7
                },
880
1
            };
881
7
            let required_ordering = PhysicalSortExpr {
882
7
                expr: col("non_nullable_col", &schema)
?0
,
883
7
                options: SortOptions {
884
7
                    descending: req_desc,
885
7
                    nulls_first: req_nulls_first,
886
7
                },
887
7
            };
888
7
            let res = physical_ordering.satisfy(&required_ordering.into(), &schema);
889
7
            assert_eq!(res, expected);
890
1
        }
891
1
892
1
        Ok(())
893
1
    }
894
895
    #[tokio::test]
896
1
    async fn test_get_window_mode_exhaustive() -> Result<()> {
897
1
        let test_schema = create_test_schema3()
?0
;
898
1
        // Columns a,c are nullable whereas b,d are not nullable.
899
1
        // Source is sorted by a ASC NULLS FIRST, b ASC NULLS FIRST, c ASC NULLS FIRST, d ASC NULLS FIRST
900
1
        // Column e is not ordered.
901
1
        let sort_exprs = vec![
902
1
            sort_expr("a", &test_schema),
903
1
            sort_expr("b", &test_schema),
904
1
            sort_expr("c", &test_schema),
905
1
            sort_expr("d", &test_schema),
906
1
        ];
907
1
        let exec_unbounded = streaming_table_exec(&test_schema, sort_exprs, true)
?0
;
908
1
909
1
        // test cases consists of vector of tuples. Where each tuple represents a single test case.
910
1
        // First field in the tuple is Vec<str> where each element in the vector represents PARTITION BY columns
911
1
        // For instance `vec!["a", "b"]` corresponds to PARTITION BY a, b
912
1
        // Second field in the tuple is Vec<str> where each element in the vector represents ORDER BY columns
913
1
        // For instance, vec!["c"], corresponds to ORDER BY c ASC NULLS FIRST, (ordering is default ordering. We do not check
914
1
        // for reversibility in this test).
915
1
        // Third field in the tuple is Option<InputOrderMode>, which corresponds to expected algorithm mode.
916
1
        // None represents that existing ordering is not sufficient to run executor with any one of the algorithms
917
1
        // (We need to add SortExec to be able to run it).
918
1
        // Some(InputOrderMode) represents, we can run algorithm with existing ordering; and algorithm should work in
919
1
        // InputOrderMode.
920
1
        let test_cases = vec![
921
1
            (vec!["a"], vec!["a"], Some(Sorted)),
922
1
            (vec!["a"], vec!["b"], Some(Sorted)),
923
1
            (vec!["a"], vec!["c"], None),
924
1
            (vec!["a"], vec!["a", "b"], Some(Sorted)),
925
1
            (vec!["a"], vec!["b", "c"], Some(Sorted)),
926
1
            (vec!["a"], vec!["a", "c"], None),
927
1
            (vec!["a"], vec!["a", "b", "c"], Some(Sorted)),
928
1
            (vec!["b"], vec!["a"], Some(Linear)),
929
1
            (vec!["b"], vec!["b"], Some(Linear)),
930
1
            (vec!["b"], vec!["c"], None),
931
1
            (vec!["b"], vec!["a", "b"], Some(Linear)),
932
1
            (vec!["b"], vec!["b", "c"], None),
933
1
            (vec!["b"], vec!["a", "c"], Some(Linear)),
934
1
            (vec!["b"], vec!["a", "b", "c"], Some(Linear)),
935
1
            (vec!["c"], vec!["a"], Some(Linear)),
936
1
            (vec!["c"], vec!["b"], None),
937
1
            (vec!["c"], vec!["c"], Some(Linear)),
938
1
            (vec!["c"], vec!["a", "b"], Some(Linear)),
939
1
            (vec!["c"], vec!["b", "c"], None),
940
1
            (vec!["c"], vec!["a", "c"], Some(Linear)),
941
1
            (vec!["c"], vec!["a", "b", "c"], Some(Linear)),
942
1
            (vec!["b", "a"], vec!["a"], Some(Sorted)),
943
1
            (vec!["b", "a"], vec!["b"], Some(Sorted)),
944
1
            (vec!["b", "a"], vec!["c"], Some(Sorted)),
945
1
            (vec!["b", "a"], vec!["a", "b"], Some(Sorted)),
946
1
            (vec!["b", "a"], vec!["b", "c"], Some(Sorted)),
947
1
            (vec!["b", "a"], vec!["a", "c"], Some(Sorted)),
948
1
            (vec!["b", "a"], vec!["a", "b", "c"], Some(Sorted)),
949
1
            (vec!["c", "b"], vec!["a"], Some(Linear)),
950
1
            (vec!["c", "b"], vec!["b"], Some(Linear)),
951
1
            (vec!["c", "b"], vec!["c"], Some(Linear)),
952
1
            (vec!["c", "b"], vec!["a", "b"], Some(Linear)),
953
1
            (vec!["c", "b"], vec!["b", "c"], Some(Linear)),
954
1
            (vec!["c", "b"], vec!["a", "c"], Some(Linear)),
955
1
            (vec!["c", "b"], vec!["a", "b", "c"], Some(Linear)),
956
1
            (vec!["c", "a"], vec!["a"], Some(PartiallySorted(vec![1]))),
957
1
            (vec!["c", "a"], vec!["b"], Some(PartiallySorted(vec![1]))),
958
1
            (vec!["c", "a"], vec!["c"], Some(PartiallySorted(vec![1]))),
959
1
            (
960
1
                vec!["c", "a"],
961
1
                vec!["a", "b"],
962
1
                Some(PartiallySorted(vec![1])),
963
1
            ),
964
1
            (
965
1
                vec!["c", "a"],
966
1
                vec!["b", "c"],
967
1
                Some(PartiallySorted(vec![1])),
968
1
            ),
969
1
            (
970
1
                vec!["c", "a"],
971
1
                vec!["a", "c"],
972
1
                Some(PartiallySorted(vec![1])),
973
1
            ),
974
1
            (
975
1
                vec!["c", "a"],
976
1
                vec!["a", "b", "c"],
977
1
                Some(PartiallySorted(vec![1])),
978
1
            ),
979
1
            (vec!["c", "b", "a"], vec!["a"], Some(Sorted)),
980
1
            (vec!["c", "b", "a"], vec!["b"], Some(Sorted)),
981
1
            (vec!["c", "b", "a"], vec!["c"], Some(Sorted)),
982
1
            (vec!["c", "b", "a"], vec!["a", "b"], Some(Sorted)),
983
1
            (vec!["c", "b", "a"], vec!["b", "c"], Some(Sorted)),
984
1
            (vec!["c", "b", "a"], vec!["a", "c"], Some(Sorted)),
985
1
            (vec!["c", "b", "a"], vec!["a", "b", "c"], Some(Sorted)),
986
1
        ];
987
49
        for (case_idx, test_case) in 
test_cases.iter().enumerate()1
{
988
49
            let (partition_by_columns, order_by_params, expected) = &test_case;
989
49
            let mut partition_by_exprs = vec![];
990
133
            for 
col_name84
in partition_by_columns {
991
84
                partition_by_exprs.push(col(col_name, &test_schema)
?0
);
992
1
            }
993
1
994
49
            let mut order_by_exprs = vec![];
995
133
            for 
col_name84
in order_by_params {
996
84
                let expr = col(col_name, &test_schema)
?0
;
997
1
                // Give default ordering, this is same with input ordering direction
998
1
                // In this test we do check for reversibility.
999
84
                let options = SortOptions::default();
1000
84
                order_by_exprs.push(PhysicalSortExpr { expr, options });
1001
1
            }
1002
49
            let res =
1003
49
                get_window_mode(&partition_by_exprs, &order_by_exprs, &exec_unbounded);
1004
49
            // Since reversibility is not important in this test. Convert Option<(bool, InputOrderMode)> to Option<InputOrderMode>
1005
49
            let res = res.map(|(_, mode)| 
mode43
);
1006
49
            assert_eq!(
1007
1
                res, *expected,
1008
1
                
"Unexpected result for in unbounded test case#: {case_idx:?}, case: {test_case:?}"0
1009
1
            );
1010
1
        }
1011
1
1012
1
        Ok(())
1013
1
    }
1014
1015
    #[tokio::test]
1016
1
    async fn test_get_window_mode() -> Result<()> {
1017
1
        let test_schema = create_test_schema3()
?0
;
1018
1
        // Columns a,c are nullable whereas b,d are not nullable.
1019
1
        // Source is sorted by a ASC NULLS FIRST, b ASC NULLS FIRST, c ASC NULLS FIRST, d ASC NULLS FIRST
1020
1
        // Column e is not ordered.
1021
1
        let sort_exprs = vec![
1022
1
            sort_expr("a", &test_schema),
1023
1
            sort_expr("b", &test_schema),
1024
1
            sort_expr("c", &test_schema),
1025
1
            sort_expr("d", &test_schema),
1026
1
        ];
1027
1
        let exec_unbounded = streaming_table_exec(&test_schema, sort_exprs, true)
?0
;
1028
1
1029
1
        // test cases consists of vector of tuples. Where each tuple represents a single test case.
1030
1
        // First field in the tuple is Vec<str> where each element in the vector represents PARTITION BY columns
1031
1
        // For instance `vec!["a", "b"]` corresponds to PARTITION BY a, b
1032
1
        // Second field in the tuple is Vec<(str, bool, bool)> where each element in the vector represents ORDER BY columns
1033
1
        // For instance, vec![("c", false, false)], corresponds to ORDER BY c ASC NULLS LAST,
1034
1
        // similarly, vec![("c", true, true)], corresponds to ORDER BY c DESC NULLS FIRST,
1035
1
        // Third field in the tuple is Option<(bool, InputOrderMode)>, which corresponds to expected result.
1036
1
        // None represents that existing ordering is not sufficient to run executor with any one of the algorithms
1037
1
        // (We need to add SortExec to be able to run it).
1038
1
        // Some((bool, InputOrderMode)) represents, we can run algorithm with existing ordering. Algorithm should work in
1039
1
        // InputOrderMode, bool field represents whether we should reverse window expressions to run executor with existing ordering.
1040
1
        // For instance, `Some((false, InputOrderMode::Sorted))`, represents that we shouldn't reverse window expressions. And algorithm
1041
1
        // should work in Sorted mode to work with existing ordering.
1042
1
        let test_cases = vec![
1043
1
            // PARTITION BY a, b ORDER BY c ASC NULLS LAST
1044
1
            (vec!["a", "b"], vec![("c", false, false)], None),
1045
1
            // ORDER BY c ASC NULLS FIRST
1046
1
            (vec![], vec![("c", false, true)], None),
1047
1
            // PARTITION BY b, ORDER BY c ASC NULLS FIRST
1048
1
            (vec!["b"], vec![("c", false, true)], None),
1049
1
            // PARTITION BY a, ORDER BY c ASC NULLS FIRST
1050
1
            (vec!["a"], vec![("c", false, true)], None),
1051
1
            // PARTITION BY b, ORDER BY c ASC NULLS FIRST
1052
1
            (
1053
1
                vec!["a", "b"],
1054
1
                vec![("c", false, true), ("e", false, true)],
1055
1
                None,
1056
1
            ),
1057
1
            // PARTITION BY a, ORDER BY b ASC NULLS FIRST
1058
1
            (vec!["a"], vec![("b", false, true)], Some((false, Sorted))),
1059
1
            // PARTITION BY a, ORDER BY a ASC NULLS FIRST
1060
1
            (vec!["a"], vec![("a", false, true)], Some((false, Sorted))),
1061
1
            // PARTITION BY a, ORDER BY a ASC NULLS LAST
1062
1
            (vec!["a"], vec![("a", false, false)], Some((false, Sorted))),
1063
1
            // PARTITION BY a, ORDER BY a DESC NULLS FIRST
1064
1
            (vec!["a"], vec![("a", true, true)], Some((false, Sorted))),
1065
1
            // PARTITION BY a, ORDER BY a DESC NULLS LAST
1066
1
            (vec!["a"], vec![("a", true, false)], Some((false, Sorted))),
1067
1
            // PARTITION BY a, ORDER BY b ASC NULLS LAST
1068
1
            (vec!["a"], vec![("b", false, false)], Some((false, Sorted))),
1069
1
            // PARTITION BY a, ORDER BY b DESC NULLS LAST
1070
1
            (vec!["a"], vec![("b", true, false)], Some((true, Sorted))),
1071
1
            // PARTITION BY a, b ORDER BY c ASC NULLS FIRST
1072
1
            (
1073
1
                vec!["a", "b"],
1074
1
                vec![("c", false, true)],
1075
1
                Some((false, Sorted)),
1076
1
            ),
1077
1
            // PARTITION BY b, a ORDER BY c ASC NULLS FIRST
1078
1
            (
1079
1
                vec!["b", "a"],
1080
1
                vec![("c", false, true)],
1081
1
                Some((false, Sorted)),
1082
1
            ),
1083
1
            // PARTITION BY a, b ORDER BY c DESC NULLS LAST
1084
1
            (
1085
1
                vec!["a", "b"],
1086
1
                vec![("c", true, false)],
1087
1
                Some((true, Sorted)),
1088
1
            ),
1089
1
            // PARTITION BY e ORDER BY a ASC NULLS FIRST
1090
1
            (
1091
1
                vec!["e"],
1092
1
                vec![("a", false, true)],
1093
1
                // For unbounded, expects to work in Linear mode. Shouldn't reverse window function.
1094
1
                Some((false, Linear)),
1095
1
            ),
1096
1
            // PARTITION BY b, c ORDER BY a ASC NULLS FIRST, c ASC NULLS FIRST
1097
1
            (
1098
1
                vec!["b", "c"],
1099
1
                vec![("a", false, true), ("c", false, true)],
1100
1
                Some((false, Linear)),
1101
1
            ),
1102
1
            // PARTITION BY b ORDER BY a ASC NULLS FIRST
1103
1
            (vec!["b"], vec![("a", false, true)], Some((false, Linear))),
1104
1
            // PARTITION BY a, e ORDER BY b ASC NULLS FIRST
1105
1
            (
1106
1
                vec!["a", "e"],
1107
1
                vec![("b", false, true)],
1108
1
                Some((false, PartiallySorted(vec![0]))),
1109
1
            ),
1110
1
            // PARTITION BY a, c ORDER BY b ASC NULLS FIRST
1111
1
            (
1112
1
                vec!["a", "c"],
1113
1
                vec![("b", false, true)],
1114
1
                Some((false, PartiallySorted(vec![0]))),
1115
1
            ),
1116
1
            // PARTITION BY c, a ORDER BY b ASC NULLS FIRST
1117
1
            (
1118
1
                vec!["c", "a"],
1119
1
                vec![("b", false, true)],
1120
1
                Some((false, PartiallySorted(vec![1]))),
1121
1
            ),
1122
1
            // PARTITION BY d, b, a ORDER BY c ASC NULLS FIRST
1123
1
            (
1124
1
                vec!["d", "b", "a"],
1125
1
                vec![("c", false, true)],
1126
1
                Some((false, PartiallySorted(vec![2, 1]))),
1127
1
            ),
1128
1
            // PARTITION BY e, b, a ORDER BY c ASC NULLS FIRST
1129
1
            (
1130
1
                vec!["e", "b", "a"],
1131
1
                vec![("c", false, true)],
1132
1
                Some((false, PartiallySorted(vec![2, 1]))),
1133
1
            ),
1134
1
            // PARTITION BY d, a ORDER BY b ASC NULLS FIRST
1135
1
            (
1136
1
                vec!["d", "a"],
1137
1
                vec![("b", false, true)],
1138
1
                Some((false, PartiallySorted(vec![1]))),
1139
1
            ),
1140
1
            // PARTITION BY b, ORDER BY b, a ASC NULLS FIRST
1141
1
            (
1142
1
                vec!["a"],
1143
1
                vec![("b", false, true), ("a", false, true)],
1144
1
                Some((false, Sorted)),
1145
1
            ),
1146
1
            // ORDER BY b, a ASC NULLS FIRST
1147
1
            (vec![], vec![("b", false, true), ("a", false, true)], None),
1148
1
        ];
1149
26
        for (case_idx, test_case) in 
test_cases.iter().enumerate()1
{
1150
26
            let (partition_by_columns, order_by_params, expected) = &test_case;
1151
26
            let mut partition_by_exprs = vec![];
1152
64
            for 
col_name38
in partition_by_columns {
1153
38
                partition_by_exprs.push(col(col_name, &test_schema)
?0
);
1154
1
            }
1155
1
1156
26
            let mut order_by_exprs = vec![];
1157
56
            for (
col_name, descending, nulls_first30
) in order_by_params {
1158
30
                let expr = col(col_name, &test_schema)
?0
;
1159
30
                let options = SortOptions {
1160
30
                    descending: *descending,
1161
30
                    nulls_first: *nulls_first,
1162
30
                };
1163
30
                order_by_exprs.push(PhysicalSortExpr { expr, options });
1164
1
            }
1165
1
1166
26
            assert_eq!(
1167
26
                get_window_mode(&partition_by_exprs, &order_by_exprs, &exec_unbounded),
1168
1
                *expected,
1169
1
                
"Unexpected result for in unbounded test case#: {case_idx:?}, case: {test_case:?}"0
1170
1
            );
1171
1
        }
1172
1
1173
1
        Ok(())
1174
1
    }
1175
}