dongjoon-hyun commented on code in PR #60163:
URL: https://github.com/apache/spark/pull/60163#discussion_r4494957281
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala:
##########
@@ -230,21 +230,51 @@ object ConstantPropagation extends Rule[LogicalPlan] {
private def safeToReplace(ar: AttributeReference, nullIsFalse: Boolean) =
!ar.nullable || nullIsFalse
+ /**
+ * Replaces attributes in the comparisons of `condition` with their constant
values.
+ *
+ * The constant is the value of an attribute only on the rows that pass the
whole condition.
+ * The rewritten condition also runs on other rows, e.g. a row with `y = 1`
for
+ * `10 / (id - y) > 0 AND y = 0`, and a part of it that becomes constant is
folded with no row
+ * at all, e.g. `10 / y` for `y = 0 AND id > 10 / y`. So that this raises no
error the original
+ * condition cannot raise, an expression that can fail keeps its attributes,
unless all of them
+ * are replaced and the result evaluates without an error.
+ */
private def replaceConstants(
condition: Expression,
equalityPredicates: AttributeMap[(Literal, BinaryComparison)]):
Expression = {
val constantsMap = AttributeMap(equalityPredicates.map { case (attr, (lit,
_)) => attr -> lit })
val predicates = equalityPredicates.values.map(_._2).toSet
def replaceInComparison(b: BinaryComparison): Expression = {
lazy val collationSafeReplacement = isSameCollationAttrRefComparison(b)
- b transform {
- case a: AttributeReference
- if isBinaryStable(a.dataType) || collationSafeReplacement =>
- constantsMap.getOrElse(a, a)
+ def canReplace(a: Attribute): Boolean =
+ constantsMap.contains(a) && (isBinaryStable(a.dataType) ||
collationSafeReplacement)
+ def replace(e: Expression): Expression = e match {
+ case a: AttributeReference if canReplace(a) => constantsMap(a)
+ case _ if ExprUtils.isTotalOperator(e) => e.mapChildren(replace)
+ case _ if e.references.nonEmpty && e.references.forall(canReplace) =>
+ val replaced = e.transform { case a: AttributeReference =>
constantsMap(a) }
Review Comment:
Thank you. Confirmed: with `INT` columns, `t1.a = 1 AND t1.b = (SELECT
max(t2.b) FROM t2 WHERE t2.a = t1.a)` fails with
`java.util.NoSuchElementException: key not found` during the optimization. With
`BIGINT` columns, `t1.a = 1` is still `t1.a = CAST(1 AS BIGINT)` in the first
iteration, and the subquery becomes a join before the rule sees a literal.
Fixed in the latest commit by replacing only the attributes that `canReplace`
accepts, so the subquery keeps its condition and is not folded. Added
`SPARK-60983: Propagate a constant into a comparison with a correlated scalar
subquery` to `SubquerySuite`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala:
##########
@@ -230,21 +230,51 @@ object ConstantPropagation extends Rule[LogicalPlan] {
private def safeToReplace(ar: AttributeReference, nullIsFalse: Boolean) =
!ar.nullable || nullIsFalse
+ /**
+ * Replaces attributes in the comparisons of `condition` with their constant
values.
+ *
+ * The constant is the value of an attribute only on the rows that pass the
whole condition.
+ * The rewritten condition also runs on other rows, e.g. a row with `y = 1`
for
+ * `10 / (id - y) > 0 AND y = 0`, and a part of it that becomes constant is
folded with no row
+ * at all, e.g. `10 / y` for `y = 0 AND id > 10 / y`. So that this raises no
error the original
+ * condition cannot raise, an expression that can fail keeps its attributes,
unless all of them
+ * are replaced and the result evaluates without an error.
+ */
private def replaceConstants(
condition: Expression,
equalityPredicates: AttributeMap[(Literal, BinaryComparison)]):
Expression = {
val constantsMap = AttributeMap(equalityPredicates.map { case (attr, (lit,
_)) => attr -> lit })
val predicates = equalityPredicates.values.map(_._2).toSet
def replaceInComparison(b: BinaryComparison): Expression = {
lazy val collationSafeReplacement = isSameCollationAttrRefComparison(b)
- b transform {
- case a: AttributeReference
- if isBinaryStable(a.dataType) || collationSafeReplacement =>
- constantsMap.getOrElse(a, a)
+ def canReplace(a: Attribute): Boolean =
+ constantsMap.contains(a) && (isBinaryStable(a.dataType) ||
collationSafeReplacement)
+ def replace(e: Expression): Expression = e match {
+ case a: AttributeReference if canReplace(a) => constantsMap(a)
+ case _ if ExprUtils.isTotalOperator(e) => e.mapChildren(replace)
+ case _ if e.references.nonEmpty && e.references.forall(canReplace) =>
+ val replaced = e.transform { case a: AttributeReference =>
constantsMap(a) }
+ if (canFold(replaced)) replaced else e
+ case _ => e
Review Comment:
Thank you. Confirmed: `d <> 0 AND n / d > 1 AND d = 5` returns `(10, 5)`
without this PR, and fails with `DIVIDE_BY_ZERO` with it. In the latest commit,
when a conjunct that can fail still refers to a propagated attribute, the rule
moves the equalities ahead of the other conjuncts, so the query returns `(10,
5)` again. An equality between an attribute and a literal cannot fail, so this
raises no new error. Conjuncts that cannot fail, such as `isnotnull(d)`, do not
trigger it, and no `PlanStabilitySuite` golden file changes. `10 / (id - y) > 0
AND y = 0` on `(5, 5)` also returns no rows again, as without this PR.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala:
##########
@@ -230,21 +230,51 @@ object ConstantPropagation extends Rule[LogicalPlan] {
private def safeToReplace(ar: AttributeReference, nullIsFalse: Boolean) =
!ar.nullable || nullIsFalse
+ /**
+ * Replaces attributes in the comparisons of `condition` with their constant
values.
+ *
+ * The constant is the value of an attribute only on the rows that pass the
whole condition.
+ * The rewritten condition also runs on other rows, e.g. a row with `y = 1`
for
+ * `10 / (id - y) > 0 AND y = 0`, and a part of it that becomes constant is
folded with no row
+ * at all, e.g. `10 / y` for `y = 0 AND id > 10 / y`. So that this raises no
error the original
+ * condition cannot raise, an expression that can fail keeps its attributes,
unless all of them
+ * are replaced and the result evaluates without an error.
+ */
private def replaceConstants(
condition: Expression,
equalityPredicates: AttributeMap[(Literal, BinaryComparison)]):
Expression = {
val constantsMap = AttributeMap(equalityPredicates.map { case (attr, (lit,
_)) => attr -> lit })
val predicates = equalityPredicates.values.map(_._2).toSet
def replaceInComparison(b: BinaryComparison): Expression = {
lazy val collationSafeReplacement = isSameCollationAttrRefComparison(b)
- b transform {
- case a: AttributeReference
- if isBinaryStable(a.dataType) || collationSafeReplacement =>
- constantsMap.getOrElse(a, a)
+ def canReplace(a: Attribute): Boolean =
+ constantsMap.contains(a) && (isBinaryStable(a.dataType) ||
collationSafeReplacement)
+ def replace(e: Expression): Expression = e match {
+ case a: AttributeReference if canReplace(a) => constantsMap(a)
+ case _ if ExprUtils.isTotalOperator(e) => e.mapChildren(replace)
+ case _ if e.references.nonEmpty && e.references.forall(canReplace) =>
+ val replaced = e.transform { case a: AttributeReference =>
constantsMap(a) }
+ if (canFold(replaced)) replaced else e
+ case _ => e
}
+ b.mapChildren(replace)
+ }
+ def replaceInCondition(e: Expression): Expression = e match {
+ case b: BinaryComparison if predicates.contains(b) => b
+ case b: BinaryComparison => replaceInComparison(b)
+ case _ if ExprUtils.isTotalOperator(e) =>
e.mapChildren(replaceInCondition)
Review Comment:
Thank you. As you mentioned, this is not introduced by this PR: `(y < 3 OR
10 / x > 1) AND y = 5` and the `HAVING` query fail the same way without it. The
reordering above does not change them either, because the rewritten operand no
longer refers to `y`. Since this needs its own policy for the operands of `OR`,
`AND`, `COALESCE` and `IN`, I'd like to handle it in a separate JIRA issue.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala:
##########
@@ -230,21 +230,51 @@ object ConstantPropagation extends Rule[LogicalPlan] {
private def safeToReplace(ar: AttributeReference, nullIsFalse: Boolean) =
!ar.nullable || nullIsFalse
+ /**
+ * Replaces attributes in the comparisons of `condition` with their constant
values.
+ *
+ * The constant is the value of an attribute only on the rows that pass the
whole condition.
+ * The rewritten condition also runs on other rows, e.g. a row with `y = 1`
for
+ * `10 / (id - y) > 0 AND y = 0`, and a part of it that becomes constant is
folded with no row
+ * at all, e.g. `10 / y` for `y = 0 AND id > 10 / y`. So that this raises no
error the original
+ * condition cannot raise, an expression that can fail keeps its attributes,
unless all of them
+ * are replaced and the result evaluates without an error.
+ */
private def replaceConstants(
condition: Expression,
equalityPredicates: AttributeMap[(Literal, BinaryComparison)]):
Expression = {
val constantsMap = AttributeMap(equalityPredicates.map { case (attr, (lit,
_)) => attr -> lit })
val predicates = equalityPredicates.values.map(_._2).toSet
def replaceInComparison(b: BinaryComparison): Expression = {
lazy val collationSafeReplacement = isSameCollationAttrRefComparison(b)
- b transform {
- case a: AttributeReference
- if isBinaryStable(a.dataType) || collationSafeReplacement =>
- constantsMap.getOrElse(a, a)
+ def canReplace(a: Attribute): Boolean =
+ constantsMap.contains(a) && (isBinaryStable(a.dataType) ||
collationSafeReplacement)
+ def replace(e: Expression): Expression = e match {
+ case a: AttributeReference if canReplace(a) => constantsMap(a)
+ case _ if ExprUtils.isTotalOperator(e) => e.mapChildren(replace)
+ case _ if e.references.nonEmpty && e.references.forall(canReplace) =>
+ val replaced = e.transform { case a: AttributeReference =>
constantsMap(a) }
+ if (canFold(replaced)) replaced else e
+ case _ => e
}
+ b.mapChildren(replace)
+ }
+ def replaceInCondition(e: Expression): Expression = e match {
+ case b: BinaryComparison if predicates.contains(b) => b
+ case b: BinaryComparison => replaceInComparison(b)
+ case _ if ExprUtils.isTotalOperator(e) =>
e.mapChildren(replaceInCondition)
+ // The children of these include the comparisons of a pulled up
correlated condition.
+ case _: InSubquery | _: ListQuery | _: Exists =>
e.mapChildren(replaceInCondition)
+ case _ => e
Review Comment:
Thank you. Added an `IF` case to the test in the latest commit. It fails if
the rule rewrites the comparison in the condition of the `IF`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala:
##########
@@ -230,21 +230,51 @@ object ConstantPropagation extends Rule[LogicalPlan] {
private def safeToReplace(ar: AttributeReference, nullIsFalse: Boolean) =
!ar.nullable || nullIsFalse
+ /**
+ * Replaces attributes in the comparisons of `condition` with their constant
values.
+ *
+ * The constant is the value of an attribute only on the rows that pass the
whole condition.
+ * The rewritten condition also runs on other rows, e.g. a row with `y = 1`
for
+ * `10 / (id - y) > 0 AND y = 0`, and a part of it that becomes constant is
folded with no row
+ * at all, e.g. `10 / y` for `y = 0 AND id > 10 / y`. So that this raises no
error the original
+ * condition cannot raise, an expression that can fail keeps its attributes,
unless all of them
+ * are replaced and the result evaluates without an error.
+ */
private def replaceConstants(
condition: Expression,
equalityPredicates: AttributeMap[(Literal, BinaryComparison)]):
Expression = {
val constantsMap = AttributeMap(equalityPredicates.map { case (attr, (lit,
_)) => attr -> lit })
val predicates = equalityPredicates.values.map(_._2).toSet
def replaceInComparison(b: BinaryComparison): Expression = {
lazy val collationSafeReplacement = isSameCollationAttrRefComparison(b)
- b transform {
- case a: AttributeReference
- if isBinaryStable(a.dataType) || collationSafeReplacement =>
- constantsMap.getOrElse(a, a)
+ def canReplace(a: Attribute): Boolean =
+ constantsMap.contains(a) && (isBinaryStable(a.dataType) ||
collationSafeReplacement)
+ def replace(e: Expression): Expression = e match {
+ case a: AttributeReference if canReplace(a) => constantsMap(a)
+ case _ if ExprUtils.isTotalOperator(e) => e.mapChildren(replace)
+ case _ if e.references.nonEmpty && e.references.forall(canReplace) =>
+ val replaced = e.transform { case a: AttributeReference =>
constantsMap(a) }
+ if (canFold(replaced)) replaced else e
Review Comment:
Thank you. In the latest commit, `canFold` became `fold`, which returns the
folded literal, so the constant is evaluated once.
--
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]
