## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/expr/src/expr_fn.rs

    Line
    Count
    Source
    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
    //! Functions for creating logical expressions
    19
    20
    use crate::expr::{
    21
        AggregateFunction, BinaryExpr, Cast, Exists, GroupingSet, InList, InSubquery,
    22
        NullTreatment, Placeholder, TryCast, Unnest, WildcardOptions, WindowFunction,
    23
    };
    24
    use crate::function::{
    25
        AccumulatorArgs, AccumulatorFactoryFunction, PartitionEvaluatorFactory,
    26
        StateFieldsArgs,
    27
    };
    28
    use crate::ptr_eq::PtrEq;
    29
    use crate::select_expr::SelectExpr;
    30
    use crate::{
    31
        conditional_expressions::CaseBuilder, expr::Sort, logical_plan::Subquery,
    32
        AggregateUDF, Expr, LimitEffect, LogicalPlan, Operator, PartitionEvaluator,
    33
        ScalarFunctionArgs, ScalarFunctionImplementation, ScalarUDF, Signature, Volatility,
    34
    };
    35
    use crate::{
    36
        AggregateUDFImpl, ColumnarValue, ScalarUDFImpl, WindowFrame, WindowUDF, WindowUDFImpl,
    37
    };
    38
    use arrow::compute::kernels::cast_utils::{
    39
        parse_interval_day_time, parse_interval_month_day_nano, parse_interval_year_month,
    40
    };
    41
    use arrow::datatypes::{DataType, Field, FieldRef};
    42
    use datafusion_common::{plan_err, Column, Result, ScalarValue, Spans, TableReference};
    43
    use datafusion_functions_window_common::field::WindowUDFFieldArgs;
    44
    use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
    45
    use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
    46
    use std::any::Any;
    47
    use std::collections::HashMap;
    48
    use std::fmt::Debug;
    49
    use std::hash::Hash;
    50
    use std::ops::Not;
    51
    use std::sync::Arc;
    52
    53
    /// Create a column expression based on a qualified or unqualified column name. Will
    54
    /// normalize unquoted identifiers according to SQL rules (identifiers will become lowercase).
    55
    ///
    56
    /// For example:
    57
    ///
    58
    /// ```rust
    59
    /// # use datafusion_expr::col;
    60
    /// let c1 = col("a");
    61
    /// let c2 = col("A");
    62
    /// assert_eq!(c1, c2);
    63
    ///
    64
    /// // note how quoting with double quotes preserves the case
    65
    /// let c3 = col(r#""A""#);
    66
    /// assert_ne!(c1, c3);
    67
    /// ```
    68
    0
    pub fn col(ident: impl Into<Column>) -> Expr {
    69
    0
        Expr::Column(ident.into())
    70
    0
    }
    71
    72
    /// Create an out reference column which hold a reference that has been resolved to a field
    73
    /// outside of the current plan.
    74
    /// The expression created by this function does not preserve the metadata of the outer column.
    75
    /// Please use `out_ref_col_with_metadata` if you want to preserve the metadata.
    76
    0
    pub fn out_ref_col(dt: DataType, ident: impl Into<Column>) -> Expr {
    77
    0
        out_ref_col_with_metadata(dt, HashMap::new(), ident)
    78
    0
    }
    79
    80
    /// Create an out reference column from an existing field (preserving metadata)
    81
    0
    pub fn out_ref_col_with_metadata(
    82
    0
        dt: DataType,
    83
    0
        metadata: HashMap<String, String>,
    84
    0
        ident: impl Into<Column>,
    85
    0
    ) -> Expr {
    86
    0
        let column = ident.into();
    87
    0
        let field: FieldRef =
    88
    0
            Arc::new(Field::new(column.name(), dt, true).with_metadata(metadata));
    89
    0
        Expr::OuterReferenceColumn(field, column)
    90
    0
    }
    91
    92
    /// Create an unqualified column expression from the provided name, without normalizing
    93
    /// the column.
    94
    ///
    95
    /// For example:
    96
    ///
    97
    /// ```rust
    98
    /// # use datafusion_expr::{col, ident};
    99
    /// let c1 = ident("A"); // not normalized staying as column 'A'
    100
    /// let c2 = col("A"); // normalized via SQL rules becoming column 'a'
    101
    /// assert_ne!(c1, c2);
    102
    ///
    103
    /// let c3 = col(r#""A""#);
    104
    /// assert_eq!(c1, c3);
    105
    ///
    106
    /// let c4 = col("t1.a"); // parses as relation 't1' column 'a'
    107
    /// let c5 = ident("t1.a"); // parses as column 't1.a'
    108
    /// assert_ne!(c4, c5);
    109
    /// ```
    110
    0
    pub fn ident(name: impl Into<String>) -> Expr {
    111
    0
        Expr::Column(Column::from_name(name))
    112
    0
    }
    113
    114
    /// Create placeholder value that will be filled in (such as `$1`)
    115
    ///
    116
    /// Note the parameter type can be inferred using [`Expr::infer_placeholder_types`]
    117
    ///
    118
    /// # Example
    119
    ///
    120
    /// ```rust
    121
    /// # use datafusion_expr::{placeholder};
    122
    /// let p = placeholder("$1"); // $1, refers to parameter 1
    123
    /// assert_eq!(p.to_string(), "$1")
    124
    /// ```
    125
    0
    pub fn placeholder(id: impl Into<String>) -> Expr {
    126
    0
        Expr::Placeholder(Placeholder {
    127
    0
            id: id.into(),
    128
    0
            field: None,
    129
    0
        })
    130
    0
    }
    131
    132
    /// Create an '*' [`Expr::Wildcard`] expression that matches all columns
    133
    ///
    134
    /// # Example
    135
    ///
    136
    /// ```rust
    137
    /// # use datafusion_expr::{wildcard};
    138
    /// let p = wildcard();
    139
    /// assert_eq!(p.to_string(), "*")
    140
    /// ```
    141
    0
    pub fn wildcard() -> SelectExpr {
    142
    0
        SelectExpr::Wildcard(WildcardOptions::default())
    143
    0
    }
    144
    145
    /// Create an '*' [`Expr::Wildcard`] expression with the wildcard options
    146
    0
    pub fn wildcard_with_options(options: WildcardOptions) -> SelectExpr {
    147
    0
        SelectExpr::Wildcard(options)
    148
    0
    }
    149
    150
    /// Create an 't.*' [`Expr::Wildcard`] expression that matches all columns from a specific table
    151
    ///
    152
    /// # Example
    153
    ///
    154
    /// ```rust
    155
    /// # use datafusion_common::TableReference;
    156
    /// # use datafusion_expr::{qualified_wildcard};
    157
    /// let p = qualified_wildcard(TableReference::bare("t"));
    158
    /// assert_eq!(p.to_string(), "t.*")
    159
    /// ```
    160
    0
    pub fn qualified_wildcard(qualifier: impl Into<TableReference>) -> SelectExpr {
    161
    0
        SelectExpr::QualifiedWildcard(qualifier.into(), WildcardOptions::default())
    162
    0
    }
    163
    164
    /// Create an 't.*' [`Expr::Wildcard`] expression with the wildcard options
    165
    0
    pub fn qualified_wildcard_with_options(
    166
    0
        qualifier: impl Into<TableReference>,
    167
    0
        options: WildcardOptions,
    168
    0
    ) -> SelectExpr {
    169
    0
        SelectExpr::QualifiedWildcard(qualifier.into(), options)
    170
    0
    }
    171
    172
    /// Return a new expression `left <op> right`
    173
    0
    pub fn binary_expr(left: Expr, op: Operator, right: Expr) -> Expr {
    174
    0
        Expr::BinaryExpr(BinaryExpr::new(Box::new(left), op, Box::new(right)))
    175
    0
    }
    176
    177
    /// Return a new expression with a logical AND
    178
    0
    pub fn and(left: Expr, right: Expr) -> Expr {
    179
    0
        Expr::BinaryExpr(BinaryExpr::new(
    180
    0
            Box::new(left),
    181
    0
            Operator::And,
    182
    0
            Box::new(right),
    183
    0
        ))
    184
    0
    }
    185
    186
    /// Return a new expression with a logical OR
    187
    0
    pub fn or(left: Expr, right: Expr) -> Expr {
    188
    0
        Expr::BinaryExpr(BinaryExpr::new(
    189
    0
            Box::new(left),
    190
    0
            Operator::Or,
    191
    0
            Box::new(right),
    192
    0
        ))
    193
    0
    }
    194
    195
    /// Return a new expression with a logical NOT
    196
    0
    pub fn not(expr: Expr) -> Expr {
    197
    0
        expr.not()
    198
    0
    }
    199
    200
    /// Return a new expression with bitwise AND
    201
    0
    pub fn bitwise_and(left: Expr, right: Expr) -> Expr {
    202
    0
        Expr::BinaryExpr(BinaryExpr::new(
    203
    0
            Box::new(left),
    204
    0
            Operator::BitwiseAnd,
    205
    0
            Box::new(right),
    206
    0
        ))
    207
    0
    }
    208
    209
    /// Return a new expression with bitwise OR
    210
    0
    pub fn bitwise_or(left: Expr, right: Expr) -> Expr {
    211
    0
        Expr::BinaryExpr(BinaryExpr::new(
    212
    0
            Box::new(left),
    213
    0
            Operator::BitwiseOr,
    214
    0
            Box::new(right),
    215
    0
        ))
    216
    0
    }
    217
    218
    /// Return a new expression with bitwise XOR
    219
    0
    pub fn bitwise_xor(left: Expr, right: Expr) -> Expr {
    220
    0
        Expr::BinaryExpr(BinaryExpr::new(
    221
    0
            Box::new(left),
    222
    0
            Operator::BitwiseXor,
    223
    0
            Box::new(right),
    224
    0
        ))
    225
    0
    }
    226
    227
    /// Return a new expression with bitwise SHIFT RIGHT
    228
    0
    pub fn bitwise_shift_right(left: Expr, right: Expr) -> Expr {
    229
    0
        Expr::BinaryExpr(BinaryExpr::new(
    230
    0
            Box::new(left),
    231
    0
            Operator::BitwiseShiftRight,
    232
    0
            Box::new(right),
    233
    0
        ))
    234
    0
    }
    235
    236
    /// Return a new expression with bitwise SHIFT LEFT
    237
    0
    pub fn bitwise_shift_left(left: Expr, right: Expr) -> Expr {
    238
    0
        Expr::BinaryExpr(BinaryExpr::new(
    239
    0
            Box::new(left),
    240
    0
            Operator::BitwiseShiftLeft,
    241
    0
            Box::new(right),
    242
    0
        ))
    243
    0
    }
    244
    245
    /// Create an in_list expression
    246
    0
    pub fn in_list(expr: Expr, list: Vec<Expr>, negated: bool) -> Expr {
    247
    0
        Expr::InList(InList::new(Box::new(expr), list, negated))
    248
    0
    }
    249
    250
    /// Create an EXISTS subquery expression
    251
    0
    pub fn exists(subquery: Arc<LogicalPlan>) -> Expr {
    252
    0
        let outer_ref_columns = subquery.all_out_ref_exprs();
    253
    0
        Expr::Exists(Exists {
    254
    0
            subquery: Subquery {
    255
    0
                subquery,
    256
    0
                outer_ref_columns,
    257
    0
                spans: Spans::new(),
    258
    0
            },
    259
    0
            negated: false,
    260
    0
        })
    261
    0
    }
    262
    263
    /// Create a NOT EXISTS subquery expression
    264
    0
    pub fn not_exists(subquery: Arc<LogicalPlan>) -> Expr {
    265
    0
        let outer_ref_columns = subquery.all_out_ref_exprs();
    266
    0
        Expr::Exists(Exists {
    267
    0
            subquery: Subquery {
    268
    0
                subquery,
    269
    0
                outer_ref_columns,
    270
    0
                spans: Spans::new(),
    271
    0
            },
    272
    0
            negated: true,
    273
    0
        })
    274
    0
    }
    275
    276
    /// Create an IN subquery expression
    277
    0
    pub fn in_subquery(expr: Expr, subquery: Arc<LogicalPlan>) -> Expr {
    278
    0
        let outer_ref_columns = subquery.all_out_ref_exprs();
    279
    0
        Expr::InSubquery(InSubquery::new(
    280
    0
            Box::new(expr),
    281
    0
            Subquery {
    282
    0
                subquery,
    283
    0
                outer_ref_columns,
    284
    0
                spans: Spans::new(),
    285
    0
            },
    286
    0
            false,
    287
    0
        ))
    288
    0
    }
    289
    290
    /// Create a NOT IN subquery expression
    291
    0
    pub fn not_in_subquery(expr: Expr, subquery: Arc<LogicalPlan>) -> Expr {
    292
    0
        let outer_ref_columns = subquery.all_out_ref_exprs();
    293
    0
        Expr::InSubquery(InSubquery::new(
    294
    0
            Box::new(expr),
    295
    0
            Subquery {
    296
    0
                subquery,
    297
    0
                outer_ref_columns,
    298
    0
                spans: Spans::new(),
    299
    0
            },
    300
    0
            true,
    301
    0
        ))
    302
    0
    }
    303
    304
    /// Create a scalar subquery expression
    305
    0
    pub fn scalar_subquery(subquery: Arc<LogicalPlan>) -> Expr {
    306
    0
        let outer_ref_columns = subquery.all_out_ref_exprs();
    307
    0
        Expr::ScalarSubquery(Subquery {
    308
    0
            subquery,
    309
    0
            outer_ref_columns,
    310
    0
            spans: Spans::new(),
    311
    0
        })
    312
    0
    }
    313
    314
    /// Create a grouping set
    315
    0
    pub fn grouping_set(exprs: Vec<Vec<Expr>>) -> Expr {
    316
    0
        Expr::GroupingSet(GroupingSet::GroupingSets(exprs))
    317
    0
    }
    318
    319
    /// Create a grouping set for all combination of `exprs`
    320
    0
    pub fn cube(exprs: Vec<Expr>) -> Expr {
    321
    0
        Expr::GroupingSet(GroupingSet::Cube(exprs))
    322
    0
    }
    323
    324
    /// Create a grouping set for rollup
    325
    0
    pub fn rollup(exprs: Vec<Expr>) -> Expr {
    326
    0
        Expr::GroupingSet(GroupingSet::Rollup(exprs))
    327
    0
    }
    328
    329
    /// Create a cast expression
    330
    0
    pub fn cast(expr: Expr, data_type: DataType) -> Expr {
    331
    0
        Expr::Cast(Cast::new(Box::new(expr), data_type))
    332
    0
    }
    333
    334
    /// Create a try cast expression
    335
    0
    pub fn try_cast(expr: Expr, data_type: DataType) -> Expr {
    336
    0
        Expr::TryCast(TryCast::new(Box::new(expr), data_type))
    337
    0
    }
    338
    339
    /// Create is null expression
    340
    0
    pub fn is_null(expr: Expr) -> Expr {
    341
    0
        Expr::IsNull(Box::new(expr))
    342
    0
    }
    343
    344
    /// Create is not null expression
    345
    0
    pub fn is_not_null(expr: Expr) -> Expr {
    346
    0
        Expr::IsNotNull(Box::new(expr))
    347
    0
    }
    348
    349
    /// Create is true expression
    350
    0
    pub fn is_true(expr: Expr) -> Expr {
    351
    0
        Expr::IsTrue(Box::new(expr))
    352
    0
    }
    353
    354
    /// Create is not true expression
    355
    0
    pub fn is_not_true(expr: Expr) -> Expr {
    356
    0
        Expr::IsNotTrue(Box::new(expr))
    357
    0
    }
    358
    359
    /// Create is false expression
    360
    0
    pub fn is_false(expr: Expr) -> Expr {
    361
    0
        Expr::IsFalse(Box::new(expr))
    362
    0
    }
    363
    364
    /// Create is not false expression
    365
    0
    pub fn is_not_false(expr: Expr) -> Expr {
    366
    0
        Expr::IsNotFalse(Box::new(expr))
    367
    0
    }
    368
    369
    /// Create is unknown expression
    370
    0
    pub fn is_unknown(expr: Expr) -> Expr {
    371
    0
        Expr::IsUnknown(Box::new(expr))
    372
    0
    }
    373
    374
    /// Create is not unknown expression
    375
    0
    pub fn is_not_unknown(expr: Expr) -> Expr {
    376
    0
        Expr::IsNotUnknown(Box::new(expr))
    377
    0
    }
    378
    379
    /// Create a CASE WHEN statement with literal WHEN expressions for comparison to the base expression.
    380
    0
    pub fn case(expr: Expr) -> CaseBuilder {
    381
    0
        CaseBuilder::new(Some(Box::new(expr)), vec![], vec![], None)
    382
    0
    }
    383
    384
    /// Create a CASE WHEN statement with boolean WHEN expressions and no base expression.
    385
    0
    pub fn when(when: Expr, then: Expr) -> CaseBuilder {
    386
    0
        CaseBuilder::new(None, vec![when], vec![then], None)
    387
    0
    }
    388
    389
    /// Create a Unnest expression
    390
    0
    pub fn unnest(expr: Expr) -> Expr {
    391
    0
        Expr::Unnest(Unnest {
    392
    0
            expr: Box::new(expr),
    393
    0
        })
    394
    0
    }
    395
    396
    /// Convenience method to create a new user defined scalar function (UDF) with a
    397
    /// specific signature and specific return type.
    398
    ///
    399
    /// Note this function does not expose all available features of [`ScalarUDF`],
    400
    /// such as
    401
    ///
    402
    /// * computing return types based on input types
    403
    /// * multiple [`Signature`]s
    404
    /// * aliases
    405
    ///
    406
    /// See [`ScalarUDF`] for details and examples on how to use the full
    407
    /// functionality.
    408
    0
    pub fn create_udf(
    409
    0
        name: &str,
    410
    0
        input_types: Vec<DataType>,
    411
    0
        return_type: DataType,
    412
    0
        volatility: Volatility,
    413
    0
        fun: ScalarFunctionImplementation,
    414
    0
    ) -> ScalarUDF {
    415
    0
        ScalarUDF::from(SimpleScalarUDF::new(
    416
    0
            name,
    417
    0
            input_types,
    418
    0
            return_type,
    419
    0
            volatility,
    420
    0
            fun,
    421
        ))
    422
    0
    }
    423
    424
    /// Implements [`ScalarUDFImpl`] for functions that have a single signature and
    425
    /// return type.
    426
    #[derive(PartialEq, Eq, Hash)]
    427
    pub struct SimpleScalarUDF {
    428
        name: String,
    429
        signature: Signature,
    430
        return_type: DataType,
    431
        fun: PtrEq<ScalarFunctionImplementation>,
    432
    }
    433
    434
    impl Debug for SimpleScalarUDF {
    435
    0
        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
    436
    0
            f.debug_struct("SimpleScalarUDF")
    437
    0
                .field("name", &self.name)
    438
    0
                .field("signature", &self.signature)
    439
    0
                .field("return_type", &self.return_type)
    440
    0
                .field("fun", &"<FUNC>")
    441
    0
                .finish()
    442
    0
        }
    443
    }
    444
    445
    impl SimpleScalarUDF {
    446
        /// Create a new `SimpleScalarUDF` from a name, input types, return type and
    447
        /// implementation. Implementing [`ScalarUDFImpl`] allows more flexibility
    448
    0
        pub fn new(
    449
    0
            name: impl Into<String>,
    450
    0
            input_types: Vec<DataType>,
    451
    0
            return_type: DataType,
    452
    0
            volatility: Volatility,
    453
    0
            fun: ScalarFunctionImplementation,
    454
    0
        ) -> Self {
    455
    0
            Self::new_with_signature(
    456
    0
                name,
    457
    0
                Signature::exact(input_types, volatility),
    458
    0
                return_type,
    459
    0
                fun,
    460
            )
    461
    0
        }
    462
    463
        /// Create a new `SimpleScalarUDF` from a name, signature, return type and
    464
        /// implementation. Implementing [`ScalarUDFImpl`] allows more flexibility
    465
    0
        pub fn new_with_signature(
    466
    0
            name: impl Into<String>,
    467
    0
            signature: Signature,
    468
    0
            return_type: DataType,
    469
    0
            fun: ScalarFunctionImplementation,
    470
    0
        ) -> Self {
    471
    0
            Self {
    472
    0
                name: name.into(),
    473
    0
                signature,
    474
    0
                return_type,
    475
    0
                fun: fun.into(),
    476
    0
            }
    477
    0
        }
    478
    }
    479
    480
    impl ScalarUDFImpl for SimpleScalarUDF {
    481
    0
        fn as_any(&self) -> &dyn Any {
    482
    0
            self
    483
    0
        }
    484
    485
    0
        fn name(&self) -> &str {
    486
    0
            &self.name
    487
    0
        }
    488
    489
    0
        fn signature(&self) -> &Signature {
    490
    0
            &self.signature
    491
    0
        }
    492
    493
    0
        fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
    494
    0
            Ok(self.return_type.clone())
    495
    0
        }
    496
    497
    0
        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    498
    0
            (self.fun)(&args.args)
    499
    0
        }
    500
    }
    501
    502
    /// Creates a new UDAF with a specific signature, state type and return type.
    503
    /// The signature and state type must match the `Accumulator's implementation`.
    504
    0
    pub fn create_udaf(
    505
    0
        name: &str,
    506
    0
        input_type: Vec<DataType>,
    507
    0
        return_type: Arc<DataType>,
    508
    0
        volatility: Volatility,
    509
    0
        accumulator: AccumulatorFactoryFunction,
    510
    0
        state_type: Arc<Vec<DataType>>,
    511
    0
    ) -> AggregateUDF {
    512
    0
        let return_type = Arc::unwrap_or_clone(return_type);
    513
    0
        let state_type = Arc::unwrap_or_clone(state_type);
    514
    0
        let state_fields = state_type
    515
    0
            .into_iter()
    516
    0
            .enumerate()
    517
    0
            .map(|(i, t)| Field::new(format!("{i}"), t, true))
    518
    0
            .map(Arc::new)
    519
    0
            .collect::<Vec<_>>();
    520
    0
        AggregateUDF::from(SimpleAggregateUDF::new(
    521
    0
            name,
    522
    0
            input_type,
    523
    0
            return_type,
    524
    0
            volatility,
    525
    0
            accumulator,
    526
    0
            state_fields,
    527
        ))
    528
    0
    }
    529
    530
    /// Implements [`AggregateUDFImpl`] for functions that have a single signature and
    531
    /// return type.
    532
    #[derive(PartialEq, Eq, Hash)]
    533
    pub struct SimpleAggregateUDF {
    534
        name: String,
    535
        signature: Signature,
    536
        return_type: DataType,
    537
        accumulator: PtrEq<AccumulatorFactoryFunction>,
    538
        state_fields: Vec<FieldRef>,
    539
    }
    540
    541
    impl Debug for SimpleAggregateUDF {
    542
    0
        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
    543
    0
            f.debug_struct("SimpleAggregateUDF")
    544
    0
                .field("name", &self.name)
    545
    0
                .field("signature", &self.signature)
    546
    0
                .field("return_type", &self.return_type)
    547
    0
                .field("fun", &"<FUNC>")
    548
    0
                .finish()
    549
    0
        }
    550
    }
    551
    552
    impl SimpleAggregateUDF {
    553
        /// Create a new `SimpleAggregateUDF` from a name, input types, return type, state type and
    554
        /// implementation. Implementing [`AggregateUDFImpl`] allows more flexibility
    555
    0
        pub fn new(
    556
    0
            name: impl Into<String>,
    557
    0
            input_type: Vec<DataType>,
    558
    0
            return_type: DataType,
    559
    0
            volatility: Volatility,
    560
    0
            accumulator: AccumulatorFactoryFunction,
    561
    0
            state_fields: Vec<FieldRef>,
    562
    0
        ) -> Self {
    563
    0
            let name = name.into();
    564
    0
            let signature = Signature::exact(input_type, volatility);
    565
    0
            Self {
    566
    0
                name,
    567
    0
                signature,
    568
    0
                return_type,
    569
    0
                accumulator: accumulator.into(),
    570
    0
                state_fields,
    571
    0
            }
    572
    0
        }
    573
    574
        /// Create a new `SimpleAggregateUDF` from a name, signature, return type, state type and
    575
        /// implementation. Implementing [`AggregateUDFImpl`] allows more flexibility
    576
    0
        pub fn new_with_signature(
    577
    0
            name: impl Into<String>,
    578
    0
            signature: Signature,
    579
    0
            return_type: DataType,
    580
    0
            accumulator: AccumulatorFactoryFunction,
    581
    0
            state_fields: Vec<FieldRef>,
    582
    0
        ) -> Self {
    583
    0
            let name = name.into();
    584
    0
            Self {
    585
    0
                name,
    586
    0
                signature,
    587
    0
                return_type,
    588
    0
                accumulator: accumulator.into(),
    589
    0
                state_fields,
    590
    0
            }
    591
    0
        }
    592
    }
    593
    594
    impl AggregateUDFImpl for SimpleAggregateUDF {
    595
    0
        fn as_any(&self) -> &dyn Any {
    596
    0
            self
    597
    0
        }
    598
    599
    0
        fn name(&self) -> &str {
    600
    0
            &self.name
    601
    0
        }
    602
    603
    0
        fn signature(&self) -> &Signature {
    604
    0
            &self.signature
    605
    0
        }
    606
    607
    0
        fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
    608
    0
            Ok(self.return_type.clone())
    609
    0
        }
    610
    611
    0
        fn accumulator(
    612
    0
            &self,
    613
    0
            acc_args: AccumulatorArgs,
    614
    0
        ) -> Result<Box<dyn crate::Accumulator>> {
    615
    0
            (self.accumulator)(acc_args)
    616
    0
        }
    617
    618
    0
        fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
    619
    0
            Ok(self.state_fields.clone())
    620
    0
        }
    621
    }
    622
    623
    /// Creates a new UDWF with a specific signature, state type and return type.
    624
    ///
    625
    /// The signature and state type must match the [`PartitionEvaluator`]'s implementation`.
    626
    ///
    627
    /// [`PartitionEvaluator`]: crate::PartitionEvaluator
    628
    0
    pub fn create_udwf(
    629
    0
        name: &str,
    630
    0
        input_type: DataType,
    631
    0
        return_type: Arc<DataType>,
    632
    0
        volatility: Volatility,
    633
    0
        partition_evaluator_factory: PartitionEvaluatorFactory,
    634
    0
    ) -> WindowUDF {
    635
    0
        let return_type = Arc::unwrap_or_clone(return_type);
    636
    0
        WindowUDF::from(SimpleWindowUDF::new(
    637
    0
            name,
    638
    0
            input_type,
    639
    0
            return_type,
    640
    0
            volatility,
    641
    0
            partition_evaluator_factory,
    642
        ))
    643
    0
    }
    644
    645
    /// Implements [`WindowUDFImpl`] for functions that have a single signature and
    646
    /// return type.
    647
    #[derive(PartialEq, Eq, Hash)]
    648
    pub struct SimpleWindowUDF {
    649
        name: String,
    650
        signature: Signature,
    651
        return_type: DataType,
    652
        partition_evaluator_factory: PtrEq<PartitionEvaluatorFactory>,
    653
    }
    654
    655
    impl Debug for SimpleWindowUDF {
    656
    0
        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
    657
    0
            f.debug_struct("WindowUDF")
    658
    0
                .field("name", &self.name)
    659
    0
                .field("signature", &self.signature)
    660
    0
                .field("return_type", &"<func>")
    661
    0
                .field("partition_evaluator_factory", &"<FUNC>")
    662
    0
                .finish()
    663
    0
        }
    664
    }
    665
    666
    impl SimpleWindowUDF {
    667
        /// Create a new `SimpleWindowUDF` from a name, input types, return type and
    668
        /// implementation. Implementing [`WindowUDFImpl`] allows more flexibility
    669
    0
        pub fn new(
    670
    0
            name: impl Into<String>,
    671
    0
            input_type: DataType,
    672
    0
            return_type: DataType,
    673
    0
            volatility: Volatility,
    674
    0
            partition_evaluator_factory: PartitionEvaluatorFactory,
    675
    0
        ) -> Self {
    676
    0
            let name = name.into();
    677
    0
            let signature = Signature::exact([input_type].to_vec(), volatility);
    678
    0
            Self {
    679
    0
                name,
    680
    0
                signature,
    681
    0
                return_type,
    682
    0
                partition_evaluator_factory: partition_evaluator_factory.into(),
    683
    0
            }
    684
    0
        }
    685
    }
    686
    687
    impl WindowUDFImpl for SimpleWindowUDF {
    688
    0
        fn as_any(&self) -> &dyn Any {
    689
    0
            self
    690
    0
        }
    691
    692
    0
        fn name(&self) -> &str {
    693
    0
            &self.name
    694
    0
        }
    695
    696
    0
        fn signature(&self) -> &Signature {
    697
    0
            &self.signature
    698
    0
        }
    699
    700
    0
        fn partition_evaluator(
    701
    0
            &self,
    702
    0
            _partition_evaluator_args: PartitionEvaluatorArgs,
    703
    0
        ) -> Result<Box<dyn PartitionEvaluator>> {
    704
    0
            (self.partition_evaluator_factory)()
    705
    0
        }
    706
    707
    0
        fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
    708
    0
            Ok(Arc::new(Field::new(
    709
    0
                field_args.name(),
    710
    0
                self.return_type.clone(),
    711
    0
                true,
    712
    0
            )))
    713
    0
        }
    714
    715
    0
        fn limit_effect(&self, _args: &[Arc<dyn PhysicalExpr>]) -> LimitEffect {
    716
    0
            LimitEffect::Unknown
    717
    0
        }
    718
    }
    719
    720
    0
    pub fn interval_year_month_lit(value: &str) -> Expr {
    721
    0
        let interval = parse_interval_year_month(value).ok();
    722
    0
        Expr::Literal(ScalarValue::IntervalYearMonth(interval), None)
    723
    0
    }
    724
    725
    0
    pub fn interval_datetime_lit(value: &str) -> Expr {
    726
    0
        let interval = parse_interval_day_time(value).ok();
    727
    0
        Expr::Literal(ScalarValue::IntervalDayTime(interval), None)
    728
    0
    }
    729
    730
    0
    pub fn interval_month_day_nano_lit(value: &str) -> Expr {
    731
    0
        let interval = parse_interval_month_day_nano(value).ok();
    732
    0
        Expr::Literal(ScalarValue::IntervalMonthDayNano(interval), None)
    733
    0
    }
    734
    735
    /// Extensions for configuring [`Expr::AggregateFunction`] or [`Expr::WindowFunction`]
    736
    ///
    737
    /// Adds methods to [`Expr`] that make it easy to set optional options
    738
    /// such as `ORDER BY`, `FILTER` and `DISTINCT`
    739
    ///
    740
    /// # Example
    741
    /// ```no_run
    742
    /// # use datafusion_common::Result;
    743
    /// # use datafusion_expr::expr::NullTreatment;
    744
    /// # use datafusion_expr::test::function_stub::count;
    745
    /// # use datafusion_expr::{ExprFunctionExt, lit, Expr, col};
    746
    /// # // first_value is an aggregate function in another crate
    747
    /// # fn first_value(_arg: Expr) -> Expr {
    748
    /// unimplemented!() }
    749
    /// # fn main() -> Result<()> {
    750
    /// // Create an aggregate count, filtering on column y > 5
    751
    /// let agg = count(col("x")).filter(col("y").gt(lit(5))).build()?;
    752
    ///
    753
    /// // Find the first value in an aggregate sorted by column y
    754
    /// // equivalent to:
    755
    /// // `FIRST_VALUE(x ORDER BY y ASC IGNORE NULLS)`
    756
    /// let sort_expr = col("y").sort(true, true);
    757
    /// let agg = first_value(col("x"))
    758
    ///     .order_by(vec![sort_expr])
    759
    ///     .null_treatment(NullTreatment::IgnoreNulls)
    760
    ///     .build()?;
    761
    ///
    762
    /// // Create a window expression for percent rank partitioned on column a
    763
    /// // equivalent to:
    764
    /// // `PERCENT_RANK() OVER (PARTITION BY a ORDER BY b ASC NULLS LAST IGNORE NULLS)`
    765
    /// // percent_rank is an udwf function in another crate
    766
    /// # fn percent_rank() -> Expr {
    767
    /// unimplemented!() }
    768
    /// let window = percent_rank()
    769
    ///     .partition_by(vec![col("a")])
    770
    ///     .order_by(vec![col("b").sort(true, true)])
    771
    ///     .null_treatment(NullTreatment::IgnoreNulls)
    772
    ///     .build()?;
    773
    /// #     Ok(())
    774
    /// # }
    775
    /// ```
    776
    pub trait ExprFunctionExt {
    777
        /// Add `ORDER BY <order_by>`
    778
        fn order_by(self, order_by: Vec<Sort>) -> ExprFuncBuilder;
    779
        /// Add `FILTER <filter>`
    780
        fn filter(self, filter: Expr) -> ExprFuncBuilder;
    781
        /// Add `DISTINCT`
    782
        fn distinct(self) -> ExprFuncBuilder;
    783
        /// Add `RESPECT NULLS` or `IGNORE NULLS`
    784
        fn null_treatment(
    785
            self,
    786
            null_treatment: impl Into<Option<NullTreatment>>,
    787
        ) -> ExprFuncBuilder;
    788
        /// Add `PARTITION BY`
    789
        fn partition_by(self, partition_by: Vec<Expr>) -> ExprFuncBuilder;
    790
        /// Add appropriate window frame conditions
    791
        fn window_frame(self, window_frame: WindowFrame) -> ExprFuncBuilder;
    792
    }
    793
    794
    #[derive(Debug, Clone)]
    795
    pub enum ExprFuncKind {
    796
        Aggregate(AggregateFunction),
    797
        Window(Box<WindowFunction>),
    798
    }
    799
    800
    /// Implementation of [`ExprFunctionExt`].
    801
    ///
    802
    /// See [`ExprFunctionExt`] for usage and examples
    803
    #[derive(Debug, Clone)]
    804
    pub struct ExprFuncBuilder {
    805
        fun: Option<ExprFuncKind>,
    806
        order_by: Option<Vec<Sort>>,
    807
        filter: Option<Expr>,
    808
        distinct: bool,
    809
        null_treatment: Option<NullTreatment>,
    810
        partition_by: Option<Vec<Expr>>,
    811
        window_frame: Option<WindowFrame>,
    812
    }
    813
    814
    impl ExprFuncBuilder {
    815
        /// Create a new `ExprFuncBuilder`, see [`ExprFunctionExt`]
    816
    0
        fn new(fun: Option<ExprFuncKind>) -> Self {
    817
    0
            Self {
    818
    0
                fun,
    819
    0
                order_by: None,
    820
    0
                filter: None,
    821
    0
                distinct: false,
    822
    0
                null_treatment: None,
    823
    0
                partition_by: None,
    824
    0
                window_frame: None,
    825
    0
            }
    826
    0
        }
    827
    828
        /// Updates and returns the in progress [`Expr::AggregateFunction`] or [`Expr::WindowFunction`]
    829
        ///
    830
        /// # Errors:
    831
        ///
    832
        /// Returns an error if this builder  [`ExprFunctionExt`] was used with an
    833
        /// `Expr` variant other than [`Expr::AggregateFunction`] or [`Expr::WindowFunction`]
    834
    0
        pub fn build(self) -> Result<Expr> {
    835
            let Self {
    836
    0
                fun,
    837
    0
                order_by,
    838
    0
                filter,
    839
    0
                distinct,
    840
    0
                null_treatment,
    841
    0
                partition_by,
    842
    0
                window_frame,
    843
    0
            } = self;
    844
    845
    0
            let Some(fun) = fun else {
    846
    0
                return plan_err!(
    847
    0
                    "ExprFunctionExt can only be used with Expr::AggregateFunction or Expr::WindowFunction"
    848
                );
    849
            };
    850
    851
    0
            let fun_expr = match fun {
    852
    0
                ExprFuncKind::Aggregate(mut udaf) => {
    853
    0
                    udaf.params.order_by = order_by.unwrap_or_default();
    854
    0
                    udaf.params.filter = filter.map(Box::new);
    855
    0
                    udaf.params.distinct = distinct;
    856
    0
                    udaf.params.null_treatment = null_treatment;
    857
    0
                    Expr::AggregateFunction(udaf)
    858
                }
    859
    0
                ExprFuncKind::Window(mut udwf) => {
    860
    0
                    let has_order_by = order_by.as_ref().map(|o| !o.is_empty());
    861
    0
                    udwf.params.partition_by = partition_by.unwrap_or_default();
    862
    0
                    udwf.params.order_by = order_by.unwrap_or_default();
    863
    0
                    udwf.params.window_frame =
    864
    0
                        window_frame.unwrap_or_else(|| WindowFrame::new(has_order_by));
    865
    0
                    udwf.params.filter = filter.map(Box::new);
    866
    0
                    udwf.params.null_treatment = null_treatment;
    867
    0
                    udwf.params.distinct = distinct;
    868
    0
                    Expr::WindowFunction(udwf)
    869
                }
    870
            };
    871
    872
    0
            Ok(fun_expr)
    873
    0
        }
    874
    }
    875
    876
    impl ExprFunctionExt for ExprFuncBuilder {
    877
        /// Add `ORDER BY <order_by>`
    878
    0
        fn order_by(mut self, order_by: Vec<Sort>) -> ExprFuncBuilder {
    879
    0
            self.order_by = Some(order_by);
    880
    0
            self
    881
    0
        }
    882
    883
        /// Add `FILTER <filter>`
    884
    0
        fn filter(mut self, filter: Expr) -> ExprFuncBuilder {
    885
    0
            self.filter = Some(filter);
    886
    0
            self
    887
    0
        }
    888
    889
        /// Add `DISTINCT`
    890
    0
        fn distinct(mut self) -> ExprFuncBuilder {
    891
    0
            self.distinct = true;
    892
    0
            self
    893
    0
        }
    894
    895
        /// Add `RESPECT NULLS` or `IGNORE NULLS`
    896
    0
        fn null_treatment(
    897
    0
            mut self,
    898
    0
            null_treatment: impl Into<Option<NullTreatment>>,
    899
    0
        ) -> ExprFuncBuilder {
    900
    0
            self.null_treatment = null_treatment.into();
    901
    0
            self
    902
    0
        }
    903
    904
    0
        fn partition_by(mut self, partition_by: Vec<Expr>) -> ExprFuncBuilder {
    905
    0
            self.partition_by = Some(partition_by);
    906
    0
            self
    907
    0
        }
    908
    909
    0
        fn window_frame(mut self, window_frame: WindowFrame) -> ExprFuncBuilder {
    910
    0
            self.window_frame = Some(window_frame);
    911
    0
            self
    912
    0
        }
    913
    }
    914
    915
    impl ExprFunctionExt for Expr {
    916
    0
        fn order_by(self, order_by: Vec<Sort>) -> ExprFuncBuilder {
    917
    0
            let mut builder = match self {
    918
    0
                Expr::AggregateFunction(udaf) => {
    919
    0
                    ExprFuncBuilder::new(Some(ExprFuncKind::Aggregate(udaf)))
    920
                }
    921
    0
                Expr::WindowFunction(udwf) => {
    922
    0
                    ExprFuncBuilder::new(Some(ExprFuncKind::Window(udwf)))
    923
                }
    924
    0
                _ => ExprFuncBuilder::new(None),
    925
            };
    926
    0
            if builder.fun.is_some() {
    927
    0
                builder.order_by = Some(order_by);
    928
    0
            }
    929
    0
            builder
    930
    0
        }
    931
    0
        fn filter(self, filter: Expr) -> ExprFuncBuilder {
    932
    0
            match self {
    933
    0
                Expr::AggregateFunction(udaf) => {
    934
    0
                    let mut builder =
    935
    0
                        ExprFuncBuilder::new(Some(ExprFuncKind::Aggregate(udaf)));
    936
    0
                    builder.filter = Some(filter);
    937
    0
                    builder
    938
                }
    939
    0
                _ => ExprFuncBuilder::new(None),
    940
            }
    941
    0
        }
    942
    0
        fn distinct(self) -> ExprFuncBuilder {
    943
    0
            match self {
    944
    0
                Expr::AggregateFunction(udaf) => {
    945
    0
                    let mut builder =
    946
    0
                        ExprFuncBuilder::new(Some(ExprFuncKind::Aggregate(udaf)));
    947
    0
                    builder.distinct = true;
    948
    0
                    builder
    949
                }
    950
    0
                _ => ExprFuncBuilder::new(None),
    951
            }
    952
    0
        }
    953
    0
        fn null_treatment(
    954
    0
            self,
    955
    0
            null_treatment: impl Into<Option<NullTreatment>>,
    956
    0
        ) -> ExprFuncBuilder {
    957
    0
            let mut builder = match self {
    958
    0
                Expr::AggregateFunction(udaf) => {
    959
    0
                    ExprFuncBuilder::new(Some(ExprFuncKind::Aggregate(udaf)))
    960
                }
    961
    0
                Expr::WindowFunction(udwf) => {
    962
    0
                    ExprFuncBuilder::new(Some(ExprFuncKind::Window(udwf)))
    963
                }
    964
    0
                _ => ExprFuncBuilder::new(None),
    965
            };
    966
    0
            if builder.fun.is_some() {
    967
    0
                builder.null_treatment = null_treatment.into();
    968
    0
            }
    969
    0
            builder
    970
    0
        }
    971
    972
    0
        fn partition_by(self, partition_by: Vec<Expr>) -> ExprFuncBuilder {
    973
    0
            match self {
    974
    0
                Expr::WindowFunction(udwf) => {
    975
    0
                    let mut builder = ExprFuncBuilder::new(Some(ExprFuncKind::Window(udwf)));
    976
    0
                    builder.partition_by = Some(partition_by);
    977
    0
                    builder
    978
                }
    979
    0
                _ => ExprFuncBuilder::new(None),
    980
            }
    981
    0
        }
    982
    983
    0
        fn window_frame(self, window_frame: WindowFrame) -> ExprFuncBuilder {
    984
    0
            match self {
    985
    0
                Expr::WindowFunction(udwf) => {
    986
    0
                    let mut builder = ExprFuncBuilder::new(Some(ExprFuncKind::Window(udwf)));
    987
    0
                    builder.window_frame = Some(window_frame);
    988
    0
                    builder
    989
                }
    990
    0
                _ => ExprFuncBuilder::new(None),
    991
            }
    992
    0
        }
    993
    }
    994
    995
    #[cfg(test)]
    996
    mod test {
    997
        use super::*;
    998
    999
        #[test]
    1000
        fn filter_is_null_and_is_not_null() {
    1001
            let col_null = col("col1");
    1002
            let col_not_null = ident("col2");
    1003
            assert_eq!(format!("{}", col_null.is_null()), "col1 IS NULL");
    1004
            assert_eq!(
    1005
                format!("{}", col_not_null.is_not_null()),
    1006
                "col2 IS NOT NULL"
    1007
            );
    1008
        }
    1009
    }