alamb commented on code in PR #9600:
URL: https://github.com/apache/arrow-datafusion/pull/9600#discussion_r1526112075
##########
datafusion/physical-expr/src/datetime_expressions.rs:
##########
@@ -351,51 +317,77 @@ pub fn make_date(args: &[ColumnarValue]) ->
Result<ColumnarValue> {
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<PrimitiveArray<Int32Type>> {
+ let value = if let Some(array_size) = len {
+ let to_primitive_array_fn = |col: &ColumnarValue| ->
PrimitiveArray<Int32Type> {
match col {
- ColumnarValue::Array(a) =>
Ok(a.as_primitive::<Int32Type>().to_owned()),
+ ColumnarValue::Array(a) =>
a.as_primitive::<Int32Type>().to_owned(),
_ => {
let v = scalar_value_fn(col).unwrap();
- Ok(PrimitiveArray::<Int32Type>::from_value(v,
scalar_count))
+ PrimitiveArray::<Int32Type>::from_value(v, array_size)
}
}
};
- 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<Date32Type> =
+ PrimitiveArray::builder(array_size);
for i in 0..array_size {
- construct_date_fn(
- &mut builder,
+ process_date(
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 {
+ // 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;
+ process_date(
+ scalar_value_fn(&years)?,
+ scalar_value_fn(&months)?,
+ scalar_value_fn(&days)?,
+ |days: i32| value = days,
+ )?;
+
+ ColumnarValue::Scalar(ScalarValue::Date32(Some(value)))
+ };
+
+ Ok(value)
+}
+
+fn process_date(
Review Comment:
Perhaps something like `make_date_inner`
```suggestion
/// Converts the year/month/day fields to a `i32`
/// representing the days from the unix epoch
/// and invokes date_consumer_fn with the value
fn make_date_inner(
```
##########
datafusion/physical-expr/src/datetime_expressions.rs:
##########
@@ -351,51 +317,77 @@ pub fn make_date(args: &[ColumnarValue]) ->
Result<ColumnarValue> {
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<PrimitiveArray<Int32Type>> {
+ let value = if let Some(array_size) = len {
+ let to_primitive_array_fn = |col: &ColumnarValue| ->
PrimitiveArray<Int32Type> {
match col {
- ColumnarValue::Array(a) =>
Ok(a.as_primitive::<Int32Type>().to_owned()),
+ ColumnarValue::Array(a) =>
a.as_primitive::<Int32Type>().to_owned(),
_ => {
let v = scalar_value_fn(col).unwrap();
- Ok(PrimitiveArray::<Int32Type>::from_value(v,
scalar_count))
+ PrimitiveArray::<Int32Type>::from_value(v, array_size)
}
}
};
- 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<Date32Type> =
+ PrimitiveArray::builder(array_size);
for i in 0..array_size {
- construct_date_fn(
- &mut builder,
+ process_date(
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 {
+ // 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;
+ process_date(
+ scalar_value_fn(&years)?,
+ scalar_value_fn(&months)?,
+ scalar_value_fn(&days)?,
+ |days: i32| value = days,
+ )?;
+
+ ColumnarValue::Scalar(ScalarValue::Date32(Some(value)))
+ };
+
+ Ok(value)
+}
+
+fn process_date(
+ year: i32,
+ month: i32,
+ day: i32,
+ mut date_consumer_fn: impl FnMut(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 {
Review Comment:
I am not quite sure what you are suggesting here @comphead I think this
code is moved from above
The fact that arrow uses `i32` for the subfields rather than `u32` is
somewhat non ideal I agree, but this code I believe just follows the standard
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
