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
}