/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 | | } |