c21 commented on a change in pull request #31595:
URL: https://github.com/apache/spark/pull/31595#discussion_r580825530
##########
File path:
sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala
##########
@@ -808,6 +813,53 @@ class DataFrameSetOperationsSuite extends QueryTest with
SharedSparkSession {
// scalastyle:on
checkAnswer(union, row1 :: row2 :: Nil)
}
+
+ test("SPARK-34474: Remove unnecessary Union under Distinct") {
+ Seq(RemoveNoopUnion.ruleName, "").map { ruleName =>
+ withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ruleName) {
+ val distinctUnionDF1 = testData.union(testData).distinct()
+ checkAnswer(distinctUnionDF1, testData.distinct())
+
+
+ val distinctUnionDF2 =
testData.union(testData).dropDuplicates(Seq("key"))
+ checkAnswer(distinctUnionDF2, testData.dropDuplicates(Seq("key")))
+
+ val distinctUnionDF3 = sql(
+ """
+ |select key, value from testData
+ |union
+ |select key, value from testData
+ |""".stripMargin)
+ checkAnswer(distinctUnionDF3, testData.distinct())
+
+ val distinctUnionDF4 = sql(
+ """
+ |select distinct key, expr
+ |from
+ |(
+ | select key, key + 1 as expr
+ | from testData
+ | union all
+ | select key, key + 2 as expr
+ | from testData
+ |)
Review comment:
After checking the optimized logical plan of this query, I found the
`collect` in `RemoveNoopUnion` relies on the fact that the query has one extra
unnecessary `Project` on one child:
```
Distinct
+- Union false, false
:- Project [key#3, (key#3 + 1) AS expr#143]
: +- SerializeFromObject [knownnotnull(assertnotnull(input[0,
org.apache.spark.sql.test.SQLTestData$TestData, true])).key AS key#3,
staticinvoke(class org.apache.spark.unsafe.types.UTF8String, StringType,
fromString, knownnotnull(assertnotnull(input[0,
org.apache.spark.sql.test.SQLTestData$TestData, true])).value, true, false) AS
value#4]
: +- ExternalRDD [obj#2]
+- Project [key#3 AS key#145, expr#144 AS expr#146]. <<-- this one
+- Project [key#3, (key#3 + 2) AS expr#144]
+- SerializeFromObject [knownnotnull(assertnotnull(input[0,
org.apache.spark.sql.test.SQLTestData$TestData, true])).key AS key#3,
staticinvoke(class org.apache.spark.unsafe.types.UTF8String, StringType,
fromString, knownnotnull(assertnotnull(input[0,
org.apache.spark.sql.test.SQLTestData$TestData, true])).value, true, false) AS
value#4]
+- ExternalRDD [obj#2]
```
So the rules works because there's one unnecessary `Project` existed before
the rule. And later on this `Project` is removed by `CollapseProject`. If this
`Project` does not exist, the correctness of the `RemoveNoopUnion` rule cannot
be guaranteed.
I feel this is dangerous and may lead to wrong query plan later. But I am
also fine that we guard this with a unit test here.
btw, because of the unnecessary `Project` in the middle, this query cannot
be optimized by the rule:
```
select distinct key, expr
from
(
select key, key as expr
from testData
union all
select key, key as expr
from testData
)
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
