Skip to content

Commit aa2efec

Browse files
authored
Merge 6cb6e76 into 4e54f1c
2 parents 4e54f1c + 6cb6e76 commit aa2efec

10 files changed

Lines changed: 484 additions & 23 deletions

File tree

‎docs/source/user-guide/latest/expressions.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -281,7 +281,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
281281
| `hour` | ✅ | Native | |
282282
| `last_day` | ✅ | Native | |
283283
| `localtimestamp` | ✅ | — | |
284-
| `make_date` | ✅ | Native | |
284+
| `make_date` | ✅ | Native | Computed arguments preserve year → month → day NULL checks. Expressions fall back when shared subexpressions or a parent could change which arguments Spark evaluates. IN probes stay native for literal/column candidates and array/struct constructors of those values; other computed candidates require Spark evaluation. |
285285
| `make_dt_interval` | ✅ | Codegen dispatch | |
286286
| `make_interval` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; intervals outside Arrow's nanosecond range are tracked by [#5279](https://github.com/apache/datafusion-comet/issues/5279); the native path is opt-in via allowIncompatible ([details](compatibility/expressions/datetime.md)) |
287287
| `make_time` | ✅ | — | Spark 4.1+; requires `spark.sql.timeType.enabled=true`, which Spark leaves off by default. Runs natively; remaining TIME type work is tracked by [#4288](https://github.com/apache/datafusion-comet/issues/4288) |

‎native/core/src/execution/planner.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ use datafusion_comet_spark_expr::{
118118
coerce_to_common_type, create_case_when, create_if_expr, decimal128_add_sub_width,
119119
jvm_udf::JvmScalarUdfExpr, spark_in_list, ArrayInsert, AtLeastNNonNulls, Cast, CheckOverflow,
120120
CreateNamedStruct, DecimalRescaleCheckOverflow, FloatOperands, GetArrayStructFields,
121-
GetStructField, IfExpr, ListExtract, NormalizeNaNAndZero, NormalizeNestedFloats,
121+
GetStructField, IfExpr, ListExtract, MakeDateExpr, NormalizeNaNAndZero, NormalizeNestedFloats,
122122
NullShortCircuit, SparkCastOptions, ToJson, UnboundColumn, WideDecimalBinaryExpr,
123123
WideDecimalOp,
124124
};
@@ -2192,6 +2192,11 @@ impl PhysicalPlanner {
21922192
} else {
21932193
scalar_expr
21942194
};
2195+
let scalar_expr = if fun_name == "make_date" {
2196+
Arc::new(MakeDateExpr::new(scalar_expr)) as Arc<dyn PhysicalExpr>
2197+
} else {
2198+
scalar_expr
2199+
};
21952200

21962201
// DF53 changed some UDFs (e.g. md5) to return StringViewArray at execution
21972202
// time (apache/datafusion#20045). Comet does not yet support view types, so

‎native/spark-expr/src/datetime_funcs/make_date.rs‎

Lines changed: 157 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,85 @@
1717

1818
use arrow::array::{Array, Date32Array, Int32Array};
1919
use arrow::compute::cast;
20-
use arrow::datatypes::DataType;
20+
use arrow::datatypes::{DataType, FieldRef, Schema};
21+
use arrow::record_batch::RecordBatch;
2122
use datafusion::common::{utils::take_function_args, DataFusionError, Result};
2223
use datafusion::logical_expr::{
2324
ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
2425
};
26+
use datafusion::physical_expr::PhysicalExpr;
27+
use std::fmt::{Display, Formatter};
2528
use std::sync::Arc;
2629

2730
use crate::{conversion_funcs::ymd_to_epoch_day, SparkError};
2831

32+
/// Skips the entire date expression on empty input, before any argument can raise an error.
33+
/// `inner` retains the planner's argument evaluation, including NULL short-circuiting when enabled.
34+
#[derive(Debug, Eq, Hash)]
35+
pub struct MakeDateExpr {
36+
inner: Arc<dyn PhysicalExpr>,
37+
}
38+
39+
impl MakeDateExpr {
40+
pub fn new(inner: Arc<dyn PhysicalExpr>) -> Self {
41+
Self { inner }
42+
}
43+
}
44+
45+
impl PartialEq for MakeDateExpr {
46+
fn eq(&self, other: &Self) -> bool {
47+
self.inner.eq(&other.inner)
48+
}
49+
}
50+
51+
impl Display for MakeDateExpr {
52+
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
53+
Display::fmt(&self.inner, f)
54+
}
55+
}
56+
57+
impl PhysicalExpr for MakeDateExpr {
58+
fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
59+
self.inner.data_type(input_schema)
60+
}
61+
62+
fn nullable(&self, input_schema: &Schema) -> Result<bool> {
63+
self.inner.nullable(input_schema)
64+
}
65+
66+
fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> {
67+
self.inner.return_field(input_schema)
68+
}
69+
70+
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
71+
if batch.num_rows() == 0 {
72+
return Ok(ColumnarValue::Array(Arc::new(Date32Array::new_null(0))));
73+
}
74+
self.inner.evaluate(batch)
75+
}
76+
77+
fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
78+
self.inner.children()
79+
}
80+
81+
fn with_new_children(
82+
self: Arc<Self>,
83+
children: Vec<Arc<dyn PhysicalExpr>>,
84+
) -> Result<Arc<dyn PhysicalExpr>> {
85+
Ok(Arc::new(Self::new(
86+
Arc::clone(&self.inner).with_new_children(children)?,
87+
)))
88+
}
89+
90+
fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
91+
self.inner.fmt_sql(f)
92+
}
93+
94+
fn is_volatile_node(&self) -> bool {
95+
self.inner.is_volatile_node()
96+
}
97+
}
98+
2999
/// Spark-compatible make_date function.
30100
/// Creates a date from year, month, and day columns.
31101
/// For an invalid `(year, month, day)` triple Spark returns NULL when `spark.sql.ansi.enabled` is
@@ -192,6 +262,92 @@ impl ScalarUDFImpl for SparkMakeDate {
192262
#[cfg(test)]
193263
mod tests {
194264
use super::*;
265+
use crate::{Cast, EvalMode, NullShortCircuit, SparkCastOptions};
266+
use arrow::datatypes::{Field, Schema};
267+
use arrow::record_batch::{RecordBatch, RecordBatchOptions};
268+
use datafusion::common::ScalarValue;
269+
use datafusion::config::ConfigOptions;
270+
use datafusion::logical_expr::ScalarUDF;
271+
use datafusion::physical_expr::expressions::{col, lit, Literal};
272+
use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr};
273+
274+
#[test]
275+
fn scalar_arguments_do_not_validate_dates_on_empty_batches() {
276+
let expression = MakeDateExpr::new(Arc::new(ScalarFunctionExpr::new(
277+
"make_date",
278+
Arc::new(ScalarUDF::new_from_impl(SparkMakeDate::new(true))),
279+
[2024, 13, 1]
280+
.into_iter()
281+
.map(|value| {
282+
Arc::new(Literal::new(ScalarValue::Int32(Some(value)))) as Arc<dyn PhysicalExpr>
283+
})
284+
.collect(),
285+
Arc::new(Field::new("make_date", DataType::Date32, true)),
286+
Arc::new(ConfigOptions::default()),
287+
)));
288+
let batch = |number_rows| {
289+
RecordBatch::try_new_with_options(
290+
Arc::new(Schema::empty()),
291+
vec![],
292+
&RecordBatchOptions::new().with_row_count(Some(number_rows)),
293+
)
294+
.unwrap()
295+
};
296+
let result = expression
297+
.evaluate(&batch(0))
298+
.unwrap()
299+
.into_array(0)
300+
.unwrap();
301+
assert_eq!(result.data_type(), &DataType::Date32);
302+
assert!(result.is_empty());
303+
assert!(expression.evaluate(&batch(1)).is_err());
304+
}
305+
306+
#[test]
307+
fn empty_batches_do_not_evaluate_make_date_arguments() {
308+
let schema = Arc::new(Schema::new(vec![Field::new("y", DataType::Int32, true)]));
309+
let batch = RecordBatch::try_new(
310+
Arc::clone(&schema),
311+
vec![Arc::new(Int32Array::from(vec![2024]))],
312+
)
313+
.unwrap();
314+
let bad_cast: Arc<dyn PhysicalExpr> = Arc::new(Cast::new(
315+
lit("bad"),
316+
DataType::Int32,
317+
SparkCastOptions::new_without_timezone(EvalMode::Ansi),
318+
None,
319+
None,
320+
));
321+
// Cover the first argument, where NullShortCircuit elides its wrapper, as well as
322+
// later arguments. The empty-input check is also required without NULL short-circuiting.
323+
for argument in 0..3 {
324+
for short_circuit in [false, true] {
325+
let mut args = vec![col("y", &schema).unwrap(), lit(1i32), lit(1i32)];
326+
args[argument] = Arc::clone(&bad_cast);
327+
let inner: Arc<dyn PhysicalExpr> = Arc::new(ScalarFunctionExpr::new(
328+
"make_date",
329+
Arc::new(ScalarUDF::new_from_impl(SparkMakeDate::new(true))),
330+
args,
331+
Arc::new(Field::new("make_date", DataType::Date32, true)),
332+
Arc::new(ConfigOptions::default()),
333+
));
334+
let inner = if short_circuit {
335+
NullShortCircuit::wrap(inner, &schema).unwrap()
336+
} else {
337+
inner
338+
};
339+
let expression = MakeDateExpr::new(inner);
340+
let ColumnarValue::Array(result) = expression.evaluate(&batch.slice(0, 0)).unwrap()
341+
else {
342+
panic!("make_date must return an empty array for an empty batch");
343+
};
344+
assert_eq!(result.data_type(), &DataType::Date32);
345+
assert!(result.is_empty());
346+
let error = expression.evaluate(&batch).unwrap_err().to_string();
347+
assert!(error.contains("CAST_INVALID_INPUT"), "{error}");
348+
}
349+
}
350+
}
195351

196352
#[test]
197353
fn test_make_date_valid() {

‎native/spark-expr/src/datetime_funcs/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ pub use extract_date_part::SparkHour;
4141
pub use extract_date_part::SparkMinute;
4242
pub use extract_date_part::SparkSecond;
4343
pub use hours::SparkHoursTransform;
44-
pub use make_date::SparkMakeDate;
44+
pub use make_date::{MakeDateExpr, SparkMakeDate};
4545
pub use make_interval::SparkMakeInterval;
4646
pub use make_time::SparkMakeTime;
4747
pub use next_day::SparkNextDay;

‎native/spark-expr/src/lib.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -87,10 +87,11 @@ pub use comet_scalar_funcs::{
8787
};
8888
pub use csv_funcs::*;
8989
pub use datetime_funcs::{
90-
spark_day_name, spark_month_name, spark_seconds_of_time, spark_to_time, SparkDateDiff,
91-
SparkDateFromUnixDate, SparkDateTrunc, SparkDayOfWeek, SparkHour, SparkHoursTransform,
92-
SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkMinute, SparkNextDay, SparkSecond,
93-
SparkSecondsToTimestamp, SparkUnixTimestamp, SparkWeekDay, TimestampTruncExpr,
90+
spark_day_name, spark_month_name, spark_seconds_of_time, spark_to_time, MakeDateExpr,
91+
SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkDayOfWeek, SparkHour,
92+
SparkHoursTransform, SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkMinute,
93+
SparkNextDay, SparkSecond, SparkSecondsToTimestamp, SparkUnixTimestamp, SparkWeekDay,
94+
TimestampTruncExpr,
9495
};
9596
pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult};
9697
pub use hash_funcs::*;

‎spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala‎

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -938,19 +938,25 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
938938
inputs: Seq[Attribute],
939939
binding: Boolean = true): Option[Expr] =
940940
CometExplainInfo.withExpressionFallbackReasonScope {
941-
val newExpr = DecimalPrecision.promote(expr)
942-
val result = exprToProtoInternal(newExpr, inputs, binding)
943-
if (!(newExpr eq expr)) {
944-
// `promote` rebuilt the tree, so the tags landed on copies that the operator does not hold.
945-
// Lift them onto `expr` so the roll-ups in `CometExecRule`, which walk the operator's own
946-
// expressions, still see them. Skipped in the common case where `promote` returned the same
947-
// tree and there is nothing to lift.
948-
liftCoverageTags(newExpr, expr)
949-
if (result.isEmpty) {
950-
liftFallbackReasons(newExpr, expr)
951-
}
941+
CometMakeDate.unsupportedParentReason(expr) match {
942+
case Some(reason) =>
943+
withFallbackReason(expr, reason)
944+
None
945+
case None =>
946+
val newExpr = DecimalPrecision.promote(expr)
947+
val result = exprToProtoInternal(newExpr, inputs, binding)
948+
if (!(newExpr eq expr)) {
949+
// `promote` rebuilt the tree, so the tags landed on copies that the operator does
950+
// not hold. Lift them onto `expr` so the roll-ups in `CometExecRule`, which walk
951+
// the operator's own expressions, still see them. Skipped in the common case where
952+
// `promote` returned the same tree and there is nothing to lift.
953+
liftCoverageTags(newExpr, expr)
954+
if (result.isEmpty) {
955+
liftFallbackReasons(newExpr, expr)
956+
}
957+
}
958+
result
952959
}
953-
result
954960
}
955961

956962
/**

‎spark/src/main/scala/org/apache/comet/serde/datetime.scala‎

Lines changed: 81 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ package org.apache.comet.serde
2121

2222
import java.util.Locale
2323

24-
import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, Cast, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, DivideDTInterval, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, MultiplyYMInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, TimestampAdd, TimestampDiff, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year}
24+
import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, BinaryComparison, CaseWhen, Cast, Coalesce, ConvertTimezone, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, DivideDTInterval, EqualNullSafe, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, If, In, InSet, LastDay, LeafExpression, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, MultiplyYMInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, SortOrder, TimestampAdd, TimestampDiff, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year}
2525
import org.apache.spark.sql.types.{CalendarIntervalType, DataType, DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType}
2626
import org.apache.spark.unsafe.types.UTF8String
2727

@@ -436,6 +436,85 @@ object CometNextDay extends CometExpressionSerde[NextDay] with CodegenDispatchFa
436436

437437
object CometMakeDate extends CometExpressionSerde[MakeDate] {
438438

439+
private val sharedArgumentReason =
440+
"make_date requires Spark evaluation when a computed argument after a nullable " +
441+
"argument cannot preserve Spark's shared-subexpression evaluation."
442+
443+
override def getUnsupportedReasons(): Seq[String] = Seq(sharedArgumentReason)
444+
445+
override def getSupportLevel(expr: MakeDate): SupportLevel = {
446+
val firstNullable = expr.children.indexWhere(_.nullable)
447+
val hasComputedMaskedArgument = firstNullable >= 0 &&
448+
expr.children.drop(firstNullable + 1).exists(!_.isInstanceOf[LeafExpression])
449+
if (hasComputedMaskedArgument && !canShortCircuitNulls(expr)) {
450+
Unsupported(Some(sharedArgumentReason))
451+
} else {
452+
Compatible()
453+
}
454+
}
455+
456+
// Native IN evaluates dynamic candidates for the whole batch. Only these constructors and
457+
// leaves are safe on rows where Spark skips a candidate after a match or a NULL probe.
458+
private def canEvaluateInCandidateEagerly(expr: Expression): Boolean = expr match {
459+
case _: Literal | _: Attribute => true
460+
case _: CreateArray | _: CreateNamedStruct =>
461+
expr.children.forall(canEvaluateInCandidateEagerly)
462+
case _ => false
463+
}
464+
465+
/** A native parent must preserve Spark's evaluation of a date and its sibling expressions. */
466+
private[serde] def unsupportedParentReason(expr: Expression): Option[String] = {
467+
var pending: List[(Expression, Option[String])] = (expr, None) :: Nil
468+
while (pending.nonEmpty) {
469+
val (current, parent) = pending.head
470+
pending = pending.tail
471+
if (current.isInstanceOf[MakeDate] && parent.isDefined) {
472+
return Some(s"make_date requires Spark evaluation under ${parent.get}")
473+
}
474+
def pushChildren(mask: Option[String]): Unit = {
475+
current.children.reverseIterator.foreach(child => pending = (child, mask) :: pending)
476+
}
477+
if (parent.isDefined) {
478+
pushChildren(parent)
479+
} else {
480+
current match {
481+
// These native expressions already select the rows passed to their children.
482+
case _: If | _: CaseWhen | _: MakeDate => pushChildren(None)
483+
// COALESCE lowers to CASE and evaluates an earlier child again in its result branch.
484+
case coalesce: Coalesce if coalesce.deterministic => pushChildren(None)
485+
// These parents evaluate every child in Spark, including NULL values.
486+
case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct =>
487+
pushChildren(None)
488+
// These parents evaluate left first and skip right only when left is NULL.
489+
case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass =>
490+
val left = current.children.head
491+
val right = current.children(1)
492+
val rightMask = if (left.nullable) Some(current.nodeName) else None
493+
pending = (left, None) :: (right, rightMask) :: pending
494+
case _: SortOrder => pushChildren(Some(current.nodeName))
495+
// Empty membership can skip its sole operand in Spark.
496+
case in: In if in.list.isEmpty => pushChildren(Some(current.nodeName))
497+
case inSet: InSet if inSet.hset.isEmpty => pushChildren(Some(current.nodeName))
498+
// Nonempty IN always evaluates its probe, but may skip candidates on individual rows.
499+
case in: In =>
500+
in.list.reverseIterator.foreach { child =>
501+
pending = (child, Some(in.nodeName)) :: pending
502+
}
503+
val probeParent = if (in.list.forall(canEvaluateInCandidateEagerly)) {
504+
None
505+
} else {
506+
Some("In with computed candidates")
507+
}
508+
pending = (in.value, probeParent) :: pending
509+
case _ if current.children.size <= 1 => pushChildren(None)
510+
// Other native multi-child expressions may evaluate operands that Spark skips.
511+
case _ => pushChildren(Some(current.nodeName))
512+
}
513+
}
514+
}
515+
None
516+
}
517+
439518
/**
440519
* `failOnError` mirrors `spark.sql.ansi.enabled`: under ANSI, Spark throws on an invalid
441520
* `(year, month, day)` triple rather than returning NULL. The resolved flag is passed to native
@@ -449,7 +528,7 @@ object CometMakeDate extends CometExpressionSerde[MakeDate] {
449528
DateType,
450529
expr.failOnError,
451530
childExpr: _*)
452-
optExpr
531+
withNullShortCircuit(expr, optExpr)
453532
}
454533
}
455534

0 commit comments

Comments
 (0)