[SPARK-31977][SQL] Returns the plan directly from NestedColumnAliasing

### What changes were proposed in this pull request?

This proposes a minor refactoring to match `NestedColumnAliasing` to `GeneratorNestedColumnAliasing` so it returns the pruned plan directly.

```scala
    case p  NestedColumnAliasing(nestedFieldToAlias, attrToAliases) =>
      NestedColumnAliasing.replaceToAliases(p, nestedFieldToAlias, attrToAliases)
```

vs

```scala
    case GeneratorNestedColumnAliasing(p) => p
```

### Why are the changes needed?

Just for readability.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Existing tests should cover.

Closes #28812 from HyukjinKwon/SPARK-31977.

Authored-by: HyukjinKwon <gurwls223@apache.org>
Signed-off-by: Takeshi Yamamuro <yamamuro@apache.org>
This commit is contained in:
HyukjinKwon 2020-06-13 07:26:37 +09:00 committed by Takeshi Yamamuro
parent 78d08a8c38
commit a620a2a7e5
2 changed files with 13 additions and 9 deletions

View file

@ -30,16 +30,21 @@ import org.apache.spark.sql.types._
*/
object NestedColumnAliasing {
def unapply(plan: LogicalPlan)
: Option[(Map[ExtractValue, Alias], Map[ExprId, Seq[Alias]])] = plan match {
def unapply(plan: LogicalPlan): Option[LogicalPlan] = plan match {
case Project(projectList, child)
if SQLConf.get.nestedSchemaPruningEnabled && canProjectPushThrough(child) =>
val exprCandidatesToPrune = projectList ++ child.expressions
getAliasSubMap(exprCandidatesToPrune, child.producedAttributes.toSeq)
getAliasSubMap(exprCandidatesToPrune, child.producedAttributes.toSeq).map {
case (nestedFieldToAlias, attrToAliases) =>
NestedColumnAliasing.replaceToAliases(plan, nestedFieldToAlias, attrToAliases)
}
case plan if SQLConf.get.nestedSchemaPruningEnabled && canPruneOn(plan) =>
val exprCandidatesToPrune = plan.expressions
getAliasSubMap(exprCandidatesToPrune, plan.producedAttributes.toSeq)
case p if SQLConf.get.nestedSchemaPruningEnabled && canPruneOn(p) =>
val exprCandidatesToPrune = p.expressions
getAliasSubMap(exprCandidatesToPrune, p.producedAttributes.toSeq).map {
case (nestedFieldToAlias, attrToAliases) =>
NestedColumnAliasing.replaceToAliases(p, nestedFieldToAlias, attrToAliases)
}
case _ => None
}
@ -47,7 +52,7 @@ object NestedColumnAliasing {
/**
* Replace nested columns to prune unused nested columns later.
*/
def replaceToAliases(
private def replaceToAliases(
plan: LogicalPlan,
nestedFieldToAlias: Map[ExtractValue, Alias],
attrToAliases: Map[ExprId, Seq[Alias]]): LogicalPlan = plan match {

View file

@ -649,8 +649,7 @@ object ColumnPruning extends Rule[LogicalPlan] {
// Can't prune the columns on LeafNode
case p @ Project(_, _: LeafNode) => p
case p @ NestedColumnAliasing(nestedFieldToAlias, attrToAliases) =>
NestedColumnAliasing.replaceToAliases(p, nestedFieldToAlias, attrToAliases)
case NestedColumnAliasing(p) => p
// for all other logical plans that inherits the output from it's children
// Project over project is handled by the first case, skip it here.