You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ur...@apache.org on 2026/09/02 15:49:05 UTC
(spark) branch master updated: [SPARK-58866][SQL] Add expression-level entry points to optimizer rules
This is an automated email from the ASF dual-hosted git repository.
uros-b pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 09515e6cb2a4 [SPARK-58866][SQL] Add expression-level entry points to optimizer rules
09515e6cb2a4 is described below
commit 09515e6cb2a4ba2ab400ad597966a5e6688cc1c8
Author: matvei-zamiatin-db <ma...@databricks.com>
AuthorDate: Wed Sep 2 15:48:50 2026 +0000
[SPARK-58866][SQL] Add expression-level entry points to optimizer rules
### What changes were proposed in this pull request?
This refactor exposes expression-level entry points on several `FinishAnalysis` rules and on the `RewriteWithExpression` rule so they can rewrite a single `Expression` tree, not just a whole `LogicalPlan`. Each rule's rewrite logic is factored into a reusable partial function, the existing plan-level `apply` delegates to it, and a new `applyForExpression` (or widened visibility) is added.
`TimeTravelSpec.resolveTimestampExpression` is converted to the new entry points, which lets it drop the fake `Project(Alias(ts), OneRowRelation())` wrapper and the two `asInstanceOf` casts it needed to unwrap the result. In this PR only `ComputeCurrentTime.applyForExpression` and `ReplaceExpressions.replace` have a production caller (`TimeTravelSpec`); the `ReplaceCurrentLike`, `SpecialDatetimeValues` and `RewriteWithExpression` expression entry points are added here, covered by unit [...]
For `ComputeCurrentTime`, `ReplaceExpressions`, `ReplaceCurrentLike` and `SpecialDatetimeValues`, the extracted partial functions are identical to the original inline bodies, so the plan-level rules produce the same output as before.
One plan-level behavior does change. `RewriteWithExpression`'s conditional-branch inlining previously did an unguarded `refToExpr(ref.id)` lookup that threw `NoSuchElementException` for a nested `With` inside a non-always-evaluated conditional branch that references an outer common expression. The refactor adds a `refToExpr.contains(ref.id)` guard so the outer reference is deferred to the enclosing `With` and inlined correctly instead. This is a latent-crash fix: the crashing shape is [...]
### Why are the changes needed?
The single-pass analyzer needs to apply these `FinishAnalysis` rewrites to individual, already-resolved expression trees rather than by running a full plan-wide rule pass. Today the logic is reachable only through each rule's plan-level apply, so it can't be invoked on a single `Expression`. Extracting shared entry points makes the rewrites reusable.
### Does this PR introduce _any_ user-facing change?
No. This is an internal refactor. All plan-level rules produce the same output as before.
### How was this patch tested?
New unit tests, 39 cases across 6 suites in sql/catalyst: `ReplaceExpressionsSuite, ReplaceCurrentLikeSuite, ComputeCurrentTimeSuite, SpecialDatetimeValuesSuite, ConstantFoldingSuite, RewriteWithExpressionSuite`
Existing coverage for the refactored rules (ComputeCurrentTimeSuite, SpecialDatetimeValuesSuite,
RewriteWithExpressionSuite, ConstantFoldingSuite, and the sql module suites) passes unchanged,
which is what guards the behavior-preserving part of the refactor.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 4.8
Closes #58119 from matvei-zamiatin-db/optimizer-rules-refactor.
Authored-by: matvei-zamiatin-db <ma...@databricks.com>
Signed-off-by: Uros Bojanic <22...@users.noreply.github.com>
---
.../sql/catalyst/analysis/TimeTravelSpec.scala | 9 +-
.../catalyst/optimizer/RewriteWithExpression.scala | 135 +++++++++-
.../sql/catalyst/optimizer/finishAnalysis.scala | 264 ++++++++++++--------
.../optimizer/ComputeCurrentTimeSuite.scala | 28 ++-
.../catalyst/optimizer/ConstantFoldingSuite.scala | 17 ++
.../optimizer/ReplaceCurrentLikeSuite.scala | 51 ++++
.../optimizer/ReplaceExpressionsSuite.scala | 46 ++++
.../optimizer/RewriteWithExpressionSuite.scala | 273 ++++++++++++++++++++-
.../optimizer/SpecialDatetimeValuesSuite.scala | 17 ++
9 files changed, 719 insertions(+), 121 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TimeTravelSpec.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TimeTravelSpec.scala
index d2fe2ad602c1..c91e848b2c9b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TimeTravelSpec.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TimeTravelSpec.scala
@@ -18,9 +18,8 @@
package org.apache.spark.sql.catalyst.analysis
import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Expression, Literal, SubqueryExpression, Unevaluable}
+import org.apache.spark.sql.catalyst.expressions.{Cast, Expression, Literal, SubqueryExpression, Unevaluable}
import org.apache.spark.sql.catalyst.optimizer.{ComputeCurrentTime, ReplaceExpressions}
-import org.apache.spark.sql.catalyst.plans.logical.{OneRowRelation, Project}
import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.types.TimestampType
import org.apache.spark.sql.util.CaseInsensitiveStringMap
@@ -47,11 +46,7 @@ object TimeTravelSpec {
throw QueryCompilationErrors.invalidTimestampExprForTimeTravel(
"INVALID_TIME_TRAVEL_TIMESTAMP_EXPR.INPUT", ts)
}
- val tsToEval = {
- val fakeProject = Project(Seq(Alias(ts, "ts")()), OneRowRelation())
- ComputeCurrentTime(ReplaceExpressions(fakeProject)).asInstanceOf[Project]
- .expressions.head.asInstanceOf[Alias].child
- }
+ val tsToEval = ComputeCurrentTime.applyForExpression(ReplaceExpressions.replace(ts))
tsToEval.foreach {
case _: Unevaluable =>
throw QueryCompilationErrors.invalidTimestampExprForTimeTravel(
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala
index 5d85e89e1eab..08119798c9aa 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala
@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.planning.PhysicalAggregation
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, PlanHelper, Project}
import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.trees.TreePattern.{COMMON_EXPR_REF, WITH_EXPRESSION}
+import org.apache.spark.sql.catalyst.trees.TreePattern.{COMMON_EXPR_REF, CURRENT_LIKE, WITH_EXPRESSION}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.util.Utils
@@ -66,6 +66,125 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
}
}
+ /**
+ * Rewrites the `With` expressions in a single expression tree by inlining their common
+ * expressions. Uses `transformUp` to handle nested `With`.
+ *
+ * Inlining duplicates a definition at every reference, so a definition referenced more than once
+ * must be safe to duplicate (see `isSafeToDuplicate`); anything else throws, as there is no plan
+ * here to pre-evaluate it into.
+ *
+ * Does not descend into subquery plans (e.g. `ScalarSubquery`). A caller whose expression
+ * may contain a subquery must rewrite those plans separately.
+ */
+ private[sql] def applyForExpression(expression: Expression): Expression =
+ inlineWith(expression, checkDuplication = true)
+
+ // The plan-level rewrite shares this to inline `With` in conditional branches, which may not be
+ // evaluated and so can't be pulled into a Project; it passes false to inline unconditionally.
+ private def inlineWith(expression: Expression, checkDuplication: Boolean): Expression = {
+ // Which definitions are safe to duplicate can only be decided outer-first, so it is collected
+ // in a separate pass before inlining bottom-up.
+ val safeIds = if (checkDuplication) {
+ safeToDuplicateIds(expression)
+ } else {
+ Set.empty[CommonExpressionId]
+ }
+ expression.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case With(child, defs) =>
+ if (checkDuplication) {
+ // Nested `With` is already rewritten, so a ref left in `child` belongs to `defs` or to
+ // an enclosing `With`. Only the ids defined here are checked.
+ val refCounts = child.collect { case ref: CommonExpressionRef => ref.id }
+ .groupBy(identity)
+ .transform((_, refs) => refs.size)
+ defs.foreach { commonExprDef =>
+ // Canonicalization re-numbers ids per `With`, so sibling `With`s reuse ids and break
+ // the global uniqueness `safeIds` relies on. Reject them, as the plan-level rewrite
+ // does.
+ if (commonExprDef.id.canonicalized) {
+ throw SparkException.internalError(
+ "Cannot inline canonicalized common expression definitions")
+ }
+ if (refCounts.getOrElse(commonExprDef.id, 0) > 1 &&
+ !safeIds.contains(commonExprDef.id)) {
+ throw SparkException.internalError(
+ "Cannot inline a common expression definition that is referenced more than " +
+ s"once and is not safe to duplicate: ${commonExprDef.child.sql}")
+ }
+ }
+ }
+ val refToExpr = defs.map(commonExprDef => commonExprDef.id -> commonExprDef.child).toMap
+ // The guard skips refs to an enclosing `With`, which are inlined when `transformUp` reaches
+ // it. Without it those refs hit `Map.apply` on a missing key and throw
+ // NoSuchElementException.
+ child.transformWithPruning(_.containsPattern(COMMON_EXPR_REF)) {
+ case ref: CommonExpressionRef if refToExpr.contains(ref.id) => refToExpr(ref.id)
+ }
+ }
+ }
+
+ /**
+ * Ids of the definitions that are safe to duplicate. Walks outer-first, so a definition that
+ * references an enclosing one is decided after it, and visits a definition's own nested `With`
+ * before the definition itself. Ids are globally unique, so one flat set is enough.
+ */
+ private def safeToDuplicateIds(expression: Expression): Set[CommonExpressionId] = {
+ val safeIds = mutable.Set.empty[CommonExpressionId]
+ def visit(e: Expression): Unit = if (e.containsPattern(WITH_EXPRESSION)) {
+ e match {
+ case With(child, defs) =>
+ defs.foreach { commonExprDef =>
+ visit(commonExprDef.child)
+ if (isSafeToDuplicate(commonExprDef.child, safeIds)) {
+ safeIds += commonExprDef.id
+ }
+ }
+ visit(child)
+ case other => other.children.foreach(visit)
+ }
+ }
+ visit(expression)
+ safeIds.toSet
+ }
+
+ /**
+ * Whether a copy of `e` at every reference always evaluates to the same value. `foldable` is not
+ * enough: `current_timestamp()` re-reads the clock, a TIME -> TIMESTAMP cast derives the current
+ * date, `aes_encrypt` is a foldable `StaticInvoke` with a random IV, and a Hive UDF is foldable
+ * on a user-declared determinism flag.
+ */
+ private def isSafeToDuplicate(
+ e: Expression,
+ safeIds: collection.Set[CommonExpressionId]): Boolean = {
+ // A whitelist: an unrecognized expression is rejected rather than assumed safe.
+ def isLiteralTree(e: Expression): Boolean = e match {
+ case _: Literal => true
+ // A ref is safe exactly when the definition it points to is.
+ case ref: CommonExpressionRef => safeIds.contains(ref.id)
+ // Any other leaf varies per row (attribute) or reads a clock.
+ case _: LeafExpression => false
+ // ComputeCurrentTime stabilizes a TIME -> TIMESTAMP cast into a DateTimeUtils.makeTimestamp*
+ // builder, which depends only on its arguments (none read a clock). Accept it here, ahead
+ // of the blanket NonSQLExpression rejection below, when every argument is a literal tree.
+ case _ if ComputeCurrentTime.isMakeTimestampBuilder(e) =>
+ e.children.forall(isLiteralTree)
+ // Arbitrary Java methods and UDFs can be foldable without being pure.
+ case _: NonSQLExpression | _: UserDefinedExpression => false
+ // A TIME -> TIMESTAMP cast has no date of its own; it derives one from the current date at
+ // eval time, so duplicated copies could resolve different dates. Reject it. Scope to exactly
+ // this source/target pair -- other timestamp-target casts (e.g. CAST('1970-01-01' AS
+ // TIMESTAMP)) are ordinary deterministic casts and fall through to the generic check below.
+ case c: Cast
+ if Cast.isTimeToTimestampNTZ(c.child.dataType, c.dataType) ||
+ Cast.isTimeToTimestampLTZ(c.child.dataType, c.dataType) => false
+ case _ => e.deterministic && e.children.forall(isLiteralTree)
+ }
+ // `current_time()` is a non-leaf whose only leaf is a literal, so the tree walk accepts it;
+ // the CURRENT_LIKE pattern is what rejects it.
+ !e.containsPattern(CURRENT_LIKE) && isLiteralTree(e)
+ }
+
private def applyInternal(p: LogicalPlan): LogicalPlan = {
val inputPlans = p.children
val commonExprIdSet = p.expressions
@@ -181,16 +300,10 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
rewriteWithExprAndInputPlans(
_, inputPlans, commonExprsPerChild, commonExprIdSet, isNestedWith))
val newExpr = c.withNewAlwaysEvaluatedInputs(newAlwaysEvaluatedInputs)
- // Use transformUp to handle nested With.
- newExpr.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
- case With(child, defs) =>
- // For With in the conditional branches, they may not be evaluated at all and we can't
- // pull the common expressions into a project which will always be evaluated. Inline it.
- val refToExpr = defs.map(d => d.id -> d.child).toMap
- child.transformWithPruning(_.containsPattern(COMMON_EXPR_REF)) {
- case ref: CommonExpressionRef => refToExpr(ref.id)
- }
- }
+ // For With in the conditional branches, they may not be evaluated at all and we can't
+ // pull the common expressions into a project which will always be evaluated. Inline it
+ // unconditionally. Use transformUp to handle nested With.
+ inlineWith(newExpr, checkDuplication = false)
case other => other.mapChildren(
rewriteWithExprAndInputPlans(
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala
index f1a9ffd54e9a..862d20870d68 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala
@@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.{CurrentUserContext, InternalRow}
import org.apache.spark.sql.catalyst.analysis.{CastSupport, ResolvedInlineTable}
import org.apache.spark.sql.catalyst.analysis.ResolveInlineTables.prepareForEval
import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.rules._
import org.apache.spark.sql.catalyst.trees.TreePattern._
@@ -53,7 +54,7 @@ object ReplaceExpressions extends Rule[LogicalPlan] {
case p => p.mapExpressions(replace)
}
- private def replace(e: Expression): Expression = e match {
+ private[sql] def replace(e: Expression): Expression = e match {
case r: RuntimeReplaceable => replace(r.replacement)
case _ => e.mapChildren(replace)
}
@@ -112,11 +113,52 @@ object EvalInlineTables extends Rule[LogicalPlan] with CastSupport {
*/
object ComputeCurrentTime extends Rule[LogicalPlan] {
def apply(plan: LogicalPlan): LogicalPlan = {
- val instant = Instant.now()
- val currentTimestampMicros = instantToMicros(instant)
- val currentTime = Literal.create(currentTimestampMicros, TimestampType)
- val currentTimeOfDayNanos = instantToNanosOfDay(instant, conf.sessionLocalTimeZone)
- val timezone = Literal.create(conf.sessionLocalTimeZone, StringType)
+ val snapshot = new TimeSnapshot(Instant.now())
+ plan.transformDownWithSubqueriesAndPruning(transformCondition) {
+ case subQuery =>
+ subQuery.transformAllExpressionsWithPruning(transformCondition)(
+ expressionTransform(snapshot))
+ }
+ }
+
+ /** Rewrites the current date/time functions in a single expression tree. */
+ private[sql] def applyForExpression(expression: Expression): Expression =
+ applyForExpression(expression, Instant.now())
+
+ /**
+ * Rewrites the current date/time functions in a single expression tree using `instant`. Callers
+ * reducing several expressions that must observe the same wall clock pass one shared instant.
+ *
+ * Does not descend into subquery plans (e.g. `ScalarSubquery`). A caller whose expression
+ * may contain a subquery must rewrite those plans separately, reusing the same `instant`.
+ */
+ private[sql] def applyForExpression(expression: Expression, instant: Instant): Expression = {
+ val snapshot = new TimeSnapshot(instant)
+ expression.transformWithPruning(transformCondition)(expressionTransform(snapshot))
+ }
+
+ // CAST_TO_TIMESTAMP is a dedicated tree-pattern bit set on Cast nodes whose target type is
+ // any timestamp type (NTZ or LTZ family). This lets the rule reach both TIME -> TIMESTAMP_NTZ
+ // and TIME -> TIMESTAMP_LTZ rewrites (which derive date fields from CURRENT_DATE) without the
+ // broad CAST pattern that previously widened traversal to nearly every plan. Node-level
+ // isTimeToTimestamp{NTZ,LTZ} guards keep rewrite semantics unchanged.
+ // We intentionally do NOT tag these casts with CURRENT_LIKE: inline-table validation treats
+ // CURRENT_LIKE as safe to defer, so tagging would let unrelated non-foldable timestamp-target
+ // casts (e.g. CAST(rand() AS TIMESTAMP_NTZ)) bypass validation (see SPARK-57618).
+ private def transformCondition(treePatternbits: TreePatternBits): Boolean = {
+ treePatternbits.containsPattern(CURRENT_LIKE) ||
+ treePatternbits.containsPattern(CAST_TO_TIMESTAMP)
+ }
+
+ /**
+ * Holds the wall-clock instant and the per-invocation literal caches shared by every rewrite, so
+ * all current date/time references within one invocation observe the same instant.
+ */
+ private class TimeSnapshot(val instant: Instant) {
+ val currentTimestampMicros: Long = instantToMicros(instant)
+ val currentTime: Literal = Literal.create(currentTimestampMicros, TimestampType)
+ val currentTimeOfDayNanos: Long = instantToNanosOfDay(instant, conf.sessionLocalTimeZone)
+ val timezone: Literal = Literal.create(conf.sessionLocalTimeZone, StringType)
val currentDates = collection.mutable.HashMap.empty[ZoneId, Literal]
val localTimestamps = collection.mutable.HashMap.empty[ZoneId, Literal]
// Nanosecond current-timestamp literals depend on the requested precision (sub-precision
@@ -124,94 +166,88 @@ object ComputeCurrentTime extends Rule[LogicalPlan] {
// by precision only; NTZ (localtimestamp) is keyed by (zone, precision) like its micro sibling.
val currentTimestampNanos = collection.mutable.HashMap.empty[Int, Literal]
val localTimestampNanos = collection.mutable.HashMap.empty[(ZoneId, Int), Literal]
+ }
- // CAST_TO_TIMESTAMP is a dedicated tree-pattern bit set on Cast nodes whose target type is
- // any timestamp type (NTZ or LTZ family). This lets the rule reach both TIME -> TIMESTAMP_NTZ
- // and TIME -> TIMESTAMP_LTZ rewrites (which derive date fields from CURRENT_DATE) without the
- // broad CAST pattern that previously widened traversal to nearly every plan. Node-level
- // isTimeToTimestamp{NTZ,LTZ} guards keep rewrite semantics unchanged.
- // We intentionally do NOT tag these casts with CURRENT_LIKE: inline-table validation treats
- // CURRENT_LIKE as safe to defer, so tagging would let unrelated non-foldable timestamp-target
- // casts (e.g. CAST(rand() AS TIMESTAMP_NTZ)) bypass validation (see SPARK-57618).
- def transformCondition(treePatternbits: TreePatternBits): Boolean = {
- treePatternbits.containsPattern(CURRENT_LIKE) ||
- treePatternbits.containsPattern(CAST_TO_TIMESTAMP)
- }
+ private def expressionTransform(ts: TimeSnapshot): PartialFunction[Expression, Expression] = {
+ case cd: CurrentDate =>
+ ts.currentDates.getOrElseUpdate(cd.zoneId, {
+ Literal.create(
+ DateTimeUtils.microsToDays(ts.currentTimestampMicros, cd.zoneId), DateType)
+ })
+ // CAST(time AS TIMESTAMP_NTZ(q)) fills the date fields from CURRENT_DATE. Rewrite it to
+ // a date+time builder anchored on the same query-stable current date literal that
+ // current_date() resolves to, so all references agree within the query. The builder's
+ // `replacement` (a StaticInvoke) is emitted directly because ReplaceExpressions has
+ // already run earlier in this batch and will not expand a fresh RuntimeReplaceable.
+ case c: Cast if Cast.isTimeToTimestampNTZ(c.child.dataType, c.dataType) =>
+ val dateLit = ts.currentDates.getOrElseUpdate(c.zoneId, {
+ Literal.create(
+ DateTimeUtils.microsToDays(ts.currentTimestampMicros, c.zoneId), DateType)
+ })
+ c.dataType match {
+ case n: TimestampNTZNanosType =>
+ MakeTimestampNTZNanos(dateLit, c.child, n.precision).replacement
+ case _: TimestampNTZType =>
+ MakeTimestampNTZ(dateLit, c.child).replacement
+ case other =>
+ // Unreachable: the outer guard `Cast.isTimeToTimestampNTZ` only matches the micro
+ // TimestampNTZType and the nanosecond TimestampNTZNanosType targets.
+ throw SparkException.internalError(
+ s"Unexpected target type in TIME -> TIMESTAMP_NTZ rewrite: $other")
+ }
+ // CAST(time AS TIMESTAMP_LTZ(q)) likewise fills the date fields from CURRENT_DATE.
+ // Rewrite it to a zone-aware date+time builder anchored on the same query-stable current
+ // date literal, so all references agree within the query.
+ case c: Cast if Cast.isTimeToTimestampLTZ(c.child.dataType, c.dataType) =>
+ val dateLit = ts.currentDates.getOrElseUpdate(c.zoneId, {
+ Literal.create(
+ DateTimeUtils.microsToDays(ts.currentTimestampMicros, c.zoneId), DateType)
+ })
+ c.dataType match {
+ case l: TimestampLTZNanosType =>
+ MakeTimestampLTZNanos(dateLit, c.child, l.precision, c.timeZoneId).replacement
+ case _: TimestampType =>
+ MakeTimestampLTZ(dateLit, c.child, c.timeZoneId).replacement
+ case other =>
+ // Unreachable: the outer guard `Cast.isTimeToTimestampLTZ` only matches the micro
+ // TimestampType and the nanosecond TimestampLTZNanosType targets.
+ throw SparkException.internalError(
+ s"Unexpected target type in TIME -> TIMESTAMP_LTZ rewrite: $other")
+ }
+ case currentTimeType : CurrentTime =>
+ val truncatedTime = truncateTimeToPrecision(ts.currentTimeOfDayNanos,
+ currentTimeType.precision)
+ Literal.create(truncatedTime, TimeType(currentTimeType.precision))
+ case CurrentTimestamp() | Now() => ts.currentTime
+ case ct: CurrentTimestampNanos =>
+ ts.currentTimestampNanos.getOrElseUpdate(ct.precision, {
+ Literal.create(
+ instantToTimestampNanos(ts.instant, ct.precision),
+ TimestampLTZNanosType(ct.precision))
+ })
+ case CurrentTimeZone() => ts.timezone
+ case localTimestamp: LocalTimestamp =>
+ ts.localTimestamps.getOrElseUpdate(localTimestamp.zoneId, {
+ val asDateTime = LocalDateTime.ofInstant(ts.instant, localTimestamp.zoneId)
+ Literal.create(localDateTimeToMicros(asDateTime), TimestampNTZType)
+ })
+ case lt: LocalTimestampNanos =>
+ ts.localTimestampNanos.getOrElseUpdate((lt.zoneId, lt.precision), {
+ val asDateTime = LocalDateTime.ofInstant(ts.instant, lt.zoneId)
+ Literal.create(
+ localDateTimeToTimestampNanos(asDateTime, lt.precision),
+ TimestampNTZNanosType(lt.precision))
+ })
+ }
- plan.transformDownWithSubqueriesAndPruning(transformCondition) {
- case subQuery =>
- subQuery.transformAllExpressionsWithPruning(transformCondition) {
- case cd: CurrentDate =>
- currentDates.getOrElseUpdate(cd.zoneId, {
- Literal.create(
- DateTimeUtils.microsToDays(currentTimestampMicros, cd.zoneId), DateType)
- })
- // CAST(time AS TIMESTAMP_NTZ(q)) fills the date fields from CURRENT_DATE. Rewrite it to
- // a date+time builder anchored on the same query-stable current date literal that
- // current_date() resolves to, so all references agree within the query. The builder's
- // `replacement` (a StaticInvoke) is emitted directly because ReplaceExpressions has
- // already run earlier in this batch and will not expand a fresh RuntimeReplaceable.
- case c: Cast if Cast.isTimeToTimestampNTZ(c.child.dataType, c.dataType) =>
- val dateLit = currentDates.getOrElseUpdate(c.zoneId, {
- Literal.create(
- DateTimeUtils.microsToDays(currentTimestampMicros, c.zoneId), DateType)
- })
- c.dataType match {
- case n: TimestampNTZNanosType =>
- MakeTimestampNTZNanos(dateLit, c.child, n.precision).replacement
- case _: TimestampNTZType =>
- MakeTimestampNTZ(dateLit, c.child).replacement
- case other =>
- // Unreachable: the outer guard `Cast.isTimeToTimestampNTZ` only matches the micro
- // TimestampNTZType and the nanosecond TimestampNTZNanosType targets.
- throw SparkException.internalError(
- s"Unexpected target type in TIME -> TIMESTAMP_NTZ rewrite: $other")
- }
- // CAST(time AS TIMESTAMP_LTZ(q)) likewise fills the date fields from CURRENT_DATE.
- // Rewrite it to a zone-aware date+time builder anchored on the same query-stable current
- // date literal, so all references agree within the query.
- case c: Cast if Cast.isTimeToTimestampLTZ(c.child.dataType, c.dataType) =>
- val dateLit = currentDates.getOrElseUpdate(c.zoneId, {
- Literal.create(
- DateTimeUtils.microsToDays(currentTimestampMicros, c.zoneId), DateType)
- })
- c.dataType match {
- case l: TimestampLTZNanosType =>
- MakeTimestampLTZNanos(dateLit, c.child, l.precision, c.timeZoneId).replacement
- case _: TimestampType =>
- MakeTimestampLTZ(dateLit, c.child, c.timeZoneId).replacement
- case other =>
- // Unreachable: the outer guard `Cast.isTimeToTimestampLTZ` only matches the micro
- // TimestampType and the nanosecond TimestampLTZNanosType targets.
- throw SparkException.internalError(
- s"Unexpected target type in TIME -> TIMESTAMP_LTZ rewrite: $other")
- }
- case currentTimeType : CurrentTime =>
- val truncatedTime = truncateTimeToPrecision(currentTimeOfDayNanos,
- currentTimeType.precision)
- Literal.create(truncatedTime, TimeType(currentTimeType.precision))
- case CurrentTimestamp() | Now() => currentTime
- case ct: CurrentTimestampNanos =>
- currentTimestampNanos.getOrElseUpdate(ct.precision, {
- Literal.create(
- instantToTimestampNanos(instant, ct.precision),
- TimestampLTZNanosType(ct.precision))
- })
- case CurrentTimeZone() => timezone
- case localTimestamp: LocalTimestamp =>
- localTimestamps.getOrElseUpdate(localTimestamp.zoneId, {
- val asDateTime = LocalDateTime.ofInstant(instant, localTimestamp.zoneId)
- Literal.create(localDateTimeToMicros(asDateTime), TimestampNTZType)
- })
- case lt: LocalTimestampNanos =>
- localTimestampNanos.getOrElseUpdate((lt.zoneId, lt.precision), {
- val asDateTime = LocalDateTime.ofInstant(instant, lt.zoneId)
- Literal.create(
- localDateTimeToTimestampNanos(asDateTime, lt.precision),
- TimestampNTZNanosType(lt.precision))
- })
- }
- }
+ /**
+ * Whether `e` is a `DateTimeUtils.makeTimestamp*` builder `StaticInvoke`, the shape
+ * `expressionTransform` emits when stabilizing a TIME -> TIMESTAMP cast.
+ */
+ private[optimizer] def isMakeTimestampBuilder(e: Expression): Boolean = e match {
+ case si: StaticInvoke if si.staticObject == classOf[DateTimeUtils.type] =>
+ si.functionName.startsWith("makeTimestamp")
+ case _ => false
}
}
@@ -221,13 +257,27 @@ object ComputeCurrentTime extends Rule[LogicalPlan] {
*/
case class ReplaceCurrentLike(catalogManager: CatalogManager) extends Rule[LogicalPlan] {
def apply(plan: LogicalPlan): LogicalPlan = {
+ plan.transformAllExpressionsWithPruning(_.containsPattern(CURRENT_LIKE))(currentLikeRewrite)
+ }
+
+ /**
+ * Replaces the current catalog/database/path/user expressions in a single expression tree.
+ *
+ * Does not descend into subquery plans (e.g. `ScalarSubquery`). A caller whose expression
+ * may contain a subquery must rewrite those plans separately.
+ */
+ private[sql] def applyForExpression(expression: Expression): Expression = {
+ expression.transformWithPruning(_.containsPattern(CURRENT_LIKE))(currentLikeRewrite)
+ }
+
+ private def currentLikeRewrite: PartialFunction[Expression, Expression] = {
import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
lazy val currentNamespace = catalogManager.currentNamespace.quoted
lazy val currentCatalog = catalogManager.currentCatalog.name()
lazy val currentUser = CurrentUserContext.getCurrentUser
lazy val currentPathStr = catalogManager.currentPathString
- plan.transformAllExpressionsWithPruning(_.containsPattern(CURRENT_LIKE)) {
+ {
case CurrentDatabase() =>
Literal.create(currentNamespace, StringType)
case CurrentCatalog() =>
@@ -250,14 +300,26 @@ object SpecialDatetimeValues extends Rule[LogicalPlan] {
TimestampType -> convertSpecialTimestamp,
TimestampNTZType -> convertSpecialTimestampNTZ)
def apply(plan: LogicalPlan): LogicalPlan = {
- plan.transformAllExpressionsWithPruning(_.containsPattern(CAST)) {
- case cast @ Cast(e, dt @ (DateType | TimestampType | TimestampNTZType), _, _)
- if e.foldable && e.dataType == StringType =>
- Option(e.eval())
- .flatMap(s => conv(dt)(s.toString, cast.zoneId))
- .map(Literal(_, dt))
- .getOrElse(cast)
- }
+ plan.transformAllExpressionsWithPruning(_.containsPattern(CAST))(specialDatetimeRewrite)
+ }
+
+ /**
+ * Replaces casts of foldable special datetime strings in a single expression tree.
+ *
+ * Does not descend into subquery plans (e.g. `ScalarSubquery`). A caller whose expression
+ * may contain a subquery must rewrite those plans separately.
+ */
+ private[sql] def applyForExpression(expression: Expression): Expression = {
+ expression.transformWithPruning(_.containsPattern(CAST))(specialDatetimeRewrite)
+ }
+
+ private val specialDatetimeRewrite: PartialFunction[Expression, Expression] = {
+ case cast @ Cast(e, dt @ (DateType | TimestampType | TimestampNTZType), _, _)
+ if e.foldable && e.dataType == StringType =>
+ Option(e.eval())
+ .flatMap(s => conv(dt)(s.toString, cast.zoneId))
+ .map(Literal(_, dt))
+ .getOrElse(cast)
}
}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ComputeCurrentTimeSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ComputeCurrentTimeSuite.scala
index dc8cfd74c785..cb93b33244ad 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ComputeCurrentTimeSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ComputeCurrentTimeSuite.scala
@@ -24,7 +24,7 @@ import scala.concurrent.duration._
import scala.jdk.CollectionConverters.MapHasAsScala
import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Cast, CurrentDate, CurrentTime, CurrentTimestamp, CurrentTimestampNanos, CurrentTimeZone, Expression, InSubquery, ListQuery, Literal, LocalTimestamp, LocalTimestampNanos, Now}
+import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Cast, CurrentDate, CurrentTime, CurrentTimestamp, CurrentTimestampNanos, CurrentTimeZone, Expression, InSubquery, ListQuery, Literal, LocalTimestamp, LocalTimestampNanos, Now, ScalarSubquery}
import org.apache.spark.sql.catalyst.plans.PlanTest
import org.apache.spark.sql.catalyst.plans.logical.{Filter, LocalRelation, LogicalPlan, Project}
import org.apache.spark.sql.catalyst.rules.RuleExecutor
@@ -338,6 +338,32 @@ class ComputeCurrentTimeSuite extends PlanTest {
checkLiterals({ zoneId: String => CurrentDate(Some(zoneId)) }, numUniqueZoneIds)
}
+ test("applyForExpression rewrites current date/time leaves to literals") {
+ val date = ComputeCurrentTime.applyForExpression(CurrentDate(Some("UTC")))
+ assert(date.isInstanceOf[Literal] && date.dataType == DateType)
+ val timestamp = ComputeCurrentTime.applyForExpression(CurrentTimestamp())
+ assert(timestamp.isInstanceOf[Literal] && timestamp.dataType == TimestampType)
+ }
+
+ test("applyForExpression shares the provided instant across expressions") {
+ val instant = Instant.now()
+ val timestamp = ComputeCurrentTime.applyForExpression(CurrentTimestamp(), instant)
+ val now = ComputeCurrentTime.applyForExpression(Now(), instant)
+ assert(timestamp.isInstanceOf[Literal] && now.isInstanceOf[Literal])
+ val micros = DateTimeUtils.instantToMicros(instant)
+ assert(timestamp.asInstanceOf[Literal].value == micros)
+ assert(now.asInstanceOf[Literal].value == micros)
+ }
+
+ test("applyForExpression does not descend into subquery plans") {
+ // current_timestamp() inside the subquery plan is left for a separate pass.
+ val subquery = ScalarSubquery(
+ Project(Seq(Alias(CurrentTimestamp(), "ts")()), LocalRelation()))
+ val rewritten = ComputeCurrentTime.applyForExpression(subquery)
+ val innerPlan = rewritten.asInstanceOf[ScalarSubquery].plan
+ assert(innerPlan.expressions.exists(_.exists(_.isInstanceOf[CurrentTimestamp])))
+ }
+
test("CAST(time AS TIMESTAMP_NTZ) is stabilized with the query current date") {
val timeLit = Literal(0L, TimeType(6))
val in = Project(Seq(
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConstantFoldingSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConstantFoldingSuite.scala
index 0489308a1ab2..eb9c43f204a8 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConstantFoldingSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConstantFoldingSuite.scala
@@ -45,6 +45,23 @@ class ConstantFoldingSuite extends PlanTest {
val testRelation = LocalRelation($"a".int, $"b".int, $"c".int)
+ test("constantFolding folds a foldable expression tree to a literal") {
+ assert(ConstantFolding.constantFolding(Literal(1) + Literal(2)) == Literal(3))
+ }
+
+ test("constantFolding folds nested foldable subtrees") {
+ assert(ConstantFolding.constantFolding((Literal(1) + Literal(2)) * Literal(3)) == Literal(9))
+ }
+
+ test("constantFolding folds only the foldable parts of a partially-foldable tree") {
+ val a = testRelation.output.head
+ assert(ConstantFolding.constantFolding(a + (Literal(1) + Literal(2))) == a + Literal(3))
+ }
+
+ test("constantFolding returns a literal unchanged") {
+ assert(ConstantFolding.constantFolding(Literal(5)) == Literal(5))
+ }
+
test("eliminate subqueries") {
val originalQuery =
testRelation
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCurrentLikeSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCurrentLikeSuite.scala
new file mode 100644
index 000000000000..e6f72b24a225
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCurrentLikeSuite.scala
@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.analysis.{EmptyFunctionRegistry, EmptyTableFunctionRegistry, FakeV2SessionCatalog}
+import org.apache.spark.sql.catalyst.catalog.{InMemoryCatalog, SessionCatalog}
+import org.apache.spark.sql.catalyst.expressions.{Coalesce, CurrentCatalog, CurrentDatabase, Literal}
+import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
+import org.apache.spark.sql.connector.catalog.DefaultCatalogManager
+import org.apache.spark.sql.types.StringType
+
+class ReplaceCurrentLikeSuite extends SparkFunSuite {
+ private val catalogManager = new DefaultCatalogManager(
+ FakeV2SessionCatalog,
+ new SessionCatalog(new InMemoryCatalog, EmptyFunctionRegistry, EmptyTableFunctionRegistry))
+ private val rule = ReplaceCurrentLike(catalogManager)
+
+ test("applyForExpression replaces current catalog/database with literals") {
+ assert(rule.applyForExpression(CurrentCatalog()) ==
+ Literal.create(catalogManager.currentCatalog.name(), StringType))
+ assert(rule.applyForExpression(CurrentDatabase()) ==
+ Literal.create(catalogManager.currentNamespace.quoted, StringType))
+ }
+
+ test("applyForExpression rewrites current-like nested in a larger expression") {
+ val expr = Coalesce(Seq(CurrentCatalog(), Literal("x")))
+ val expected = Coalesce(Seq(
+ Literal.create(catalogManager.currentCatalog.name(), StringType), Literal("x")))
+ assert(rule.applyForExpression(expr) == expected)
+ }
+
+ test("applyForExpression leaves expressions without current-like unchanged") {
+ assert(rule.applyForExpression(Literal("x")) == Literal("x"))
+ }
+}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceExpressionsSuite.scala
new file mode 100644
index 000000000000..7c919c53d057
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceExpressionsSuite.scala
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.expressions.{Add, Coalesce, Literal, Nvl, RuntimeReplaceable}
+
+class ReplaceExpressionsSuite extends SparkFunSuite {
+
+ test("replace expands RuntimeReplaceable expressions") {
+ val replaced = ReplaceExpressions.replace(new Nvl(Literal(1), Literal(2)))
+ assert(!replaced.exists(_.isInstanceOf[RuntimeReplaceable]))
+ assert(replaced == Coalesce(Seq(Literal(1), Literal(2))))
+ }
+
+ test("replace expands nested RuntimeReplaceable expressions") {
+ val nested = new Nvl(new Nvl(Literal(1), Literal(2)), Literal(3))
+ val replaced = ReplaceExpressions.replace(nested)
+ assert(!replaced.exists(_.isInstanceOf[RuntimeReplaceable]))
+ assert(replaced == Coalesce(Seq(Coalesce(Seq(Literal(1), Literal(2))), Literal(3))))
+ }
+
+ test("replace recurses into children of non-RuntimeReplaceable expressions") {
+ val e = Add(Literal(1), new Nvl(Literal(2), Literal(3)))
+ assert(ReplaceExpressions.replace(e) == Add(Literal(1), Coalesce(Seq(Literal(2), Literal(3)))))
+ }
+
+ test("replace leaves expressions without RuntimeReplaceable unchanged") {
+ assert(ReplaceExpressions.replace(Literal(1)) == Literal(1))
+ }
+}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala
index 8918b58ca1b5..e0930afc801e 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala
@@ -17,13 +17,21 @@
package org.apache.spark.sql.catalyst.optimizer
-import org.apache.spark.sql.catalyst.analysis.TempResolvedColumn
+import java.time.Instant
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.analysis.{EmptyFunctionRegistry, EmptyTableFunctionRegistry, FakeV2SessionCatalog, TempResolvedColumn}
+import org.apache.spark.sql.catalyst.catalog.{InMemoryCatalog, SessionCatalog}
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.dsl.plans._
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.PlanTest
import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.catalyst.util.DateTimeUtils
+import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
+import org.apache.spark.sql.connector.catalog.DefaultCatalogManager
+import org.apache.spark.sql.types.{DateType, StringType, TimestampNTZType, TimestampType, TimeType}
class RewriteWithExpressionSuite extends PlanTest {
@@ -36,6 +44,22 @@ class RewriteWithExpressionSuite extends PlanTest {
private val testRelation = LocalRelation($"a".int, $"b".int)
private val testRelation2 = LocalRelation($"x".int, $"y".int)
+ private val catalogManager = new DefaultCatalogManager(
+ FakeV2SessionCatalog,
+ new SessionCatalog(new InMemoryCatalog, EmptyFunctionRegistry, EmptyTableFunctionRegistry))
+
+ /**
+ * Runs the expression-level rewrites in the same order as the optimizer's `FinishAnalysis`
+ * batch, then inlines the `With`.
+ */
+ private def finishAnalysisAndInline(expr: Expression, instant: Instant): Expression = {
+ val afterReplace = ReplaceExpressions.replace(expr)
+ val afterCurrentTime = ComputeCurrentTime.applyForExpression(afterReplace, instant)
+ val afterCurrentLike = ReplaceCurrentLike(catalogManager).applyForExpression(afterCurrentTime)
+ RewriteWithExpression.applyForExpression(
+ SpecialDatetimeValues.applyForExpression(afterCurrentLike))
+ }
+
private def normalizeCommonExpressionIds(plan: LogicalPlan): LogicalPlan = {
plan.transformAllExpressions {
case a: Alias if a.name.startsWith("_common_expr") =>
@@ -59,6 +83,239 @@ class RewriteWithExpressionSuite extends PlanTest {
comparePlans(Optimizer.execute(plan), testRelation.select((a + a).as("col")))
}
+ test("applyForExpression inlines a literal referenced more than once") {
+ val expr = With(Literal(1)) { case Seq(ref) =>
+ ref + ref
+ }
+ val rewritten = RewriteWithExpression.applyForExpression(expr)
+ assert(!rewritten.isInstanceOf[With])
+ assert(rewritten == Literal(1) + Literal(1))
+ }
+
+ test("applyForExpression inlines a non-literal referenced at most once") {
+ // Inlining does not duplicate the definition here, so it is safe whatever the definition is.
+ val a = testRelation.output.head
+ val referencedOnce = With(a + a) { case Seq(ref) =>
+ ref * Literal(2)
+ }
+ assert(RewriteWithExpression.applyForExpression(referencedOnce) == (a + a) * Literal(2))
+ val neverReferenced = With(a + a) { case Seq(_) =>
+ Literal(3)
+ }
+ assert(RewriteWithExpression.applyForExpression(neverReferenced) == Literal(3))
+ }
+
+ test("applyForExpression rejects a non-literal referenced more than once") {
+ // Inlining would duplicate `a + a` at every reference, violating the evaluate-once contract.
+ val a = testRelation.output.head
+ val expr = With(a + a) { case Seq(ref) =>
+ ref * ref
+ }
+ intercept[SparkException] {
+ RewriteWithExpression.applyForExpression(expr)
+ }
+ }
+
+ test("applyForExpression inlines a tree of literals referenced more than once") {
+ val expr = With(Literal(1) + Literal(1)) { case Seq(ref) =>
+ ref * ref
+ }
+ val rewritten = RewriteWithExpression.applyForExpression(expr)
+ assert(rewritten == (Literal(1) + Literal(1)) * (Literal(1) + Literal(1)))
+ }
+
+ test("applyForExpression rejects an impure foldable definition referenced more than once") {
+ // aes_encrypt becomes a foldable StaticInvoke that draws a fresh random IV on every eval, so
+ // two inlined copies would encrypt to different values.
+ val aes = ReplaceExpressions.replace(
+ new AesEncrypt(Literal("abc".getBytes), Literal("1234567890123456".getBytes)))
+ assert(aes.foldable, "the AES rewrite is only interesting while it stays foldable")
+ val expr = With(aes) { case Seq(ref) =>
+ EqualTo(ref, ref)
+ }
+ intercept[SparkException] {
+ RewriteWithExpression.applyForExpression(expr)
+ }
+ }
+
+ test("applyForExpression rejects canonicalized common expression ids") {
+ // Canonicalization re-numbers ids per `With`, starting from 1, so these two siblings both get
+ // id 1: the safe (literal) definition would otherwise mark id 1 safe in the flat `safeIds` set
+ // and authorize duplicating the unsafe (per-row attribute) definition with the same id.
+ val a = testRelation.output.head
+ val safe = With(Literal(1)) { case Seq(ref) => ref + ref }
+ val unsafe = With(a) { case Seq(ref) => ref + ref }
+ val expr = (safe + unsafe).canonicalized
+ intercept[SparkException] {
+ RewriteWithExpression.applyForExpression(expr)
+ }
+ }
+
+ test("applyForExpression rejects current_time referenced more than once") {
+ // CurrentTime is not a leaf and its only leaf is a literal precision, so the structural check
+ // alone would accept it.
+ val expr = With(CurrentTime()) { case Seq(ref) =>
+ EqualTo(ref, ref)
+ }
+ intercept[SparkException] {
+ RewriteWithExpression.applyForExpression(expr)
+ }
+ }
+
+ test("applyForExpression rejects a TIME -> TIMESTAMP cast referenced more than once") {
+ // The cast fills its date fields from the current date, so it is not safe to duplicate even
+ // though its only leaf is a literal.
+ val expr = With(Cast(Literal(0L, TimeType(6)), TimestampNTZType, Some("UTC"))) {
+ case Seq(ref) => EqualTo(ref, ref)
+ }
+ intercept[SparkException] {
+ RewriteWithExpression.applyForExpression(expr)
+ }
+ }
+
+ test("applyForExpression rejects a TIME -> TIMESTAMP_LTZ cast referenced more than once") {
+ // Same TimestampType target as the inlined string cast below; only the TIME source makes it
+ // unsafe, so the node-level check must key on the source type, not the target.
+ val expr = With(Cast(Literal(0L, TimeType(6)), TimestampType, Some("UTC"))) {
+ case Seq(ref) => EqualTo(ref, ref)
+ }
+ intercept[SparkException] {
+ RewriteWithExpression.applyForExpression(expr)
+ }
+ }
+
+ test("applyForExpression inlines a non-TIME timestamp-target cast referenced more than once") {
+ // A string -> TIMESTAMP cast carries the CAST_TO_TIMESTAMP tree pattern (keyed on target type
+ // alone) but is an ordinary deterministic cast, so it must inline rather than be rejected.
+ val cast = Cast(Literal("1970-01-01"), TimestampType, Some("UTC"))
+ val expr = With(cast) { case Seq(ref) =>
+ EqualTo(ref, ref)
+ }
+ assert(RewriteWithExpression.applyForExpression(expr) == EqualTo(cast, cast))
+ }
+
+ test("applyForExpression inlines an inner def referencing a foldable outer def") {
+ val outer = With(Literal(1)) { case Seq(outerRef) =>
+ With(outerRef + Literal(1)) { case Seq(innerRef) =>
+ innerRef + innerRef
+ }
+ }
+ val rewritten = RewriteWithExpression.applyForExpression(outer)
+ assert(!rewritten.exists(_.isInstanceOf[With]))
+ assert(!rewritten.exists(_.isInstanceOf[CommonExpressionRef]))
+ val inlinedInner = Literal(1) + Literal(1)
+ assert(rewritten == inlinedInner + inlinedInner)
+ }
+
+ test("applyForExpression rejects an inner def referencing a per-row outer def") {
+ // A ref is only safe to duplicate when its definition is: here the outer definition is an
+ // attribute, so duplicating `a + 1` would evaluate it once per reference.
+ val a = testRelation.output.head
+ val outer = With(a) { case Seq(outerRef) =>
+ With(outerRef + Literal(1)) { case Seq(innerRef) =>
+ innerRef + innerRef
+ }
+ }
+ intercept[SparkException] {
+ RewriteWithExpression.applyForExpression(outer)
+ }
+ }
+
+ test("applyForExpression handles nested With") {
+ val inner = With(Literal(1)) { case Seq(ref) =>
+ ref * ref
+ }
+ val outer = With(inner + Literal(2)) { case Seq(ref) =>
+ ref + ref
+ }
+ val rewritten = RewriteWithExpression.applyForExpression(outer)
+ assert(!rewritten.exists(_.isInstanceOf[With]))
+ assert(!rewritten.exists(_.isInstanceOf[CommonExpressionRef]))
+ val inlinedInner = Literal(1) * Literal(1) + Literal(2)
+ assert(rewritten == inlinedInner + inlinedInner)
+ }
+
+ test("applyForExpression leaves expressions without With unchanged") {
+ val a = testRelation.output.head
+ assert(RewriteWithExpression.applyForExpression(a + a) == a + a)
+ }
+
+ test("applyForExpression defers a ref to an outer common expression to the enclosing With") {
+ // The inner `With` is rewritten first, so the outer ref it references is not yet in scope
+ // and must be left for the enclosing `With` to inline.
+ val outer = With(Literal(1)) { case Seq(outerRef) =>
+ With(Literal(2)) { case Seq(innerRef) =>
+ outerRef + innerRef
+ }
+ }
+ val rewritten = RewriteWithExpression.applyForExpression(outer)
+ assert(!rewritten.exists(_.isInstanceOf[With]))
+ assert(!rewritten.exists(_.isInstanceOf[CommonExpressionRef]))
+ assert(rewritten == Literal(1) + Literal(2))
+ }
+
+ test("applyForExpression rejects current_timestamp that ComputeCurrentTime has not folded") {
+ // current_timestamp() is foldable, but each eval re-reads the clock, so inlining it at two
+ // references could produce two different values.
+ val expr = With(CurrentTimestamp()) { case Seq(ref) =>
+ EqualTo(ref, ref)
+ }
+ intercept[SparkException] {
+ RewriteWithExpression.applyForExpression(expr)
+ }
+ }
+
+ test("applyForExpression inlines current_timestamp folded to one shared literal") {
+ val instant = Instant.now()
+ val expr = With(CurrentTimestamp()) { case Seq(ref) =>
+ EqualTo(ref, ref)
+ }
+ val expected = Literal.create(DateTimeUtils.instantToMicros(instant), TimestampType)
+ assert(finishAnalysisAndInline(expr, instant) == EqualTo(expected, expected))
+ }
+
+ test("applyForExpression inlines current_database folded by ReplaceCurrentLike") {
+ val expr = With(CurrentDatabase()) { case Seq(ref) =>
+ EqualTo(ref, ref)
+ }
+ val expected = Literal.create(catalogManager.currentNamespace.quoted, StringType)
+ assert(finishAnalysisAndInline(expr, Instant.now()) == EqualTo(expected, expected))
+ }
+
+ test("applyForExpression inlines a special datetime value folded by SpecialDatetimeValues") {
+ val expr = With(Cast(Literal("epoch"), DateType, Some("UTC"))) { case Seq(ref) =>
+ EqualTo(ref, ref)
+ }
+ val expected = Literal(0, DateType)
+ assert(finishAnalysisAndInline(expr, Instant.now()) == EqualTo(expected, expected))
+ }
+
+ test("applyForExpression inlines a TIME -> TIMESTAMP cast stabilized by ComputeCurrentTime") {
+ // ComputeCurrentTime rewrites the cast to a pure makeTimestamp* builder StaticInvoke anchored
+ // on a query-stable current-date literal. Even though ConstantFolding (not part of this
+ // pipeline) has not folded it to a single literal, inlining it at both references is safe.
+ val expr = With(Cast(Literal(0L, TimeType(6)), TimestampNTZType, Some("UTC"))) {
+ case Seq(ref) => EqualTo(ref, ref)
+ }
+ val rewritten = finishAnalysisAndInline(expr, Instant.now())
+ assert(!rewritten.exists(_.isInstanceOf[With]))
+ val EqualTo(left, right) = rewritten: @unchecked
+ assert(ComputeCurrentTime.isMakeTimestampBuilder(left))
+ assert(left == right)
+ }
+
+ test("applyForExpression inlines a stabilized TIME -> TIMESTAMP_LTZ cast") {
+ // The LTZ target reaches isMakeTimestampBuilder through the same gate as the NTZ case above.
+ val expr = With(Cast(Literal(0L, TimeType(6)), TimestampType, Some("UTC"))) {
+ case Seq(ref) => EqualTo(ref, ref)
+ }
+ val rewritten = finishAnalysisAndInline(expr, Instant.now())
+ assert(!rewritten.exists(_.isInstanceOf[With]))
+ val EqualTo(left, right) = rewritten: @unchecked
+ assert(ComputeCurrentTime.isMakeTimestampBuilder(left))
+ assert(left == right)
+ }
+
test("non-cheap common expression") {
val a = testRelation.output.head
val expr = With(a + a) { case Seq(ref) =>
@@ -250,6 +507,20 @@ class RewriteWithExpressionSuite extends PlanTest {
)
}
+ test("WITH in a conditional branch referencing an outer common expression") {
+ val a = testRelation.output.head
+ // The conditional branch holds a nested `With` referencing the outer common expression.
+ // Both are inlined; the outer ref must survive the inner rewrite.
+ val expr = With(a + a) { case Seq(outerRef) =>
+ Coalesce(Seq(a, With(a * a) { case Seq(innerRef) =>
+ outerRef + innerRef
+ }))
+ }
+ val plan = testRelation.select(expr.as("col"))
+ val inlinedExpr = Coalesce(Seq(a, (a + a) + (a * a)))
+ comparePlans(Optimizer.execute(plan), testRelation.select(inlinedExpr.as("col")))
+ }
+
test("WITH expression in grouping exprs") {
val a = testRelation.output.head
val expr1 = With(a + 1) { case Seq(ref) =>
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SpecialDatetimeValuesSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SpecialDatetimeValuesSuite.scala
index 12dd8c9ae0f9..d140e99a975c 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SpecialDatetimeValuesSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SpecialDatetimeValuesSuite.scala
@@ -58,6 +58,23 @@ class SpecialDatetimeValuesSuite extends PlanTest {
}
}
+ test("applyForExpression rewrites a foldable special datetime cast to a literal") {
+ val zoneId = ZoneId.systemDefault()
+ val cast = Cast(Literal("epoch"), DateType, Some(zoneId.getId))
+ assert(SpecialDatetimeValues.applyForExpression(cast) == Literal(0, DateType))
+ }
+
+ test("applyForExpression folds special timestamp values") {
+ val cast = Cast(Literal("epoch"), TimestampType, Some("UTC"))
+ assert(SpecialDatetimeValues.applyForExpression(cast) == Literal(0L, TimestampType))
+ }
+
+ test("applyForExpression leaves non-special casts and other expressions unchanged") {
+ val cast = Cast(Literal("2021-01-01"), DateType, Some("UTC"))
+ assert(SpecialDatetimeValues.applyForExpression(cast) == cast)
+ assert(SpecialDatetimeValues.applyForExpression(Literal(1)) == Literal(1))
+ }
+
private def testSpecialTs(tsType: AtomicType, expected: Set[Long], zoneId: ZoneId): Unit = {
val in = Project(Seq(
Alias(Cast(Literal("epoch"), tsType, Some(zoneId.getId)), "epoch")(),
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org