diff --git a/datafusion/functions/src/datetime/make_date.rs b/datafusion/functions/src/datetime/make_date.rs index b8010f64d270a..8f892aafb4e1c 100644 --- a/datafusion/functions/src/datetime/make_date.rs +++ b/datafusion/functions/src/datetime/make_date.rs @@ -80,47 +80,10 @@ impl ScalarUDFImpl for MakeDateFunc { ColumnarValue::Array(a) => Some(a.len()), }); - let is_scalar = len.is_none(); - let array_size = if is_scalar { 1 } else { len.unwrap() }; - let years = args[0].cast_to(&DataType::Int32, None)?; let months = args[1].cast_to(&DataType::Int32, None)?; let days = args[2].cast_to(&DataType::Int32, None)?; - // since the epoch for the date32 datatype is the unix epoch - // we need to subtract the unix epoch from the current date - // note this can result in a negative value - let unix_days_from_ce = NaiveDate::from_ymd_opt(1970, 1, 1) - .unwrap() - .num_days_from_ce(); - - let mut builder: PrimitiveBuilder = - PrimitiveArray::builder(array_size); - - let construct_date_fn = |builder: &mut PrimitiveBuilder, - year: i32, - month: i32, - day: i32, - unix_days_from_ce: i32| - -> Result<()> { - let Ok(m) = u32::try_from(month) else { - return exec_err!("Month value '{month:?}' is out of range"); - }; - let Ok(d) = u32::try_from(day) else { - return exec_err!("Day value '{day:?}' is out of range"); - }; - - let date = NaiveDate::from_ymd_opt(year, m, d); - - match date { - Some(d) => builder.append_value(d.num_days_from_ce() - unix_days_from_ce), - None => { - return exec_err!("Unable to parse date from {year}, {month}, {day}") - } - }; - Ok(()) - }; - let scalar_value_fn = |col: &ColumnarValue| -> Result { let ColumnarValue::Scalar(s) = col else { return exec_err!("Expected scalar value"); @@ -131,54 +94,83 @@ impl ScalarUDFImpl for MakeDateFunc { Ok(*i) }; - // For scalar only columns the operation is faster without using the PrimitiveArray - if is_scalar { - construct_date_fn( - &mut builder, - scalar_value_fn(&years)?, - scalar_value_fn(&months)?, - scalar_value_fn(&days)?, - unix_days_from_ce, - )?; - } else { - let to_primitive_array = |col: &ColumnarValue, - scalar_count: usize| - -> Result> { - match col { - ColumnarValue::Array(a) => { - Ok(a.as_primitive::().to_owned()) + let value = if let Some(array_size) = len { + let to_primitive_array_fn = + |col: &ColumnarValue| -> PrimitiveArray { + match col { + ColumnarValue::Array(a) => { + a.as_primitive::().to_owned() + } + _ => { + let v = scalar_value_fn(col).unwrap(); + PrimitiveArray::::from_value(v, array_size) + } } - _ => { - let v = scalar_value_fn(col).unwrap(); - Ok(PrimitiveArray::::from_value(v, scalar_count)) - } - } - }; + }; - let years = to_primitive_array(&years, array_size).unwrap(); - let months = to_primitive_array(&months, array_size).unwrap(); - let days = to_primitive_array(&days, array_size).unwrap(); + let years = to_primitive_array_fn(&years); + let months = to_primitive_array_fn(&months); + let days = to_primitive_array_fn(&days); + + let mut builder: PrimitiveBuilder = + PrimitiveArray::builder(array_size); for i in 0..array_size { - construct_date_fn( - &mut builder, + make_date_inner( years.value(i), months.value(i), days.value(i), - unix_days_from_ce, + |days: i32| builder.append_value(days), )?; } - } - let arr = builder.finish(); + let arr = builder.finish(); - if is_scalar { - // If all inputs are scalar, keeps output as scalar - Ok(ColumnarValue::Scalar(ScalarValue::Date32(Some( - arr.value(0), - )))) + ColumnarValue::Array(Arc::new(arr)) } else { - Ok(ColumnarValue::Array(Arc::new(arr))) - } + // For scalar only columns the operation is faster without using the PrimitiveArray. + // Also, keep the output as scalar since all inputs are scalar. + let mut value = 0; + make_date_inner( + scalar_value_fn(&years)?, + scalar_value_fn(&months)?, + scalar_value_fn(&days)?, + |days: i32| value = days, + )?; + + ColumnarValue::Scalar(ScalarValue::Date32(Some(value))) + }; + + Ok(value) + } +} + +/// Converts the year/month/day fields to an `i32` representing the days from +/// the unix epoch and invokes `date_consumer_fn` with the value +fn make_date_inner( + year: i32, + month: i32, + day: i32, + mut date_consumer_fn: F, +) -> Result<()> { + let Ok(m) = u32::try_from(month) else { + return exec_err!("Month value '{month:?}' is out of range"); + }; + let Ok(d) = u32::try_from(day) else { + return exec_err!("Day value '{day:?}' is out of range"); + }; + + if let Some(date) = NaiveDate::from_ymd_opt(year, m, d) { + // The number of days until the start of the unix epoch in the proleptic Gregorian calendar + // (with January 1, Year 1 (CE) as day 1). See [Datelike::num_days_from_ce]. + const UNIX_DAYS_FROM_CE: i32 = 719_163; + + // since the epoch for the date32 datatype is the unix epoch + // we need to subtract the unix epoch from the current date + // note that this can result in a negative value + date_consumer_fn(date.num_days_from_ce() - UNIX_DAYS_FROM_CE); + Ok(()) + } else { + exec_err!("Unable to parse date from {year}, {month}, {day}") } }