sunchao commented on code in PR #5889:
URL: https://github.com/apache/datafusion-comet/pull/5889#discussion_r4316013895
##########
spark/src/main/scala/org/apache/comet/serde/CometScalarSubquery.scala:
##########
@@ -19,30 +19,271 @@
package org.apache.comet.serde
-import org.apache.spark.sql.catalyst.expressions.Attribute
-import org.apache.spark.sql.execution.ScalarSubquery
+import scala.util.DynamicVariable
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, DayOfWeek,
Expression, ExprId, GetDateField, Hour, LastDay, Minute, NamedExpression,
NextDay, Second, TimeZoneAwareExpression, ToUnixTimestamp, UnixTimestamp,
WeekDay}
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{ExecSubqueryExpression, ExpandExec,
ScalarSubquery, SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
QueryStageExec}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.types._
import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
import org.apache.comet.serde.QueryPlanSerde.{serializeDataType,
supportedDataType}
object CometScalarSubquery extends CometExpressionSerde[ScalarSubquery] {
+ private val structSubqueryInputs =
+ new TreeNodeTag[Set[ExprId]]("CometStructScalarSubqueryInputs")
+ private val currentStructSubqueryInputs = new
DynamicVariable[Set[ExprId]](Set.empty)
+
+ /**
+ * Trace struct scalar values before native conversion replaces their
producing projections.
+ * Attribute object identity is insufficient: aliases, nullable join outputs
and reused stages
+ * can create fresh attributes. Keep provenance by ExprId on the owning plan
instead.
+ *
+ * This is recomputed on every rule application, including the plans hidden
by AQE/reuse nodes.
+ * No provenance is inferred from a column's data type or carried in a
global ExprId registry.
+ */
+ private[comet] def tagPlanInputs(root: SparkPlan): Unit = {
+ val visited = new java.util.IdentityHashMap[SparkPlan, Set[ExprId]]()
+ def remember(plan: SparkPlan, sources: Set[ExprId]): Unit = {
+ val _ = visited.put(plan, sources)
+ }
+
+ def visit(plan: SparkPlan): Set[ExprId] = {
+ if (visited.containsKey(plan)) return visited.get(plan)
+ remember(plan, Set.empty)
+ plan.unsetTagValue(structSubqueryInputs)
+
+ val original = plan match {
+ case comet: CometExec => comet.originalPlan
+ case other => other
+ }
+ val children = plan match {
+ case reused: ReusedExchangeExec => Seq(reused.child)
+ case stage: QueryStageExec => Seq(stage.plan)
+ case adaptive: AdaptiveSparkPlanExec => Seq(adaptive.executedPlan)
+ case comet: CometExec if comet.children.isEmpty =>
Seq(comet.originalPlan)
+ case other => other.children
+ }
+ val childSources = children.map(visit)
+ var sources = childSources.flatten.toSet
+ if (sources.isEmpty &&
+ !original.expressions.exists(_.containsPattern(PLAN_EXPRESSION))) {
+ return Set.empty
+ }
+
+ // Union and reused stages can expose different ExprIds for the same
output positions.
+ original match {
+ case _: UnionExec | _: ReusedExchangeExec | _: QueryStageExec |
+ _: AdaptiveSparkPlanExec =>
+ plan.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (children.zip(childSources).exists { case (child, ids) =>
+ child.output.lift(ordinal).exists(a => ids.contains(a.exprId))
+ }) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+
+ val named = original.expressions.flatMap(_.collect { case e:
NamedExpression => e })
+ var changed = true
+ while (changed) {
+ val previous = sources
+ named.foreach { expression =>
+ if (hasStructSubquerySource(expression, sources)) sources +=
expression.exprId
+ }
+ original match {
+ case aggregate: BaseAggregateExec =>
+
aggregate.aggregateExpressions.zip(aggregate.aggregateAttributes).foreach {
+ case (expression, output) =>
+ val function = expression.aggregateFunction
+ if (hasStructSubquerySource(expression, sources) ||
+ function.inputAggBufferAttributes.exists(a =>
sources.contains(a.exprId))) {
+ sources += output.exprId
+ sources ++= function.aggBufferAttributes.map(_.exprId)
+ sources ++= function.inputAggBufferAttributes.map(_.exprId)
+ }
+ }
+ case expand: ExpandExec =>
+ expand.output.zipWithIndex.foreach { case (output, ordinal) =>
+ if (expand.projections.exists(p =>
hasStructSubquerySource(p(ordinal), sources))) {
+ sources += output.exprId
+ }
+ }
+ case _ =>
+ }
+ changed = previous != sources
+ }
+
+ plan.setTagValue(structSubqueryInputs, sources)
+ val outputs =
plan.output.iterator.map(_.exprId).filter(sources.contains).toSet
+ remember(plan, outputs)
+ outputs
+ }
+
+ val _ = visit(root)
+ }
+
+ /** Keep the lineage local to this conversion, even when an attribute was
copied or rebound. */
+ private[comet] def withPlanInputs[T](plan: SparkPlan)(f: => T): T =
+ currentStructSubqueryInputs.withValue(
+ plan.getTagValue(structSubqueryInputs).getOrElse(Set.empty))(f)
+
+ private val calendarConsumerFallbackReason =
+ "Calendar operations on struct scalar subqueries require Spark because
their runtime values " +
+ "may exceed the native calendar range"
+
+ private val timezoneConsumerFallbackReason =
+ "Timezone-dependent scalar subquery operations require Spark for non-fixed
timezones"
+
+ private val dateUnixTimestampFallbackReason =
+ "Date scalar subqueries in unix_timestamp require Spark because their
runtime values " +
+ "may exceed the native calendar range"
+
+ private def containsCalendarType(dt: DataType): Boolean = dt match {
+ case DateType | TimestampType | TimestampNTZType => true
+ case s: StructType => s.fields.exists(f =>
containsCalendarType(f.dataType))
+ case ArrayType(elementType, _) => containsCalendarType(elementType)
+ case MapType(keyType, valueType, _) =>
+ containsCalendarType(keyType) || containsCalendarType(valueType)
+ case _ => false
+ }
+
+ private def isTimestamp(dt: DataType): Boolean =
+ dt == TimestampType || dt == TimestampNTZType
+
+ private def isEpochCast(from: DataType, to: DataType): Boolean =
+ (isTimestamp(from) && to == LongType) || (from == LongType &&
isTimestamp(to))
+
+ private def hasStructSubquerySource(expr: Expression, inputs: Set[ExprId]):
Boolean =
+ expr.references.exists(a => inputs.contains(a.exprId)) ||
+ (expr.containsPattern(PLAN_EXPRESSION) && expr.exists {
+ case subquery: ExecSubqueryExpression =>
subquery.dataType.isInstanceOf[StructType]
+ case _ => false
+ })
+
+ private def hasStructSubquery(expr: Expression): Boolean =
+ hasStructSubquerySource(expr, currentStructSubqueryInputs.value)
+
+ private def clockFallbackReason(
+ child: Expression,
+ timeZoneId: Option[String]): Option[String] = {
+ if (child.dataType == TimestampNTZType ||
CometTimeZone.isFixedOffset(timeZoneId)) {
Review Comment:
[P2] Extend fixed-offset range handling to array-producing consumers. This
exemption accepts `hour`, `minute`, and `second` around row-dependent
expressions, but the new full-range branch in `extract_date_part.rs` only
handles `ColumnarValue::Scalar`. With session timezone `+01:00`, subquery reuse
enabled, and one two-row Parquet batch containing `(id, timestamp_micros) = (0,
9223372036854775807), (1, 0)`, run `SELECT id, hour(IF(id = 0, (SELECT max(ts)
FROM t), ts)), hour((SELECT min(ts) FROM t)) FROM t`. Spark merges the
subqueries into a struct. The mixed-row `IF` produces an array, so its extreme
timestamp reaches Arrow’s calendar conversion and returns NULL instead of
Spark’s hour `5`. Before this PR, the unsupported struct kept the projection on
Spark. `minute` and `second` similarly return NULL instead of `0` and `54`,
potentially changing downstream filtering and aggregation. Please apply
fixed-offset arithmetic to array inputs too, or retain fallback for wrappers
that
can produce arrays.
Evidence: An exact-head disposable Rust test evaluated `GetStructField` over
a struct literal containing `TimestampMicrosecond(i64::MAX, UTC)`, then
`IfExpr` with predicates `[true, false]` and an epoch-zero alternative, then
the three clock kernels in `+01:00`. Actual arrays were hour `[NULL, 1]`,
minute `[NULL, 0]`, and second `[NULL, 0]`. Expected arrays were `[5, 1]`, `[0,
0]`, and `[54, 0]`. Spark 3.4.3–4.2.0 sources use
`getLocalDateTime(...).getHour/getMinute/getSecond`; executing the equivalent
Java conversion confirmed these values. The native reproduction failed its
expected-value assertion. SQL admission and prior fallback were traced through
source rather than tested end to end.
--
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]
