comphead commented on code in PR #5292:
URL: https://github.com/apache/datafusion-comet/pull/5292#discussion_r4488343916
##########
native/spark-expr/src/datetime_funcs/make_interval.rs:
##########
@@ -16,72 +16,270 @@
// under the License.
use crate::arithmetic_overflow_error;
-use arrow::array::Array;
-use arrow::datatypes::DataType;
-use datafusion::common::Result;
-use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs,
ScalarUDFImpl, Signature};
-use datafusion_spark::function::datetime::make_interval::SparkMakeInterval as
DataFusionMakeInterval;
+use arrow::array::{Array, ArrayRef, Decimal128Array, Int32Array, Int64Array,
StructArray};
+use arrow::buffer::NullBuffer;
+use arrow::datatypes::{DataType, Field, Fields};
+use datafusion::common::{DataFusionError, Result};
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+use std::collections::HashMap;
+use std::sync::Arc;
+
+const CALENDAR_INTERVAL_STRUCT_KEY: &str = "SPARK::calendarInterval::struct";
+const MICROS_PER_HOUR: i64 = 3_600_000_000;
+const MICROS_PER_MINUTE: i64 = 60_000_000;
+
+pub fn calendar_interval_type() -> DataType {
+ let months = Field::new("months", DataType::Int32,
false).with_metadata(HashMap::from([(
+ CALENDAR_INTERVAL_STRUCT_KEY.to_string(),
+ "true".to_string(),
+ )]));
+ DataType::Struct(Fields::from(vec![
+ months,
+ Field::new("days", DataType::Int32, false),
+ Field::new("microseconds", DataType::Int64, false),
+ ]))
+}
+
+/// Whether `fields` are the children of a Spark `CalendarIntervalType` value,
i.e. of the tagged
+/// struct built by [`calendar_interval_type`].
+///
+/// The `SPARK::calendarInterval::struct` marker on `months` is what
identifies an interval: a
+/// user's own `struct<months:int, days:int, microseconds:bigint>` has the
same shape and must keep
+/// ordinary struct semantics. Child nullability is deliberately not compared:
the struct's own
+/// validity decides whether an interval is null, and a nullability widening
that keeps field
+/// metadata must not turn an interval into a plain struct, which for the hash
kernels would mean
+/// silently hashing it as one. The JVM counterpart is
`Utils.isCalendarIntervalStructField`.
+pub fn is_calendar_interval_fields(fields: &Fields) -> bool {
+ let child = |index: usize, name: &str, data_type: &DataType| {
+ fields[index].name() == name && fields[index].data_type() == data_type
+ };
+ fields.len() == 3
+ && child(0, "months", &DataType::Int32)
+ && child(1, "days", &DataType::Int32)
+ && child(2, "microseconds", &DataType::Int64)
+ && fields[0]
+ .metadata()
+ .get(CALENDAR_INTERVAL_STRUCT_KEY)
+ .is_some_and(|value| value == "true")
+}
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct SparkMakeInterval {
- inner: DataFusionMakeInterval,
+ signature: Signature,
fail_on_error: bool,
}
impl SparkMakeInterval {
pub fn new(fail_on_error: bool) -> Self {
Self {
- inner: DataFusionMakeInterval::new(),
+ signature: Signature::exact(
+ vec![
+ DataType::Int32,
+ DataType::Int32,
+ DataType::Int32,
+ DataType::Int32,
+ DataType::Int32,
+ DataType::Int32,
+ DataType::Decimal128(18, 6),
+ ],
+ Volatility::Immutable,
+ ),
fail_on_error,
}
}
}
+fn make_interval(
+ years: i32,
+ months: i32,
+ weeks: i32,
+ days: i32,
+ hours: i32,
+ minutes: i32,
+ seconds_micros: i128,
+) -> std::result::Result<(i32, i32, i64), &'static str> {
+ let months = years
+ .checked_mul(12)
+ .and_then(|years| years.checked_add(months))
+ .ok_or("integer")?;
+ let days = weeks
+ .checked_mul(7)
+ .and_then(|weeks| weeks.checked_add(days))
+ .ok_or("integer")?;
+ let micros = i64::try_from(seconds_micros)
+ .ok()
+ .and_then(|seconds| {
+ i64::from(hours)
+ .checked_mul(MICROS_PER_HOUR)
+ .and_then(|hours| seconds.checked_add(hours))
+ })
+ .and_then(|micros| {
+ i64::from(minutes)
+ .checked_mul(MICROS_PER_MINUTE)
+ .and_then(|minutes| micros.checked_add(minutes))
+ })
+ .ok_or("long")?;
+ Ok((months, days, micros))
+}
+
impl ScalarUDFImpl for SparkMakeInterval {
fn name(&self) -> &str {
- self.inner.name()
+ "make_interval"
}
fn signature(&self) -> &Signature {
- self.inner.signature()
+ &self.signature
}
- fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
- self.inner.return_type(arg_types)
+ fn return_type(&self, _: &[DataType]) -> Result<DataType> {
+ Ok(calendar_interval_type())
}
fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
- let inputs = if self.fail_on_error {
- Some(args.args.clone())
- } else {
- None
- };
- let result = self.inner.invoke_with_args(args)?;
-
- if let Some(inputs) = inputs {
- let inputs_are_valid = |i| {
- inputs.iter().all(|input| match input {
- ColumnarValue::Array(values) => values.is_valid(i),
- ColumnarValue::Scalar(value) => !value.is_null(),
- })
- };
- let overflow = match &result {
- ColumnarValue::Array(values) =>
values.nulls().is_some_and(|nulls| {
- nulls.null_count() != 0
- && nulls
- .iter()
- .enumerate()
- .any(|(i, is_valid)| !is_valid &&
inputs_are_valid(i))
- }),
- ColumnarValue::Scalar(value) => value.is_null() &&
inputs_are_valid(0),
- };
- if overflow {
- // Spark identifies the integer or long operation that
overflowed. The native
- // wrapper only sees the result null mask, so it can only
report interval overflow.
- return Err(arithmetic_overflow_error("interval").into());
+ let number_rows = args.number_rows;
+ let arrays = args
+ .args
+ .into_iter()
+ .map(|arg| arg.into_array(number_rows))
+ .collect::<Result<Vec<_>>>()?;
+ let years = arrays[0]
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .ok_or_else(|| {
+ DataFusionError::Execution("make_interval: expected Int32
years".into())
+ })?;
+ let months = arrays[1]
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .ok_or_else(|| {
+ DataFusionError::Execution("make_interval: expected Int32
months".into())
+ })?;
+ let weeks = arrays[2]
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .ok_or_else(|| {
+ DataFusionError::Execution("make_interval: expected Int32
weeks".into())
+ })?;
+ let days = arrays[3]
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .ok_or_else(|| {
+ DataFusionError::Execution("make_interval: expected Int32
days".into())
+ })?;
+ let hours = arrays[4]
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .ok_or_else(|| {
+ DataFusionError::Execution("make_interval: expected Int32
hours".into())
+ })?;
+ let minutes = arrays[5]
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .ok_or_else(|| {
+ DataFusionError::Execution("make_interval: expected Int32
minutes".into())
+ })?;
+ let seconds = arrays[6]
+ .as_any()
+ .downcast_ref::<Decimal128Array>()
+ .ok_or_else(|| {
+ DataFusionError::Execution("make_interval: expected Decimal128
seconds".into())
+ })?;
+
+ let mut result_months = Vec::with_capacity(years.len());
+ let mut result_days = Vec::with_capacity(years.len());
+ let mut result_micros = Vec::with_capacity(years.len());
+ let mut valid = Vec::with_capacity(years.len());
+
+ for i in 0..years.len() {
+ if arrays.iter().any(|array| array.is_null(i)) {
+ result_months.push(0);
+ result_days.push(0);
+ result_micros.push(0);
+ valid.push(false);
+ continue;
+ }
Review Comment:
This checks all seven arrays for NULL on every row through `dyn Array`, and
builds a `Vec<bool>` for the result validity. Could the validity come from
folding the seven null buffers with `NullBuffer::union` once per batch, so the
loop only skips the rows that buffer marks NULL and then passes it straight to
`StructArray::new`? The seven downcast blocks above could also be
`datafusion::common::cast::as_int32_array` and `as_decimal128_array`, which
return an error the same way and are already used in `list_extract.rs` and
`abs.rs`. A rerun of the benchmark was requested anyway, so the `make_interval`
criterion bench would show whether it matters. I have not measured it.
##########
spark/src/main/scala/org/apache/comet/serde/datetime.scala:
##########
@@ -907,35 +907,60 @@ object CometMakeDTInterval extends
CometCodegenDispatch[MakeDTInterval]
object CometDivideDTInterval extends CometCodegenDispatch[DivideDTInterval]
object CometMakeInterval extends CometExpressionSerde[MakeInterval] with
CodegenDispatchFallback {
- private val incompatReason =
- "The native implementation converts seconds to `Float64`, which can lose
microsecond" +
- " precision, and stores time in nanoseconds, which overflows for large
time components" +
- " (hours, minutes, seconds) that Spark can represent."
- override def getCompatibleNotes(): Seq[String] = Seq(
- "Both the default JVM codegen-dispatch path and the native path currently
limit the" +
- " elapsed-time component to about 292 years in either direction. This
only affects" +
- " extreme intervals and is tracked in" +
- " [#5279](https://github.com/apache/datafusion-comet/issues/5279).")
+ private val eagerArgumentReason =
+ "An argument that follows a nullable argument and is not a column
reference, a literal, or " +
+ "a lossless up-cast of one, such as an ANSI `CAST` that can fail: Spark
stops at the " +
+ "first NULL argument without evaluating the rest, while the native
kernel evaluates every " +
+ "argument first"
- override def getIncompatibleReasons(): Seq[String] = Seq(incompatReason)
+ override def getUnsupportedReasons(): Seq[String] = Seq(eagerArgumentReason)
override def getSupportLevel(expr: MakeInterval): SupportLevel =
- Incompatible(Some(incompatReason))
+ if (argumentsCanBeEvaluatedEagerly(expr)) Compatible()
+ else Unsupported(Some(eagerArgumentReason))
+
+ /**
+ * Spark's `MakeInterval` is a `SeptenaryExpression`: it evaluates its
arguments left to right
+ * and returns NULL at the first NULL one, so an argument that follows a
nullable argument only
+ * runs on the rows where every earlier argument is non-null. The native
scalar function
+ * evaluates every argument over the whole batch before its kernel checks
for nulls. Running
+ * such an argument on the extra rows is unobservable only when it cannot
throw, carry state or
+ * have a side effect. That holds for a column reference, a literal, and a
lossless up-cast of
+ * either, which covers the implicit casts Spark adds to widen integral and
decimal arguments.
+ * Anything else routes through the JVM codegen dispatcher, which runs
Spark's own nested
+ * evaluation. `foldable` is not a usable test: an unfolded foldable
argument can still throw.
+ * Arguments before the first nullable one run on every row in Spark too, so
they are
+ * unrestricted.
+ */
+ private def argumentsCanBeEvaluatedEagerly(expr: MakeInterval): Boolean = {
+ def unobservable(arg: Expression): Boolean = arg match {
+ case _: Literal | _: Attribute | _: BoundReference => true
+ case cast: Cast =>
+ Cast.canUpCast(cast.child.dataType, cast.dataType) &&
unobservable(cast.child)
+ case _ => false
+ }
+ val firstNullable = expr.children.indexWhere(_.nullable)
+ firstNullable < 0 || expr.children.drop(firstNullable +
1).forall(unobservable)
+ }
Review Comment:
Since this was written, #6716 added `QueryPlanSerde.withNullShortCircuit`,
which sets `null_short_circuit` on the `ScalarFunc` so that native
`NullShortCircuit::wrap` evaluates each argument after the first only on rows
where no earlier argument is NULL. That is the `SeptenaryExpression` evaluation
order this check protects. `adding_a_new_expression.md` now points serdes at
it, and `CometArrayContains` and `CometSlice` already use it. Could `convert`
wrap the `scalarFunctionExprToProtoWithReturnType` call in
`withNullShortCircuit(expr, ...)` instead? Then `make_interval(y, CAST(s AS
INT))` would stay native rather than going to the dispatcher, and
`argumentsCanBeEvaluatedEagerly`, `eagerArgumentReason` and the `Unsupported`
arm could go. The `expect_dispatch(make_interval)` short-circuit query in
`make_interval_ansi.sql` would then become a native one. I have not run it.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala:
##########
@@ -97,6 +97,8 @@ private[arrow] object ArrowWriter {
case (_: DayTimeIntervalType, vector: DurationVector) => new
DurationWriter(vector)
case (CalendarIntervalType, vector: IntervalMonthDayNanoVector) =>
new IntervalMonthDayNanoWriter(vector)
Review Comment:
With `toArrowField` now building the struct, I could not find anything that
still creates an `IntervalMonthDayNanoVector` for a `CalendarIntervalType`
column. Every `ArrowWriter.create` caller builds its root from
`Utils.toArrowSchema`, and no native type maps to `Interval(MonthDayNano)` any
more. That leaves the #4898 read and write paths unused: this arm and
`IntervalMonthDayNanoWriter`, `IntervalChildVector` and `intervalChildren` in
`CometPlainVector.java`, the `IntervalMonthDayNanoVector` entries in
`CometBatchKernelCodegen.vectorClassBySimpleName`,
`CometBatchKernelCodegenInput` (the class list and the
`IntervalMonthDayNanoVector` condition of the `getInterval` case) and
`CometScalaUDFCodegen.specFor`, the `MONTH_DAY_NANO` arm of
`Utils.fromArrowType`, and the `IntervalMonthDayNano(None)` null literal in the
native planner (now in `planner/literal.rs` on main). Could they go in this PR,
since it retires that layout? Then a stray `Interval(MonthDayNano)` would fail
loudly ins
tead of being read with the old nanosecond conversion.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
