@@ -21,7 +21,7 @@ package org.apache.comet.serde
2121
2222import 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 }
2525import org .apache .spark .sql .types .{CalendarIntervalType , DataType , DateType , DoubleType , FloatType , IntegerType , LongType , StringType , TimestampNTZType , TimestampType }
2626import org .apache .spark .unsafe .types .UTF8String
2727
@@ -436,6 +436,85 @@ object CometNextDay extends CometExpressionSerde[NextDay] with CodegenDispatchFa
436436
437437object 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