## Coverage Report

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

    /Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/window/ntile.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
    //! Defines physical expression for `ntile` that can evaluated
    19
    //! at runtime during query execution
    20
    21
    use crate::expressions::Column;
    22
    use crate::window::BuiltInWindowFunctionExpr;
    23
    use crate::{PhysicalExpr, PhysicalSortExpr};
    24
    25
    use arrow::array::{ArrayRef, UInt64Array};
    26
    use arrow::datatypes::Field;
    27
    use arrow_schema::{DataType, SchemaRef, SortOptions};
    28
    use datafusion_common::Result;
    29
    use datafusion_expr::PartitionEvaluator;
    30
    31
    use std::any::Any;
    32
    use std::sync::Arc;
    33
    34
    #[derive(Debug)]
    35
    pub struct Ntile {
    36
        name: String,
    37
        n: u64,
    38
        /// Output data type
    39
        data_type: DataType,
    40
    }
    41
    42
    impl Ntile {
    43
    0
        pub fn new(name: String, n: u64, data_type: &DataType) -> Self {
    44
    0
            Self {
    45
    0
                name,
    46
    0
                n,
    47
    0
                data_type: data_type.clone(),
    48
    0
            }
    49
    0
        }
    50
    51
    0
        pub fn get_n(&self) -> u64 {
    52
    0
            self.n
    53
    0
        }
    54
    }
    55
    56
    impl BuiltInWindowFunctionExpr for Ntile {
    57
    0
        fn as_any(&self) -> &dyn Any {
    58
    0
            self
    59
    0
        }
    60
    61
    0
        fn field(&self) -> Result<Field> {
    62
    0
            let nullable = false;
    63
    0
            Ok(Field::new(self.name(), self.data_type.clone(), nullable))
    64
    0
        }
    65
    66
    0
        fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
    67
    0
            vec![]
    68
    0
        }
    69
    70
    0
        fn name(&self) -> &str {
    71
    0
            &self.name
    72
    0
        }
    73
    74
    0
        fn create_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
    75
    0
            Ok(Box::new(NtileEvaluator { n: self.n }))
    76
    0
        }
    77
    78
    0
        fn get_result_ordering(&self, schema: &SchemaRef) -> Option<PhysicalSortExpr> {
    79
    0
            // The built-in NTILE window function introduces a new ordering:
    80
    0
            schema.column_with_name(self.name()).map(|(idx, field)| {
    81
    0
                let expr = Arc::new(Column::new(field.name(), idx));
    82
    0
                let options = SortOptions {
    83
    0
                    descending: false,
    84
    0
                    nulls_first: false,
    85
    0
                }; // ASC, NULLS LAST
    86
    0
                PhysicalSortExpr { expr, options }
    87
    0
            })
    88
    0
        }
    89
    }
    90
    91
    #[derive(Debug)]
    92
    pub(crate) struct NtileEvaluator {
    93
        n: u64,
    94
    }
    95
    96
    impl PartitionEvaluator for NtileEvaluator {
    97
    0
        fn evaluate_all(
    98
    0
            &mut self,
    99
    0
            _values: &[ArrayRef],
    100
    0
            num_rows: usize,
    101
    0
        ) -> Result<ArrayRef> {
    102
    0
            let num_rows = num_rows as u64;
    103
    0
            let mut vec: Vec<u64> = Vec::new();
    104
    0
            let n = u64::min(self.n, num_rows);
    105
    0
            for i in 0..num_rows {
    106
    0
                let res = i * n / num_rows;
    107
    0
                vec.push(res + 1)
    108
            }
    109
    0
            Ok(Arc::new(UInt64Array::from(vec)))
    110
    0
        }
    111
    }