## Coverage Report

#### Created: 2024-10-13 08:39

    /Users/andrewlamb/Software/datafusion/datafusion/expr/src/udwf.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
    //! [`WindowUDF`]: User Defined Window Functions
    19
    20
    use arrow::compute::SortOptions;
    21
    use std::cmp::Ordering;
    22
    use std::hash::{DefaultHasher, Hash, Hasher};
    23
    use std::{
    24
        any::Any,
    25
        fmt::{self, Debug, Display, Formatter},
    26
        sync::Arc,
    27
    };
    28
    29
    use arrow::datatypes::{DataType, Field};
    30
    31
    use datafusion_common::{not_impl_err, Result};
    32
    use datafusion_functions_window_common::field::WindowUDFFieldArgs;
    33
    34
    use crate::expr::WindowFunction;
    35
    use crate::{
    36
        function::WindowFunctionSimplification, Documentation, Expr, PartitionEvaluator,
    37
        Signature,
    38
    };
    39
    40
    /// Logical representation of a user-defined window function (UDWF)
    41
    /// A UDWF is different from a UDF in that it is stateful across batches.
    42
    ///
    43
    /// See the documentation on [`PartitionEvaluator`] for more details
    44
    ///
    45
    /// 1. For simple use cases, use [`create_udwf`] (examples in
    46
    ///    [`simple_udwf.rs`]).
    47
    ///
    48
    /// 2. For advanced use cases, use [`WindowUDFImpl`] which provides full API
    49
    ///    access (examples in [`advanced_udwf.rs`]).
    50
    ///
    51
    /// # API Note
    52
    /// This is a separate struct from `WindowUDFImpl` to maintain backwards
    53
    /// compatibility with the older API.
    54
    ///
    55
    /// [`PartitionEvaluator`]: crate::PartitionEvaluator
    56
    /// [`create_udwf`]: crate::expr_fn::create_udwf
    57
    /// [`simple_udwf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/simple_udwf.rs
    58
    /// [`advanced_udwf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/advanced_udwf.rs
    59
    #[derive(Debug, Clone, PartialOrd)]
    60
    pub struct WindowUDF {
    61
        inner: Arc<dyn WindowUDFImpl>,
    62
    }
    63
    64
    /// Defines how the WindowUDF is shown to users
    65
    impl Display for WindowUDF {
    66
    0
        fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    67
    0
            write!(f, "{}", self.name())
    68
    0
        }
    69
    }
    70
    71
    impl PartialEq for WindowUDF {
    72
    0
        fn eq(&self, other: &Self) -> bool {
    73
    0
            self.inner.equals(other.inner.as_ref())
    74
    0
        }
    75
    }
    76
    77
    impl Eq for WindowUDF {}
    78
    79
    impl Hash for WindowUDF {
    80
    0
        fn hash<H: Hasher>(&self, state: &mut H) {
    81
    0
            self.inner.hash_value().hash(state)
    82
    0
        }
    83
    }
    84
    85
    impl WindowUDF {
    86
        /// Create a new `WindowUDF` from a `[WindowUDFImpl]` trait object
    87
        ///
    88
        /// Note this is the same as using the `From` impl (`WindowUDF::from`)
    89
    0
        pub fn new_from_impl<F>(fun: F) -> WindowUDF
    90
    0
        where
    91
    0
            F: WindowUDFImpl + 'static,
    92
    0
        {
    93
    0
            Self {
    94
    0
                inner: Arc::new(fun),
    95
    0
            }
    96
    0
        }
    97
    98
        /// Return the underlying [`WindowUDFImpl`] trait object for this function
    99
    0
        pub fn inner(&self) -> &Arc<dyn WindowUDFImpl> {
    100
    0
            &self.inner
    101
    0
        }
    102
    103
        /// Adds additional names that can be used to invoke this function, in
    104
        /// addition to `name`
    105
        ///
    106
        /// If you implement [`WindowUDFImpl`] directly you should return aliases directly.
    107
    0
        pub fn with_aliases(self, aliases: impl IntoIterator<Item = &'static str>) -> Self {
    108
    0
            Self::new_from_impl(AliasedWindowUDFImpl::new(Arc::clone(&self.inner), aliases))
    109
    0
        }
    110
    111
        /// creates a [`Expr`] that calls the window function with default
    112
        /// values for `order_by`, `partition_by`, `window_frame`.
    113
        ///
    114
        /// See [`ExprFunctionExt`] for details on setting these values.
    115
        ///
    116
        /// This utility allows using a user defined window function without
    117
        /// requiring access to the registry, such as with the DataFrame API.
    118
        ///
    119
        /// [`ExprFunctionExt`]: crate::expr_fn::ExprFunctionExt
    120
    0
        pub fn call(&self, args: Vec<Expr>) -> Expr {
    121
    0
            let fun = crate::WindowFunctionDefinition::WindowUDF(Arc::new(self.clone()));
    122
    0
    123
    0
            Expr::WindowFunction(WindowFunction::new(fun, args))
    124
    0
        }
    125
    126
        /// Returns this function's name
    127
        ///
    128
        /// See [`WindowUDFImpl::name`] for more details.
    129
    0
        pub fn name(&self) -> &str {
    130
    0
            self.inner.name()
    131
    0
        }
    132
    133
        /// Returns the aliases for this function.
    134
    0
        pub fn aliases(&self) -> &[String] {
    135
    0
            self.inner.aliases()
    136
    0
        }
    137
    138
        /// Returns this function's signature (what input types are accepted)
    139
        ///
    140
        /// See [`WindowUDFImpl::signature`] for more details.
    141
    0
        pub fn signature(&self) -> &Signature {
    142
    0
            self.inner.signature()
    143
    0
        }
    144
    145
        /// Do the function rewrite
    146
        ///
    147
        /// See [`WindowUDFImpl::simplify`] for more details.
    148
    0
        pub fn simplify(&self) -> Option<WindowFunctionSimplification> {
    149
    0
            self.inner.simplify()
    150
    0
        }
    151
    152
        /// Return a `PartitionEvaluator` for evaluating this window function
    153
    0
        pub fn partition_evaluator_factory(&self) -> Result<Box<dyn PartitionEvaluator>> {
    154
    0
            self.inner.partition_evaluator()
    155
    0
        }
    156
    157
        /// Returns the field of the final result of evaluating this window function.
    158
        ///
    159
        /// See [`WindowUDFImpl::field`] for more details.
    160
    0
        pub fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
    161
    0
            self.inner.field(field_args)
    162
    0
        }
    163
    164
        /// Returns custom result ordering introduced by this window function
    165
        /// which is used to update ordering equivalences.
    166
        ///
    167
        /// See [`WindowUDFImpl::sort_options`] for more details.
    168
    0
        pub fn sort_options(&self) -> Option<SortOptions> {
    169
    0
            self.inner.sort_options()
    170
    0
        }
    171
    172
        /// See [`WindowUDFImpl::coerce_types`] for more details.
    173
    0
        pub fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
    174
    0
            self.inner.coerce_types(arg_types)
    175
    0
        }
    176
    177
        /// Returns the reversed user-defined window function when the
    178
        /// order of evaluation is reversed.
    179
        ///
    180
        /// See [`WindowUDFImpl::reverse_expr`] for more details.
    181
    0
        pub fn reverse_expr(&self) -> ReversedUDWF {
    182
    0
            self.inner.reverse_expr()
    183
    0
        }
    184
    185
        /// Returns the documentation for this Window UDF.
    186
        ///
    187
        /// Documentation can be accessed programmatically as well as
    188
        /// generating publicly facing documentation.
    189
    0
        pub fn documentation(&self) -> Option<&Documentation> {
    190
    0
            self.inner.documentation()
    191
    0
        }
    192
    }
    193
    194
    impl<F> From<F> for WindowUDF
    195
    where
    196
        F: WindowUDFImpl + Send + Sync + 'static,
    197
    {
    198
    0
        fn from(fun: F) -> Self {
    199
    0
            Self::new_from_impl(fun)
    200
    0
        }
    201
    }
    202
    203
    /// Trait for implementing [`WindowUDF`].
    204
    ///
    205
    /// This trait exposes the full API for implementing user defined window functions and
    206
    /// can be used to implement any function.
    207
    ///
    208
    /// See [`advanced_udwf.rs`] for a full example with complete implementation and
    209
    /// [`WindowUDF`] for other available options.
    210
    ///
    211
    ///
    212
    /// [`advanced_udwf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/advanced_udwf.rs
    213
    /// # Basic Example
    214
    /// ```
    215
    /// # use std::any::Any;
    216
    /// # use std::sync::OnceLock;
    217
    /// # use arrow::datatypes::{DataType, Field};
    218
    /// # use datafusion_common::{DataFusionError, plan_err, Result};
    219
    /// # use datafusion_expr::{col, Signature, Volatility, PartitionEvaluator, WindowFrame, ExprFunctionExt, Documentation};
    220
    /// # use datafusion_expr::{WindowUDFImpl, WindowUDF};
    221
    /// # use datafusion_expr::window_doc_sections::DOC_SECTION_ANALYTICAL;
    222
    /// # use datafusion_functions_window_common::field::WindowUDFFieldArgs;
    223
    ///
    224
    /// #[derive(Debug, Clone)]
    225
    /// struct SmoothIt {
    226
    ///   signature: Signature,
    227
    /// }
    228
    ///
    229
    /// impl SmoothIt {
    230
    ///   fn new() -> Self {
    231
    ///     Self {
    232
    ///       signature: Signature::uniform(1, vec![DataType::Int32], Volatility::Immutable),
    233
    ///      }
    234
    ///   }
    235
    /// }
    236
    ///
    237
    /// static DOCUMENTATION: OnceLock<Documentation> = OnceLock::new();
    238
    ///
    239
    /// fn get_doc() -> &'static Documentation {
    240
    ///     DOCUMENTATION.get_or_init(|| {
    241
    ///         Documentation::builder()
    242
    ///             .with_doc_section(DOC_SECTION_ANALYTICAL)
    243
    ///             .with_description("smooths the windows")
    244
    ///             .with_syntax_example("smooth_it(2)")
    245
    ///             .with_argument("arg1", "The int32 number to smooth by")
    246
    ///             .build()
    247
    ///             .unwrap()
    248
    ///     })
    249
    /// }
    250
    ///
    251
    /// /// Implement the WindowUDFImpl trait for SmoothIt
    252
    /// impl WindowUDFImpl for SmoothIt {
    253
    ///    fn as_any(&self) -> &dyn Any { self }
    254
    ///    fn name(&self) -> &str { "smooth_it" }
    255
    ///    fn signature(&self) -> &Signature { &self.signature }
    256
    ///    // The actual implementation would smooth the window
    257
    ///    fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> { unimplemented!() }
    258
    ///    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
    259
    ///      if let Some(DataType::Int32) = field_args.get_input_type(0) {
    260
    ///        Ok(Field::new(field_args.name(), DataType::Int32, false))
    261
    ///      } else {
    262
    ///        plan_err!("smooth_it only accepts Int32 arguments")
    263
    ///      }
    264
    ///    }
    265
    ///    fn documentation(&self) -> Option<&Documentation> {
    266
    ///      Some(get_doc())
    267
    ///    }
    268
    /// }
    269
    ///
    270
    /// // Create a new WindowUDF from the implementation
    271
    /// let smooth_it = WindowUDF::from(SmoothIt::new());
    272
    ///
    273
    /// // Call the function `add_one(col)`
    274
    /// // smooth_it(speed) OVER (PARTITION BY car ORDER BY time ASC)
    275
    /// let expr = smooth_it.call(vec![col("speed")])
    276
    ///     .partition_by(vec![col("car")])
    277
    ///     .order_by(vec![col("time").sort(true, true)])
    278
    ///     .window_frame(WindowFrame::new(None))
    279
    ///     .build()
    280
    ///     .unwrap();
    281
    /// ```
    282
    pub trait WindowUDFImpl: Debug + Send + Sync {
    283
        // Note: When adding any methods (with default implementations), remember to add them also
    284
        // into the AliasedWindowUDFImpl below!
    285
    286
        /// Returns this object as an [`Any`] trait object
    287
        fn as_any(&self) -> &dyn Any;
    288
    289
        /// Returns this function's name
    290
        fn name(&self) -> &str;
    291
    292
        /// Returns the function's [`Signature`] for information about what input
    293
        /// types are accepted and the function's Volatility.
    294
        fn signature(&self) -> &Signature;
    295
    296
        /// Invoke the function, returning the [`PartitionEvaluator`] instance
    297
        fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>>;
    298
    299
        /// Returns any aliases (alternate names) for this function.
    300
        ///
    301
        /// Note: `aliases` should only include names other than [`Self::name`].
    302
        /// Defaults to `[]` (no aliases)
    303
    0
        fn aliases(&self) -> &[String] {
    304
    0
            &[]
    305
    0
        }
    306
    307
        /// Optionally apply per-UDWF simplification / rewrite rules.
    308
        ///
    309
        /// This can be used to apply function specific simplification rules during
    310
        /// optimization. The default implementation does nothing.
    311
        ///
    312
        /// Note that DataFusion handles simplifying arguments and  "constant
    313
        /// folding" (replacing a function call with constant arguments such as
    314
        /// `my_add(1,2) --> 3` ). Thus, there is no need to implement such
    315
        /// optimizations manually for specific UDFs.
    316
        ///
    317
        /// Example:
    318
        /// [`simplify_udwf_expression.rs`]: <https://github.com/apache/arrow-datafusion/blob/main/datafusion-examples/examples/simplify_udwf_expression.rs>
    319
        ///
    320
        /// # Returns
    321
        /// [None] if simplify is not defined or,
    322
        ///
    323
        /// Or, a closure with two arguments:
    324
        /// * 'window_function': [crate::expr::WindowFunction] for which simplified has been invoked
    325
        /// * 'info': [crate::simplify::SimplifyInfo]
    326
    0
        fn simplify(&self) -> Option<WindowFunctionSimplification> {
    327
    0
            None
    328
    0
        }
    329
    330
        /// Return true if this window UDF is equal to the other.
    331
        ///
    332
        /// Allows customizing the equality of window UDFs.
    333
        /// Must be consistent with [`Self::hash_value`] and follow the same rules as [`Eq`]:
    334
        ///
    335
        /// - reflexive: `a.equals(a)`;
    336
        /// - symmetric: `a.equals(b)` implies `b.equals(a)`;
    337
        /// - transitive: `a.equals(b)` and `b.equals(c)` implies `a.equals(c)`.
    338
        ///
    339
        /// By default, compares [`Self::name`] and [`Self::signature`].
    340
    0
        fn equals(&self, other: &dyn WindowUDFImpl) -> bool {
    341
    0
            self.name() == other.name() && self.signature() == other.signature()
    342
    0
        }
    343
    344
        /// Returns a hash value for this window UDF.
    345
        ///
    346
        /// Allows customizing the hash code of window UDFs. Similarly to [`Hash`] and [`Eq`],
    347
        /// if [`Self::equals`] returns true for two UDFs, their `hash_value`s must be the same.
    348
        ///
    349
        /// By default, hashes [`Self::name`] and [`Self::signature`].
    350
    0
        fn hash_value(&self) -> u64 {
    351
    0
            let hasher = &mut DefaultHasher::new();
    352
    0
            self.name().hash(hasher);
    353
    0
            self.signature().hash(hasher);
    354
    0
            hasher.finish()
    355
    0
        }
    356
    357
        /// The [`Field`] of the final result of evaluating this window function.
    358
        ///
    359
        /// Call `field_args.name()` to get the fully qualified name for defining
    360
        /// the [`Field`]. For a complete example see the implementation in the
    361
        /// [Basic Example](WindowUDFImpl#basic-example) section.
    362
        fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field>;
    363
    364
        /// Allows the window UDF to define a custom result ordering.
    365
        ///
    366
        /// By default, a window UDF doesn't introduce an ordering.
    367
        /// But when specified by a window UDF this is used to update
    368
        /// ordering equivalences.
    369
    0
        fn sort_options(&self) -> Option<SortOptions> {
    370
    0
            None
    371
    0
        }
    372
    373
        /// Coerce arguments of a function call to types that the function can evaluate.
    374
        ///
    375
        /// This function is only called if [`WindowUDFImpl::signature`] returns [`crate::TypeSignature::UserDefined`]. Most
    376
        /// UDWFs should return one of the other variants of `TypeSignature` which handle common
    377
        /// cases
    378
        ///
    379
        /// See the [type coercion module](crate::type_coercion)
    380
        /// documentation for more details on type coercion
    381
        ///
    382
        /// For example, if your function requires a floating point arguments, but the user calls
    383
        /// it like `my_func(1::int)` (aka with `1` as an integer), coerce_types could return `[DataType::Float64]`
    384
        /// to ensure the argument was cast to `1::double`
    385
        ///
    386
        /// # Parameters
    387
        /// * `arg_types`: The argument types of the arguments  this function with
    388
        ///
    389
        /// # Return value
    390
        /// A Vec the same length as `arg_types`. DataFusion will `CAST` the function call
    391
        /// arguments to these specific types.
    392
    0
        fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> {
    393
    0
            not_impl_err!("Function {} does not implement coerce_types", self.name())
    394
    0
        }
    395
    396
        /// Allows customizing the behavior of the user-defined window
    397
        /// function when it is evaluated in reverse order.
    398
    0
        fn reverse_expr(&self) -> ReversedUDWF {
    399
    0
            ReversedUDWF::NotSupported
    400
    0
        }
    401
    402
        /// Returns the documentation for this Window UDF.
    403
        ///
    404
        /// Documentation can be accessed programmatically as well as
    405
        /// generating publicly facing documentation.
    406
    0
        fn documentation(&self) -> Option<&Documentation> {
    407
    0
            None
    408
    0
        }
    409
    }
    410
    411
    pub enum ReversedUDWF {
    412
        /// The result of evaluating the user-defined window function
    413
        /// remains identical when reversed.
    414
        Identical,
    415
        /// A window function which does not support evaluating the result
    416
        /// in reverse order.
    417
        NotSupported,
    418
        /// Customize the user-defined window function for evaluating the
    419
        /// result in reverse order.
    420
        Reversed(Arc<WindowUDF>),
    421
    }
    422
    423
    impl PartialEq for dyn WindowUDFImpl {
    424
    0
        fn eq(&self, other: &Self) -> bool {
    425
    0
            self.equals(other)
    426
    0
        }
    427
    }
    428
    429
    impl PartialOrd for dyn WindowUDFImpl {
    430
    0
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
    431
    0
            match self.name().partial_cmp(other.name()) {
    432
    0
                Some(Ordering::Equal) => self.signature().partial_cmp(other.signature()),
    433
    0
                cmp => cmp,
    434
            }
    435
    0
        }
    436
    }
    437
    438
    /// WindowUDF that adds an alias to the underlying function. It is better to
    439
    /// implement [`WindowUDFImpl`], which supports aliases, directly if possible.
    440
    #[derive(Debug)]
    441
    struct AliasedWindowUDFImpl {
    442
        inner: Arc<dyn WindowUDFImpl>,
    443
        aliases: Vec<String>,
    444
    }
    445
    446
    impl AliasedWindowUDFImpl {
    447
    0
        pub fn new(
    448
    0
            inner: Arc<dyn WindowUDFImpl>,
    449
    0
            new_aliases: impl IntoIterator<Item = &'static str>,
    450
    0
        ) -> Self {
    451
    0
            let mut aliases = inner.aliases().to_vec();
    452
    0
            aliases.extend(new_aliases.into_iter().map(|s| s.to_string()));
    453
    0
    454
    0
            Self { inner, aliases }
    455
    0
        }
    456
    }
    457
    458
    impl WindowUDFImpl for AliasedWindowUDFImpl {
    459
    0
        fn as_any(&self) -> &dyn Any {
    460
    0
            self
    461
    0
        }
    462
    463
    0
        fn name(&self) -> &str {
    464
    0
            self.inner.name()
    465
    0
        }
    466
    467
    0
        fn signature(&self) -> &Signature {
    468
    0
            self.inner.signature()
    469
    0
        }
    470
    471
    0
        fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
    472
    0
            self.inner.partition_evaluator()
    473
    0
        }
    474
    475
    0
        fn aliases(&self) -> &[String] {
    476
    0
            &self.aliases
    477
    0
        }
    478
    479
    0
        fn simplify(&self) -> Option<WindowFunctionSimplification> {
    480
    0
            self.inner.simplify()
    481
    0
        }
    482
    483
    0
        fn equals(&self, other: &dyn WindowUDFImpl) -> bool {
    484
    0
            if let Some(other) = other.as_any().downcast_ref::<AliasedWindowUDFImpl>() {
    485
    0
                self.inner.equals(other.inner.as_ref()) && self.aliases == other.aliases
    486
            } else {
    487
    0
                false
    488
            }
    489
    0
        }
    490
    491
    0
        fn hash_value(&self) -> u64 {
    492
    0
            let hasher = &mut DefaultHasher::new();
    493
    0
            self.inner.hash_value().hash(hasher);
    494
    0
            self.aliases.hash(hasher);
    495
    0
            hasher.finish()
    496
    0
        }
    497
    498
    0
        fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
    499
    0
            self.inner.field(field_args)
    500
    0
        }
    501
    502
    0
        fn sort_options(&self) -> Option<SortOptions> {
    503
    0
            self.inner.sort_options()
    504
    0
        }
    505
    506
    0
        fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
    507
    0
            self.inner.coerce_types(arg_types)
    508
    0
        }
    509
    510
    0
        fn documentation(&self) -> Option<&Documentation> {
    511
    0
            self.inner.documentation()
    512
    0
        }
    513
    }
    514
    515
    // Window UDF doc sections for use in public documentation
    516
    pub mod window_doc_sections {
    517
        use crate::DocSection;
    518
    519
        pub fn doc_sections() -> Vec<DocSection> {
    520
            vec![
    521
                DOC_SECTION_AGGREGATE,
    522
                DOC_SECTION_RANKING,
    523
                DOC_SECTION_ANALYTICAL,
    524
            ]
    525
        }
    526
    527
        pub const DOC_SECTION_AGGREGATE: DocSection = DocSection {
    528
            include: true,
    529
            label: "Aggregate Functions",
    530
            description: Some("All aggregate functions can be used as window functions."),
    531
        };
    532
    533
        pub const DOC_SECTION_RANKING: DocSection = DocSection {
    534
            include: true,
    535
            label: "Ranking Functions",
    536
            description: None,
    537
        };
    538
    539
        pub const DOC_SECTION_ANALYTICAL: DocSection = DocSection {
    540
            include: true,
    541
            label: "Analytical Functions",
    542
            description: None,
    543
        };
    544
    }
    545
    546
    #[cfg(test)]
    547
    mod test {
    548
        use crate::{PartitionEvaluator, WindowUDF, WindowUDFImpl};
    549
        use arrow::datatypes::{DataType, Field};
    550
        use datafusion_common::Result;
    551
        use datafusion_expr_common::signature::{Signature, Volatility};
    552
        use datafusion_functions_window_common::field::WindowUDFFieldArgs;
    553
        use std::any::Any;
    554
        use std::cmp::Ordering;
    555
    556
        #[derive(Debug, Clone)]
    557
        struct AWindowUDF {
    558
            signature: Signature,
    559
        }
    560
    561
        impl AWindowUDF {
    562
            fn new() -> Self {
    563
                Self {
    564
                    signature: Signature::uniform(
    565
                        1,
    566
                        vec![DataType::Int32],
    567
                        Volatility::Immutable,
    568
                    ),
    569
                }
    570
            }
    571
        }
    572
    573
        /// Implement the WindowUDFImpl trait for AddOne
    574
        impl WindowUDFImpl for AWindowUDF {
    575
            fn as_any(&self) -> &dyn Any {
    576
                self
    577
            }
    578
            fn name(&self) -> &str {
    579
                "a"
    580
            }
    581
            fn signature(&self) -> &Signature {
    582
                &self.signature
    583
            }
    584
            fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
    585
                unimplemented!()
    586
            }
    587
            fn field(&self, _field_args: WindowUDFFieldArgs) -> Result<Field> {
    588
                unimplemented!()
    589
            }
    590
        }
    591
    592
        #[derive(Debug, Clone)]
    593
        struct BWindowUDF {
    594
            signature: Signature,
    595
        }
    596
    597
        impl BWindowUDF {
    598
            fn new() -> Self {
    599
                Self {
    600
                    signature: Signature::uniform(
    601
                        1,
    602
                        vec![DataType::Int32],
    603
                        Volatility::Immutable,
    604
                    ),
    605
                }
    606
            }
    607
        }
    608
    609
        /// Implement the WindowUDFImpl trait for AddOne
    610
        impl WindowUDFImpl for BWindowUDF {
    611
            fn as_any(&self) -> &dyn Any {
    612
                self
    613
            }
    614
            fn name(&self) -> &str {
    615
                "b"
    616
            }
    617
            fn signature(&self) -> &Signature {
    618
                &self.signature
    619
            }
    620
            fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
    621
                unimplemented!()
    622
            }
    623
            fn field(&self, _field_args: WindowUDFFieldArgs) -> Result<Field> {
    624
                unimplemented!()
    625
            }
    626
        }
    627
    628
        #[test]
    629
        fn test_partial_ord() {
    630
            let a1 = WindowUDF::from(AWindowUDF::new());
    631
            let a2 = WindowUDF::from(AWindowUDF::new());
    632
            assert_eq!(a1.partial_cmp(&a2), Some(Ordering::Equal));
    633
    634
            let b1 = WindowUDF::from(BWindowUDF::new());
    635
            assert!(a1 < b1);
    636
            assert!(!(a1 == b1));
    637
        }
    638
    }