datafusion_functions_window/
row_number.rs

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//! `row_number` window function implementation
19
20use datafusion_common::arrow::array::ArrayRef;
21use datafusion_common::arrow::array::UInt64Array;
22use datafusion_common::arrow::compute::SortOptions;
23use datafusion_common::arrow::datatypes::DataType;
24use datafusion_common::arrow::datatypes::Field;
25use datafusion_common::{Result, ScalarValue};
26use datafusion_expr::{
27    Documentation, PartitionEvaluator, Signature, Volatility, WindowUDFImpl,
28};
29use datafusion_functions_window_common::field;
30use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
31use datafusion_macros::user_doc;
32use field::WindowUDFFieldArgs;
33use std::any::Any;
34use std::fmt::Debug;
35use std::ops::Range;
36
37define_udwf_and_expr!(
38    RowNumber,
39    row_number,
40    "Returns a unique row number for each row in window partition beginning at 1."
41);
42
43/// row_number expression
44#[user_doc(
45    doc_section(label = "Ranking Functions"),
46    description = "Number of the current row within its partition, counting from 1.",
47    syntax_example = "row_number()"
48)]
49#[derive(Debug)]
50pub struct RowNumber {
51    signature: Signature,
52}
53
54impl RowNumber {
55    /// Create a new `row_number` function
56    pub fn new() -> Self {
57        Self {
58            signature: Signature::nullary(Volatility::Immutable),
59        }
60    }
61}
62
63impl Default for RowNumber {
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69impl WindowUDFImpl for RowNumber {
70    fn as_any(&self) -> &dyn Any {
71        self
72    }
73
74    fn name(&self) -> &str {
75        "row_number"
76    }
77
78    fn signature(&self) -> &Signature {
79        &self.signature
80    }
81
82    fn partition_evaluator(
83        &self,
84        _partition_evaluator_args: PartitionEvaluatorArgs,
85    ) -> Result<Box<dyn PartitionEvaluator>> {
86        Ok(Box::<NumRowsEvaluator>::default())
87    }
88
89    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
90        Ok(Field::new(field_args.name(), DataType::UInt64, false))
91    }
92
93    fn sort_options(&self) -> Option<SortOptions> {
94        Some(SortOptions {
95            descending: false,
96            nulls_first: false,
97        })
98    }
99
100    fn documentation(&self) -> Option<&Documentation> {
101        self.doc()
102    }
103}
104
105/// State for the `row_number` built-in window function.
106#[derive(Debug, Default)]
107struct NumRowsEvaluator {
108    n_rows: usize,
109}
110
111impl PartitionEvaluator for NumRowsEvaluator {
112    fn is_causal(&self) -> bool {
113        // The row_number function doesn't need "future" values to emit results:
114        true
115    }
116
117    fn evaluate_all(
118        &mut self,
119        _values: &[ArrayRef],
120        num_rows: usize,
121    ) -> Result<ArrayRef> {
122        Ok(std::sync::Arc::new(UInt64Array::from_iter_values(
123            1..(num_rows as u64) + 1,
124        )))
125    }
126
127    fn evaluate(
128        &mut self,
129        _values: &[ArrayRef],
130        _range: &Range<usize>,
131    ) -> Result<ScalarValue> {
132        self.n_rows += 1;
133        Ok(ScalarValue::UInt64(Some(self.n_rows as u64)))
134    }
135
136    fn supports_bounded_execution(&self) -> bool {
137        true
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use std::sync::Arc;
144
145    use datafusion_common::arrow::array::{Array, BooleanArray};
146    use datafusion_common::cast::as_uint64_array;
147
148    use super::*;
149
150    #[test]
151    fn row_number_all_null() -> Result<()> {
152        let values: ArrayRef = Arc::new(BooleanArray::from(vec![
153            None, None, None, None, None, None, None, None,
154        ]));
155        let num_rows = values.len();
156
157        let actual = RowNumber::default()
158            .partition_evaluator(PartitionEvaluatorArgs::default())?
159            .evaluate_all(&[values], num_rows)?;
160        let actual = as_uint64_array(&actual)?;
161
162        assert_eq!(vec![1, 2, 3, 4, 5, 6, 7, 8], *actual.values());
163        Ok(())
164    }
165
166    #[test]
167    fn row_number_all_values() -> Result<()> {
168        let values: ArrayRef = Arc::new(BooleanArray::from(vec![
169            true, false, true, false, false, true, false, true,
170        ]));
171        let num_rows = values.len();
172
173        let actual = RowNumber::default()
174            .partition_evaluator(PartitionEvaluatorArgs::default())?
175            .evaluate_all(&[values], num_rows)?;
176        let actual = as_uint64_array(&actual)?;
177
178        assert_eq!(vec![1, 2, 3, 4, 5, 6, 7, 8], *actual.values());
179        Ok(())
180    }
181}
datafusion_functions_window/row_number.rs

datafusion_functions_window/
row_number.rs
