[SPARK-34638][SQL] Single field nested column prune on generator output by viirya · Pull Request #31966 · apache/spark

viirya · 2021-03-26T02:15:27Z

What changes were proposed in this pull request?

This patch proposes an improvement on nested column pruning if the pruning target is generator's output. Previously we disallow such case. This patch allows to prune on it if there is only one single nested column is accessed after Generate.

E.g., df.select(explode($"items").as('item)).select($"item.itemId"). As we only need itemId from item, we can prune other fields out and only keep itemId.

In this patch, we only address explode-like generators. We will address other generators in followups.

Why are the changes needed?

This helps to extend the availability of nested column pruning.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Unit test

SparkQA · 2021-03-26T03:33:10Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41125/

SparkQA · 2021-03-26T03:41:31Z

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41125/

SparkQA · 2021-03-26T08:22:43Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41134/

SparkQA · 2021-03-26T08:31:15Z

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41134/

SparkQA · 2021-03-26T10:17:58Z

Test build #136549 has finished for PR 31966 at commit 4758f96.

This patch fails Spark unit tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2021-03-26T11:05:21Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41137/

SparkQA · 2021-03-26T12:28:23Z

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41137/

SparkQA · 2021-03-26T14:03:48Z

Test build #136553 has finished for PR 31966 at commit 7908788.

This patch fails SparkR unit tests.
This patch merges cleanly.
This patch adds no public classes.

viirya · 2021-03-26T16:27:42Z

cc @cloud-fan @maropu @dongjoon-hyun

dongjoon-hyun

Thank you for pinging me, @viirya . I agree that we can handle multiple field case later.

For now I only add single field support. Multiple field case would be more complicated. We can consider to add it later.

maropu

This behaivour (the col name that is internally auto-generated is visible to a user) is expected one?

scala> sql("select __auto_generated_subquery_name.col.a from (select explode(items) from t)").show()
+---+
| a|
+---+
| 1|
| 2|
+---+

maropu · 2021-03-29T01:44:20Z

- NestedColumnAliasing.getAliasSubMap(exprsToPrune, g.qualifiedGeneratorOutput).map {
+ NestedColumnAliasing.getAliasSubMap(exprsToPrune).map {
 case (nestedFieldToAlias, attrToAliases) =>
 // Defer updating `Generate.unrequiredChildIndex` to next round of `ColumnPruning`.


Move this comment into L275-276?

viirya · 2021-04-05T07:21:44Z

This behaivour (the col name that is internally auto-generated is visible to a user) is expected one?
scala> sql("select __auto_generated_subquery_name.col.a from (select explode(items) from t)").show()
+---+
| a|
+---+
| 1|
| 2|
+---+

I can see we use it in test:

spark/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

Line 2788 in 0494dc9

 checkAnswer(sql("SELECT __auto_generated_subquery_name.i from (SELECT i FROM v)"), Row(1)) 

viirya · 2021-04-05T07:32:01Z

Addressed above comments except for #31966 (comment). For case-sensitive test case, I need to figure out a suspicious bug in nested column pruning first.

SparkQA · 2021-04-05T08:46:04Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41484/

SparkQA · 2021-04-05T08:46:06Z

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41484/

SparkQA · 2021-04-05T12:19:30Z

Test build #136907 has finished for PR 31966 at commit fe286df.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

viirya · 2021-04-10T21:00:27Z

@maropu @dongjoon-hyun Thanks for your review and patience. I added test-sensitivity test. Once you have time, please take another look. Thanks.

SparkQA · 2021-04-10T22:06:53Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41748/

SparkQA · 2021-04-10T22:06:54Z

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41748/

SparkQA · 2021-04-11T01:48:25Z

Test build #137170 has finished for PR 31966 at commit 6c9d839.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2021-04-15T08:43:29Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41982/

SparkQA · 2021-04-15T08:43:30Z

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/41982/

SparkQA · 2021-04-15T12:32:22Z

Test build #137407 has finished for PR 31966 at commit ad3d191.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

maropu · 2021-04-19T01:28:03Z

Looks fine otherwise. I think it's better to have @dongjoon-hyun check this before merging.

viirya · 2021-04-19T01:29:25Z

Thanks @maropu!

dongjoon-hyun · 2021-04-19T04:32:43Z

+
+ // As we change the child of the generator, its output data type must be updated.
+ val updatedGeneratorOutput = rewrittenG.generatorOutput
+ .zip(rewrittenG.generator.elementSchema.toAttributes)


Indentation? It seems that two-space is enough in this case.

dongjoon-hyun · 2021-04-19T04:42:00Z

+ newAttr.withExprId(oldAttr.exprId).withName(oldAttr.name)
+ }
+ assert(updatedGeneratorOutput.length == rewrittenG.generatorOutput.length,
+ "Updated generator output must have same length as original generator output.")


Maybe, same length as -> the same length with?

dongjoon-hyun · 2021-04-19T05:53:08Z

Hi, @viirya . While testing this PR, I found the following regression.

BEFORE (3.1.1)

scala> sql("select * from values array(array(named_struct('a', 1, 'b', 3), named_struct('a', 2, 'b', 4))) T(items)").write.parquet("/tmp/nested_array")
scala> spark.read.parquet("/tmp/nested_array").createOrReplaceTempView("t")
scala> sql("select d.a from (select explode(c) d from (select explode(items) c from t))").show()
+---+| a|+---+|1||2|+---+

BEFORE(master)

scala> spark.read.parquet("/tmp/nested_array").createOrReplaceTempView("t")
scala> sql("select d.a from (select explode(c) d from (select explode(items) c from t))").show()
+---+| a|+---+|1||2|+---+

AFTER (This PR)

scala> spark.read.parquet("/tmp/nested_array").createOrReplaceTempView("t")
scala> sql("select d.a from (select explode(c) d from (select explode(items) c from t))").show()
java.lang.ClassCastException: org.apache.spark.sql.types.ArrayType cannot be cast to org.apache.spark.sql.types.StructType
at org.apache.spark.sql.catalyst.expressions.SelectedField$.selectField(SelectedField.scala:81)
at org.apache.spark.sql.catalyst.expressions.SelectedField$.unapply(SelectedField.scala:62)
at org.apache.spark.sql.catalyst.expressions.SchemaPruning$.getRootFields(SchemaPruning.scala:124)
at org.apache.spark.sql.catalyst.expressions.SchemaPruning$.$anonfun$identifyRootFields$1(SchemaPruning.scala:81)

Could you double-check this and add some test coverage?

viirya · 2021-04-19T06:45:03Z

Thanks @dongjoon-hyun. Let me take a look.

viirya · 2021-04-25T00:47:16Z

@dongjoon-hyun Sorry for late. I addressed the cases and added test for it.

SparkQA · 2021-04-25T02:04:49Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/42423/

SparkQA · 2021-04-25T02:04:50Z

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/42423/

dongjoon-hyun · 2021-04-25T02:49:10Z

No problem. Thank you for updates, @viirya .

dongjoon-hyun · 2021-04-25T02:51:20Z

+ // It is because we don't allow field extractor on two-level array,
+ // i.e., attr.field when attr is a ArrayType(ArrayType(...)).
+ // Similarily, we also cannot push through if the child of generator is `MapType`.
+ g.generator.children.head.dataType match {


Got it. Let me play more with this PR for a while. It seems I need more tests.

SparkQA · 2021-04-25T05:46:38Z

Test build #137898 has finished for PR 31966 at commit a719409.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

dongjoon-hyun

+1, LGTM. Thank you, @viirya and @maropu .

viirya · 2021-04-26T16:29:40Z

Thank you @dongjoon-hyun @maropu

viirya · 2021-04-26T16:31:40Z

Merging to master.

Nested column prune on generator output for one field.
5221be3

github-actionsBot added the SQL label Mar 26, 2021

This comment has been minimized.
Sign in to view

viirya added 2 commits March 25, 2021 23:57

Add e2e test cases.
4758f96

Fix wrong transform logic.
7908788

viirya commented Mar 26, 2021
View reviewed changes

Comment threadsql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/NestedColumnAliasing.scala Outdated