Skip to content

[SPARK-32941][SQL] Optimize UpdateFields expression chain and put the rule early in Analysis phase #29812

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
viirya wants to merge 6 commits into from
Closed

[SPARK-32941][SQL] Optimize UpdateFields expression chain and put the rule early in Analysis phase #29812

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
74cf2dd
Optimize WithFields expression chain.
viirya Sep 20, 2020
00acff9
Use resolver.
viirya Sep 22, 2020
cb8872c
Address comment.
viirya Sep 22, 2020
82ad8c8
Merge remote-tracking branch 'upstream/master' into SPARK-32941
viirya Oct 18, 2020
38bdefd
Simplify UpdateFields in analysis too.
viirya Oct 18, 2020
f41900c
Skip the rule if possible.
viirya Oct 19, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,11 @@ object SimplifyExtractValueOps extends Rule[LogicalPlan] {
// `$"struct_col".withField("b", lit(1)).withField("b", lit(2)).getField("b")`
// we want to return `lit(2)` (and not `lit(1)`).
val expr = matches.last._2
If(IsNull(struct), Literal(null, expr.dataType), expr)
if (struct.nullable) {
If(IsNull(struct), Literal(null, expr.dataType), expr)
} else {
expr
}
} else {
GetStructField(struct, ordinal, maybeName)
}
Expand Down
2 changes: 1 addition & 1 deletion sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
RemoveRedundantAliases,
UnwrapCastInBinaryComparison,
RemoveNoopOperators,
CombineWithFields,
OptimizeWithFields,
SimplifyExtractValueOps,
CombineConcats) ++
extendedOperatorOptimizationRules
Expand Down
22 changes: 19 additions & 3 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/WithFields.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,32 @@

package org.apache.spark.sql.catalyst.optimizer

import org.apache.spark.sql.catalyst.expressions.WithFields
import scala.collection.mutable

import org.apache.spark.sql.catalyst.expressions.{Expression, WithFields}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.internal.SQLConf


/**
* Combines all adjacent [[WithFields]] expression into a single [[WithFields]] expression.
* Optimizes [[WithFields]] expression chains.
*/
object CombineWithFields extends Rule[LogicalPlan] {
object OptimizeWithFields extends Rule[LogicalPlan] {
lazy val resolver = SQLConf.get.resolver

def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
case WithFields(structExpr, names, values) if names.distinct.length != names.length =>

@fqaiser94 fqaiser94 Sep 21, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could this case statement be after the next case statement? So that we combine the chains first before deduplicating?

@viirya viirya Sep 22, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't run this rule just once, so the order should be fine.

val newNames = mutable.ArrayBuffer.empty[String]
val newValues = mutable.ArrayBuffer.empty[Expression]
names.zip(values).reverse.foreach { case (name, value) =>
if (newNames.find(resolver(_, name)).isEmpty) {

@cloud-fan cloud-fan Sep 22, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is a bit inefficient. Shall we build a set with lowercased names if case sensitivity is false?

@viirya viirya Sep 22, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a set for case-sensitive case.

newNames += name
newValues += value
}
}
WithFields(structExpr, names = newNames.reverse.toSeq, valExprs = newValues.reverse.toSeq)

@fqaiser94 fqaiser94 Sep 21, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For my understanding, can you explain how we expect to benefit from this optimization?

I ask because we do this kind of deduplication inside of WithFields already as part of the foldLeft operation here. It will only keep the last valExpr for each name. So I think the optimized logical plan will be the same with or without this optimization in all scenarios? CMIIW

@viirya viirya Sep 22, 2020 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right. It is eventually the same. But for some cases, before we extend WithFields, the expression tree might be very complex. This is coming from improving scalability of #29587. This is applied during I fixed the scalability issue. I found this is useful to reduce the complex of WithFields expression tree.

I will run these rules in #29587 to simplify expression tree before entering optimizer.

@fqaiser94 fqaiser94 Sep 23, 2020 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, so I took a look at the PR you linked and left a related comment there. I don't think you actually need this optimization for #29587

This optimization is only useful if someone uses WithFields to update the same field multiple times. However, it would simply be better to not update the same field multiple times. At the very least, we should not do this when we re-use this Expression internally within Spark.

Unfortunately, "bad" end-users might still update the same field multiple times. Assuming we should optimize for such users (not sure), since this batch is only applied half-way through the optimization cycle anyway, I think we could just move up the Batch("ReplaceWithFieldsExpression", Once, ReplaceWithFieldsExpression) to get the same benefit (which is just simplified tree). What do you reckon?

@viirya viirya Sep 23, 2020 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually I'd like to run these rules to simplify WithFields tree early in analysis stage. During fixing scale issue of #29587, I thought that it is very likely to write bad WithFields tree. Once hitting that, it is very hard to debug and the analyzer/optimizer spend a lot of time traversing expression tree. So I think it is very useful keep this rule to simplify the expression tree, but I don't think we want to do ReplaceWithFieldsExpression in analysis stage.

@fqaiser94 fqaiser94 Sep 23, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ahh I see, yes, in the analysis stage this would likely be helpful!

Okay in that case, could this PR wait till #29795 goes in? I'm refactoring WithFields so this optimization would need to change anyway.

@viirya viirya Sep 23, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm fine to wait until #29795.


case WithFields(WithFields(struct, names1, valExprs1), names2, valExprs2) =>
WithFields(struct, names1 ++ names2, valExprs1 ++ valExprs2)
}
Expand Down
43 changes: 40 additions & 3 deletions ...st/optimizer/CombineWithFieldsSuite.scala → ...t/optimizer/OptimizeWithFieldsSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,19 +19,21 @@ package org.apache.spark.sql.catalyst.optimizer

import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.dsl.plans._
import org.apache.spark.sql.catalyst.expressions.{Alias, Literal, WithFields}
import org.apache.spark.sql.catalyst.expressions.{Alias, GetStructField, Literal, WithFields}
import org.apache.spark.sql.catalyst.plans.PlanTest
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.rules._


class CombineWithFieldsSuite extends PlanTest {
class OptimizeWithFieldsSuite extends PlanTest {

object Optimize extends RuleExecutor[LogicalPlan] {
val batches = Batch("CombineWithFields", FixedPoint(10), CombineWithFields) :: Nil
val batches = Batch("OptimizeWithFields", FixedPoint(10),
OptimizeWithFields, SimplifyExtractValueOps) :: Nil
}

private val testRelation = LocalRelation('a.struct('a1.int))
private val testRelation2 = LocalRelation('a.struct('a1.int).notNull)

test("combines two WithFields") {
val originalQuery = testRelation
Expand Down Expand Up @@ -73,4 +75,39 @@ class CombineWithFieldsSuite extends PlanTest {

comparePlans(optimized, correctAnswer)
}

test("SPARK-32941: optimize WithFields followed by GetStructField") {
val originalQuery = testRelation2
.select(Alias(
GetStructField(WithFields(
'a,
Seq("b1"),
Seq(Literal(4))), 1), "out")())

val optimized = Optimize.execute(originalQuery.analyze)
val correctAnswer = testRelation2
.select(Alias(Literal(4), "out")())
.analyze

comparePlans(optimized, correctAnswer)
}

test("SPARK-32941: optimize WithFields chain") {
val originalQuery = testRelation
.select(Alias(
WithFields(
WithFields(
'a,
Seq("b1"),
Seq(Literal(4))),
Seq("b1"),
Seq(Literal(5))), "out")())

val optimized = Optimize.execute(originalQuery.analyze)
val correctAnswer = testRelation
.select(Alias(WithFields('a, Seq("b1"), Seq(Literal(5))), "out")())
.analyze

comparePlans(optimized, correctAnswer)
}
}