vladimirg-db commented on code in PR #57629:
URL: https://github.com/apache/spark/pull/57629#discussion_r3728480327
##########
sql/core/src/test/resources/sql-tests/inputs/distinct-map-aggregates.sql:
##########
@@ -0,0 +1,95 @@
+-- Test DISTINCT aggregates with MapType arguments.
+
+CREATE OR REPLACE TEMPORARY VIEW distinct_map_data AS SELECT * FROM VALUES
+ (2, map('a', 1, 'b', 2), 1, true),
+ (2, map('b', 2, 'a', 1), 1, true),
+ (1, map('a', 1, 'b', 2), 1, true),
+ (1, map('a', 3), 2, false)
+AS distinct_map_data(g, m, id, should_keep);
+
+SELECT COUNT(DISTINCT m) FROM distinct_map_data;
+
+SELECT SIZE(COLLECT_LIST(DISTINCT m)) FROM distinct_map_data;
+
+SELECT map_entries(m)
+FROM (
+ SELECT EXPLODE(COLLECT_LIST(DISTINCT m)) AS m
+ FROM distinct_map_data
+) AS collected_maps
+ORDER BY element_at(m, 'a');
+
+SELECT map_entries(FIRST(DISTINCT m)), map_entries(LAST(DISTINCT m)),
COUNT(DISTINCT m)
+FROM VALUES (map('b', 2, 'a', 1)) AS single_map_data(m);
+
+SELECT COUNT(DISTINCT m, id) FROM distinct_map_data;
+
+SELECT COUNT(DISTINCT m), COUNT(DISTINCT id) FROM distinct_map_data;
+
+SELECT g, COUNT(DISTINCT m)
+FROM distinct_map_data
+GROUP BY g
+ORDER BY g;
+
+SELECT m, COUNT(DISTINCT m), COLLECT_LIST(DISTINCT m)
+FROM distinct_map_data
+GROUP BY m
+ORDER BY element_at(m, 'a');
+
+SELECT COUNT(DISTINCT m) FILTER (WHERE should_keep) FROM distinct_map_data;
+
+SELECT MAX(map_values(m)[0])
+FROM distinct_map_data
+WHERE id = 1;
+
+SELECT MAX(map_values(m)[0]), COUNT(DISTINCT m)
+FROM distinct_map_data
+WHERE id = 1;
+
+SELECT g
+FROM distinct_map_data
+GROUP BY g
+ORDER BY COUNT(DISTINCT m), g;
+
+SELECT g
+FROM distinct_map_data
+GROUP BY g
+HAVING COUNT(DISTINCT m) = 1
+ORDER BY g;
+
+SELECT COUNT(DISTINCT named_struct('m', m)) FROM distinct_map_data;
+
+SELECT COUNT(DISTINCT array(m)) FROM distinct_map_data;
+
+SELECT COUNT(DISTINCT map('m', m)) FROM distinct_map_data;
+
+SELECT COUNT(DISTINCT m), COLLECT_LIST(DISTINCT m)
+FROM VALUES
+ (CAST(map() AS MAP<STRING, INT>)),
+ (CAST(map() AS MAP<STRING, INT>)),
+ (CAST(NULL AS MAP<STRING, INT>))
+AS null_and_empty_map_data(m);
+
+SELECT g, GROUPING(g), COUNT(DISTINCT m)
+FROM distinct_map_data
+GROUP BY GROUPING SETS ((g), ())
+ORDER BY GROUPING(g), g;
+
+SELECT COUNT(DISTINCT named_struct('m', m, 'n', n))
+FROM VALUES
+ (map('a', 1, 'b', 2), map('x', 1, 'y', 2)),
+ (map('b', 2, 'a', 1), map('y', 2, 'x', 1))
+AS grouped_distinct_map_data(m, n)
+GROUP BY m;
+
+SET spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled=false;
+
+SELECT COUNT(DISTINCT m), COLLECT_LIST(DISTINCT m) FROM distinct_map_data;
Review Comment:
Addressed in ac48274ceee. The flag-off case now keeps COUNT(DISTINCT m)
separate, explodes COLLECT_LIST(DISTINCT m), renders each value with
map_entries, and orders the rows by map contents. This makes the legacy
map-entry order visible without asserting nondeterministic COLLECT_LIST array
order.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortInAggregateSuite.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions.{Alias, MapSort}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LocalRelation,
LogicalPlan, Project}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StringType}
+
+class InsertMapSortInAggregateSuite extends PlanTest {
+ private val input = LocalRelation(Symbol("m").map(StringType, IntegerType))
+ private val mapAttribute = input.output.head
+
+ private def aliasesNamed(plan: LogicalPlan, name: String): Seq[Alias] = {
+ plan.flatMap { node =>
+ node.expressions.flatMap(_.collect {
+ case alias @ Alias(_, aliasName) if aliasName == name => alias
+ })
+ }
+ }
+
+ test("reuse map sort when a grouping key is also a distinct argument") {
+ val plan = Aggregate(
+ Seq(mapAttribute),
+ Seq(mapAttribute, countDistinct(mapAttribute).as("count")),
+ input)
+ val rewritten = InsertMapSortInAggregate(plan)
+ val groupingAliases = aliasesNamed(rewritten, "_groupingmapsort")
+
+ assert(groupingAliases.size == 1)
+ assert(groupingAliases.head.child.isInstanceOf[MapSort])
+ rewritten match {
+ case Aggregate(Seq(groupingExpression), aggregateExpressions, _:
Project, _) =>
+
assert(groupingExpression.semanticEquals(groupingAliases.head.toAttribute))
+ val distinctChildren = aggregateExpressions.flatMap(_.collect {
+ case expression: AggregateExpression if expression.isDistinct =>
+ expression.aggregateFunction.children
+ }).flatten
+ assert(distinctChildren.size == 1)
+
assert(distinctChildren.head.semanticEquals(groupingAliases.head.toAttribute))
+ case other =>
+ fail(s"Unexpected plan:\n$other")
+ }
+ }
+
+ test("project complex distinct arguments only when needed") {
+ val attributePlan = Aggregate(
+ Nil,
+ Seq(countDistinct(mapAttribute).as("count")),
+ input)
+ val complexPlan = Aggregate(
+ Nil,
+ Seq(countDistinct(namedStruct("m", mapAttribute)).as("count")),
+ input)
+
+ val rewrittenAttributePlan = InsertMapSortInAggregate(attributePlan)
+ assert(rewrittenAttributePlan.collect { case _: Project => 1 }.size == 1)
+ assert(aliasesNamed(rewrittenAttributePlan,
"_distinctaggregateexpression").isEmpty)
+ assert(aliasesNamed(rewrittenAttributePlan, "_distinctmapsort").size == 1)
+
+ val rewrittenComplexPlan = InsertMapSortInAggregate(complexPlan)
+ assert(rewrittenComplexPlan.collect { case _: Project => 1 }.size == 2)
+ assert(aliasesNamed(rewrittenComplexPlan,
"_distinctaggregateexpression").size == 1)
+ assert(aliasesNamed(rewrittenComplexPlan, "_distinctmapsort").size == 1)
+ }
+
+ test("skip distinct argument normalization when disabled") {
+ val plan = Aggregate(
+ Nil,
+ Seq(countDistinct(mapAttribute).as("count")),
+ input)
+
+ withSQLConf(SQLConf.INSERT_MAP_SORT_IN_DISTINCT_AGGREGATES_ENABLED.key ->
"false") {
+ comparePlans(InsertMapSortInAggregate(plan), plan)
+ }
+ }
+}
Review Comment:
Addressed in ac48274ceee. Added a PlanTest with a map-free Aggregate and a
map-carrying Aggregate under Union; it compares the map-free aggregate
unchanged, asserts it has no Project, and confirms the map aggregate is
rewritten.
--
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]
