diff --git a/datafusion/functions/src/datetime/make_date.rs b/datafusion/functions/src/datetime/make_date.rs index dc1328742f24e..3a6b76ed86eb0 100644 --- a/datafusion/functions/src/datetime/make_date.rs +++ b/datafusion/functions/src/datetime/make_date.rs @@ -17,10 +17,10 @@ use std::sync::Arc; -use arrow::array::builder::PrimitiveBuilder; use arrow::array::cast::AsArray; use arrow::array::types::{Date32Type, Int32Type}; use arrow::array::{Array, PrimitiveArray}; +use arrow::buffer::NullBuffer; use arrow::datatypes::DataType; use arrow::datatypes::DataType::Date32; use chrono::prelude::*; @@ -139,24 +139,27 @@ impl ScalarUDFImpl for MakeDateFunc { let months = months.as_primitive::(); let days = days.as_primitive::(); - let mut builder: PrimitiveBuilder = - PrimitiveArray::builder(len); + let nulls = + NullBuffer::union_many([years.nulls(), months.nulls(), days.nulls()]); + let mut values = Vec::with_capacity(len); for i in 0..len { // match postgresql behaviour which returns null for any null input - if years.is_null(i) || months.is_null(i) || days.is_null(i) { - builder.append_null(); + if nulls.as_ref().is_some_and(|n| n.is_null(i)) { + values.push(0); } else { make_date_inner( years.value(i), months.value(i), days.value(i), - |days: i32| builder.append_value(days), + |days: i32| values.push(days), )?; } } - Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + Ok(ColumnarValue::Array(Arc::new( + PrimitiveArray::::new(values.into(), nulls), + ))) } } } @@ -197,3 +200,88 @@ fn make_date_inner( exec_err!("Unable to parse date from {year}, {month}, {day}") } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + use arrow::datatypes::Field; + use datafusion_common::config::ConfigOptions; + + fn invoke(args: Vec, number_rows: usize) -> Result { + let arg_fields = args + .iter() + .map(|a| Field::new("a", a.data_type(), true).into()) + .collect::>(); + MakeDateFunc::new().invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Field::new("f", Date32, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }) + } + + #[test] + fn test_make_date_array() { + let years = ColumnarValue::Array(Arc::new(Int32Array::from(vec![ + Some(1970), + Some(1970), + ]))); + let months = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(1)]))); + let days = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(2)]))); + + let ColumnarValue::Array(arr) = invoke(vec![years, months, days], 2).unwrap() + else { + panic!("expected array result"); + }; + let arr = arr.as_primitive::(); + // Days since the unix epoch. + assert_eq!(arr.value(0), 0); + assert_eq!(arr.value(1), 1); + } + + #[test] + fn test_make_date_null_propagation() { + // A NULL in any component column yields a NULL row. + let years = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(2000), None]))); + let months = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(6), Some(6)]))); + let days = ColumnarValue::Array(Arc::new(Int32Array::from(vec![None, Some(15)]))); + + let ColumnarValue::Array(arr) = invoke(vec![years, months, days], 2).unwrap() + else { + panic!("expected array result"); + }; + let arr = arr.as_primitive::(); + assert!(arr.is_null(0)); + assert!(arr.is_null(1)); + } + + #[test] + fn test_make_date_scalar_array_mix() { + let year = ColumnarValue::Scalar(ScalarValue::Int32(Some(1970))); + let month = ColumnarValue::Scalar(ScalarValue::Int32(Some(1))); + let days = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(3)]))); + + let ColumnarValue::Array(arr) = invoke(vec![year, month, days], 2).unwrap() + else { + panic!("expected array result"); + }; + let arr = arr.as_primitive::(); + assert_eq!(arr.value(0), 0); + assert_eq!(arr.value(1), 2); + } + + #[test] + fn test_make_date_out_of_range_errors() { + let years = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(2000)]))); + let months = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(13)]))); + let days = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1)]))); + assert!(invoke(vec![years, months, days], 1).is_err()); + } +}