|
17 | 17 |
|
18 | 18 | use std::sync::Arc; |
19 | 19 |
|
20 | | -use arrow::array::builder::PrimitiveBuilder; |
21 | 20 | use arrow::array::cast::AsArray; |
22 | 21 | use arrow::array::types::{Date32Type, Int32Type}; |
23 | 22 | use arrow::array::{Array, PrimitiveArray}; |
| 23 | +use arrow::buffer::NullBuffer; |
24 | 24 | use arrow::datatypes::DataType; |
25 | 25 | use arrow::datatypes::DataType::Date32; |
26 | 26 | use chrono::prelude::*; |
@@ -139,24 +139,27 @@ impl ScalarUDFImpl for MakeDateFunc { |
139 | 139 | let months = months.as_primitive::<Int32Type>(); |
140 | 140 | let days = days.as_primitive::<Int32Type>(); |
141 | 141 |
|
142 | | - let mut builder: PrimitiveBuilder<Date32Type> = |
143 | | - PrimitiveArray::builder(len); |
| 142 | + let nulls = |
| 143 | + NullBuffer::union_many([years.nulls(), months.nulls(), days.nulls()]); |
144 | 144 |
|
| 145 | + let mut values = Vec::with_capacity(len); |
145 | 146 | for i in 0..len { |
146 | 147 | // match postgresql behaviour which returns null for any null input |
147 | | - if years.is_null(i) || months.is_null(i) || days.is_null(i) { |
148 | | - builder.append_null(); |
| 148 | + if nulls.as_ref().is_some_and(|n| n.is_null(i)) { |
| 149 | + values.push(0); |
149 | 150 | } else { |
150 | 151 | make_date_inner( |
151 | 152 | years.value(i), |
152 | 153 | months.value(i), |
153 | 154 | days.value(i), |
154 | | - |days: i32| builder.append_value(days), |
| 155 | + |days: i32| values.push(days), |
155 | 156 | )?; |
156 | 157 | } |
157 | 158 | } |
158 | 159 |
|
159 | | - Ok(ColumnarValue::Array(Arc::new(builder.finish()))) |
| 160 | + Ok(ColumnarValue::Array(Arc::new( |
| 161 | + PrimitiveArray::<Date32Type>::new(values.into(), nulls), |
| 162 | + ))) |
160 | 163 | } |
161 | 164 | } |
162 | 165 | } |
@@ -197,3 +200,88 @@ fn make_date_inner<F: FnMut(i32)>( |
197 | 200 | exec_err!("Unable to parse date from {year}, {month}, {day}") |
198 | 201 | } |
199 | 202 | } |
| 203 | + |
| 204 | +#[cfg(test)] |
| 205 | +mod tests { |
| 206 | + use super::*; |
| 207 | + use arrow::array::Int32Array; |
| 208 | + use arrow::datatypes::Field; |
| 209 | + use datafusion_common::config::ConfigOptions; |
| 210 | + |
| 211 | + fn invoke(args: Vec<ColumnarValue>, number_rows: usize) -> Result<ColumnarValue> { |
| 212 | + let arg_fields = args |
| 213 | + .iter() |
| 214 | + .map(|a| Field::new("a", a.data_type(), true).into()) |
| 215 | + .collect::<Vec<_>>(); |
| 216 | + MakeDateFunc::new().invoke_with_args(ScalarFunctionArgs { |
| 217 | + args, |
| 218 | + arg_fields, |
| 219 | + number_rows, |
| 220 | + return_field: Field::new("f", Date32, true).into(), |
| 221 | + config_options: Arc::new(ConfigOptions::default()), |
| 222 | + }) |
| 223 | + } |
| 224 | + |
| 225 | + #[test] |
| 226 | + fn test_make_date_array() { |
| 227 | + let years = ColumnarValue::Array(Arc::new(Int32Array::from(vec![ |
| 228 | + Some(1970), |
| 229 | + Some(1970), |
| 230 | + ]))); |
| 231 | + let months = |
| 232 | + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(1)]))); |
| 233 | + let days = |
| 234 | + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(2)]))); |
| 235 | + |
| 236 | + let ColumnarValue::Array(arr) = invoke(vec![years, months, days], 2).unwrap() |
| 237 | + else { |
| 238 | + panic!("expected array result"); |
| 239 | + }; |
| 240 | + let arr = arr.as_primitive::<Date32Type>(); |
| 241 | + // Days since the unix epoch. |
| 242 | + assert_eq!(arr.value(0), 0); |
| 243 | + assert_eq!(arr.value(1), 1); |
| 244 | + } |
| 245 | + |
| 246 | + #[test] |
| 247 | + fn test_make_date_null_propagation() { |
| 248 | + // A NULL in any component column yields a NULL row. |
| 249 | + let years = |
| 250 | + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(2000), None]))); |
| 251 | + let months = |
| 252 | + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(6), Some(6)]))); |
| 253 | + let days = ColumnarValue::Array(Arc::new(Int32Array::from(vec![None, Some(15)]))); |
| 254 | + |
| 255 | + let ColumnarValue::Array(arr) = invoke(vec![years, months, days], 2).unwrap() |
| 256 | + else { |
| 257 | + panic!("expected array result"); |
| 258 | + }; |
| 259 | + let arr = arr.as_primitive::<Date32Type>(); |
| 260 | + assert!(arr.is_null(0)); |
| 261 | + assert!(arr.is_null(1)); |
| 262 | + } |
| 263 | + |
| 264 | + #[test] |
| 265 | + fn test_make_date_scalar_array_mix() { |
| 266 | + let year = ColumnarValue::Scalar(ScalarValue::Int32(Some(1970))); |
| 267 | + let month = ColumnarValue::Scalar(ScalarValue::Int32(Some(1))); |
| 268 | + let days = |
| 269 | + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(3)]))); |
| 270 | + |
| 271 | + let ColumnarValue::Array(arr) = invoke(vec![year, month, days], 2).unwrap() |
| 272 | + else { |
| 273 | + panic!("expected array result"); |
| 274 | + }; |
| 275 | + let arr = arr.as_primitive::<Date32Type>(); |
| 276 | + assert_eq!(arr.value(0), 0); |
| 277 | + assert_eq!(arr.value(1), 2); |
| 278 | + } |
| 279 | + |
| 280 | + #[test] |
| 281 | + fn test_make_date_out_of_range_errors() { |
| 282 | + let years = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(2000)]))); |
| 283 | + let months = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(13)]))); |
| 284 | + let days = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1)]))); |
| 285 | + assert!(invoke(vec![years, months, days], 1).is_err()); |
| 286 | + } |
| 287 | +} |
0 commit comments