|
| 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 `row_number` that can evaluated at runtime during query execution |
| 19 | +
|
| 20 | +use crate::error::Result; |
| 21 | +use crate::physical_plan::{ |
| 22 | + window_functions::BuiltInWindowFunctionExpr, PhysicalExpr, WindowAccumulator, |
| 23 | +}; |
| 24 | +use crate::scalar::ScalarValue; |
| 25 | +use arrow::array::{ArrayRef, UInt64Array}; |
| 26 | +use arrow::datatypes::{DataType, Field}; |
| 27 | +use std::any::Any; |
| 28 | +use std::sync::Arc; |
| 29 | + |
| 30 | +/// row_number expression |
| 31 | +#[derive(Debug)] |
| 32 | +pub struct RowNumber { |
| 33 | + name: String, |
| 34 | +} |
| 35 | + |
| 36 | +impl RowNumber { |
| 37 | + /// Create a new ROW_NUMBER function |
| 38 | + pub fn new(name: String) -> Self { |
| 39 | + Self { name } |
| 40 | + } |
| 41 | +} |
| 42 | + |
| 43 | +impl BuiltInWindowFunctionExpr for RowNumber { |
| 44 | + /// Return a reference to Any that can be used for downcasting |
| 45 | + fn as_any(&self) -> &dyn Any { |
| 46 | + self |
| 47 | + } |
| 48 | + |
| 49 | + fn field(&self) -> Result<Field> { |
| 50 | + let nullable = false; |
| 51 | + let data_type = DataType::UInt64; |
| 52 | + Ok(Field::new(&self.name(), data_type, nullable)) |
| 53 | + } |
| 54 | + |
| 55 | + fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> { |
| 56 | + vec![] |
| 57 | + } |
| 58 | + |
| 59 | + fn name(&self) -> &str { |
| 60 | + self.name.as_str() |
| 61 | + } |
| 62 | + |
| 63 | + fn create_accumulator(&self) -> Result<Box<dyn WindowAccumulator>> { |
| 64 | + Ok(Box::new(RowNumberAccumulator::new())) |
| 65 | + } |
| 66 | +} |
| 67 | + |
| 68 | +#[derive(Debug)] |
| 69 | +struct RowNumberAccumulator { |
| 70 | + row_number: u64, |
| 71 | +} |
| 72 | + |
| 73 | +impl RowNumberAccumulator { |
| 74 | + /// new row_number accumulator |
| 75 | + pub fn new() -> Self { |
| 76 | + // row number is 1 based |
| 77 | + Self { row_number: 1 } |
| 78 | + } |
| 79 | +} |
| 80 | + |
| 81 | +impl WindowAccumulator for RowNumberAccumulator { |
| 82 | + fn scan(&mut self, _values: &[ScalarValue]) -> Result<Option<ScalarValue>> { |
| 83 | + let result = Some(ScalarValue::UInt64(Some(self.row_number))); |
| 84 | + self.row_number += 1; |
| 85 | + Ok(result) |
| 86 | + } |
| 87 | + |
| 88 | + fn scan_batch( |
| 89 | + &mut self, |
| 90 | + num_rows: usize, |
| 91 | + _values: &[ArrayRef], |
| 92 | + ) -> Result<Option<ArrayRef>> { |
| 93 | + let new_row_number = self.row_number + (num_rows as u64); |
| 94 | + // TODO: probably would be nice to have a (optimized) kernel for this at some point to |
| 95 | + // generate an array like this. |
| 96 | + let result = UInt64Array::from_iter_values(self.row_number..new_row_number); |
| 97 | + self.row_number = new_row_number; |
| 98 | + Ok(Some(Arc::new(result))) |
| 99 | + } |
| 100 | + |
| 101 | + fn evaluate(&self) -> Result<Option<ScalarValue>> { |
| 102 | + Ok(None) |
| 103 | + } |
| 104 | +} |
| 105 | + |
| 106 | +#[cfg(test)] |
| 107 | +mod tests { |
| 108 | + use super::*; |
| 109 | + use crate::error::Result; |
| 110 | + use arrow::record_batch::RecordBatch; |
| 111 | + use arrow::{array::*, datatypes::*}; |
| 112 | + |
| 113 | + #[test] |
| 114 | + fn row_number_all_null() -> Result<()> { |
| 115 | + let arr: ArrayRef = Arc::new(BooleanArray::from(vec![ |
| 116 | + None, None, None, None, None, None, None, None, |
| 117 | + ])); |
| 118 | + let schema = Schema::new(vec![Field::new("arr", DataType::Boolean, false)]); |
| 119 | + let batch = RecordBatch::try_new(Arc::new(schema), vec![arr])?; |
| 120 | + |
| 121 | + let row_number = Arc::new(RowNumber::new("row_number".to_owned())); |
| 122 | + |
| 123 | + let mut acc = row_number.create_accumulator()?; |
| 124 | + let expr = row_number.expressions(); |
| 125 | + let values = expr |
| 126 | + .iter() |
| 127 | + .map(|e| e.evaluate(&batch)) |
| 128 | + .map(|r| r.map(|v| v.into_array(batch.num_rows()))) |
| 129 | + .collect::<Result<Vec<_>>>()?; |
| 130 | + |
| 131 | + let result = acc.scan_batch(batch.num_rows(), &values)?; |
| 132 | + assert_eq!(true, result.is_some()); |
| 133 | + |
| 134 | + let result = result.unwrap(); |
| 135 | + let result = result.as_any().downcast_ref::<UInt64Array>().unwrap(); |
| 136 | + let result = result.values(); |
| 137 | + assert_eq!(vec![1, 2, 3, 4, 5, 6, 7, 8], result); |
| 138 | + |
| 139 | + let result = acc.evaluate()?; |
| 140 | + assert_eq!(false, result.is_some()); |
| 141 | + Ok(()) |
| 142 | + } |
| 143 | + |
| 144 | + #[test] |
| 145 | + fn row_number_all_values() -> Result<()> { |
| 146 | + let arr: ArrayRef = Arc::new(BooleanArray::from(vec![ |
| 147 | + true, false, true, false, false, true, false, true, |
| 148 | + ])); |
| 149 | + let schema = Schema::new(vec![Field::new("arr", DataType::Boolean, false)]); |
| 150 | + let batch = RecordBatch::try_new(Arc::new(schema), vec![arr])?; |
| 151 | + |
| 152 | + let row_number = Arc::new(RowNumber::new("row_number".to_owned())); |
| 153 | + |
| 154 | + let mut acc = row_number.create_accumulator()?; |
| 155 | + let expr = row_number.expressions(); |
| 156 | + let values = expr |
| 157 | + .iter() |
| 158 | + .map(|e| e.evaluate(&batch)) |
| 159 | + .map(|r| r.map(|v| v.into_array(batch.num_rows()))) |
| 160 | + .collect::<Result<Vec<_>>>()?; |
| 161 | + |
| 162 | + let result = acc.scan_batch(batch.num_rows(), &values)?; |
| 163 | + assert_eq!(true, result.is_some()); |
| 164 | + |
| 165 | + let result = result.unwrap(); |
| 166 | + let result = result.as_any().downcast_ref::<UInt64Array>().unwrap(); |
| 167 | + let result = result.values(); |
| 168 | + assert_eq!(vec![1, 2, 3, 4, 5, 6, 7, 8], result); |
| 169 | + |
| 170 | + let result = acc.evaluate()?; |
| 171 | + assert_eq!(false, result.is_some()); |
| 172 | + Ok(()) |
| 173 | + } |
| 174 | +} |
0 commit comments