You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2024/08/02 06:28:36 UTC
(spark) branch master updated: [SPARK-49093][SQL] GROUP BY with MapType nested inside complex type
This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new aca0d24c8724 [SPARK-49093][SQL] GROUP BY with MapType nested inside complex type
aca0d24c8724 is described below
commit aca0d24c8724f1de4e65ad6c663b8e2255dbccd0
Author: Nebojsa Savic <ne...@databricks.com>
AuthorDate: Fri Aug 2 14:28:20 2024 +0800
[SPARK-49093][SQL] GROUP BY with MapType nested inside complex type
### What changes were proposed in this pull request?
Currently we are supporting GROUP BY <column>, where column is of MapType, but we don't support scenarios where column type contains MapType nested in some complex type (i.e. ARRAY<MAP<INT,INT>>), this PR addresses this issue.
### Why are the changes needed?
We are extending support for MapType columns in GROUP BY clause.
### Does this PR introduce _any_ user-facing change?
Customer would be able to use MapType nested in complex type as part of GROUP BY clause.
### How was this patch tested?
Added tests.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #47331 from nebojsa-db/SC-170296.
Authored-by: Nebojsa Savic <ne...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../InsertMapSortInGroupingExpressions.scala | 53 +++++++-
.../sql-tests/analyzer-results/group-by.sql.out | 130 +++++++++++++++++++
.../test/resources/sql-tests/inputs/group-by.sql | 64 +++++++++
.../resources/sql-tests/results/group-by.sql.out | 143 +++++++++++++++++++++
4 files changed, 386 insertions(+), 4 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortInGroupingExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortInGroupingExpressions.scala
index 3d883fa9d9ae..51f6ca374909 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortInGroupingExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortInGroupingExpressions.scala
@@ -17,11 +17,12 @@
package org.apache.spark.sql.catalyst.optimizer
-import org.apache.spark.sql.catalyst.expressions.MapSort
+import org.apache.spark.sql.catalyst.expressions.{ArrayTransform, CreateNamedStruct, Expression, GetStructField, If, IsNull, LambdaFunction, Literal, MapFromArrays, MapKeys, MapSort, MapValues, NamedLambdaVariable}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.AGGREGATE
-import org.apache.spark.sql.types.MapType
+import org.apache.spark.sql.types.{ArrayType, MapType, StructType}
+import org.apache.spark.util.ArrayImplicits.SparkArrayOps
/**
* Adds MapSort to group expressions containing map columns, as the key/value paris need to be
@@ -34,12 +35,56 @@ object InsertMapSortInGroupingExpressions extends Rule[LogicalPlan] {
_.containsPattern(AGGREGATE), ruleId) {
case a @ Aggregate(groupingExpr, _, _) =>
val newGrouping = groupingExpr.map { expr =>
- if (!expr.isInstanceOf[MapSort] && expr.dataType.isInstanceOf[MapType]) {
- MapSort(expr)
+ if (!expr.exists(_.isInstanceOf[MapSort])
+ && expr.dataType.existsRecursively(_.isInstanceOf[MapType])) {
+ insertMapSortRecursively(expr)
} else {
expr
}
}
a.copy(groupingExpressions = newGrouping)
}
+
+ /*
+ Inserts MapSort recursively taking into account when
+ it is nested inside a struct or array.
+ */
+ private def insertMapSortRecursively(e: Expression): Expression = {
+ e.dataType match {
+ case m: MapType =>
+ // Check if value type of MapType contains MapType (possibly nested)
+ // and special handle this case.
+ val mapSortExpr = if (m.valueType.existsRecursively(_.isInstanceOf[MapType])) {
+ MapFromArrays(MapKeys(e), insertMapSortRecursively(MapValues(e)))
+ } else {
+ e
+ }
+
+ MapSort(mapSortExpr)
+
+ case StructType(fields)
+ if fields.exists(_.dataType.existsRecursively(_.isInstanceOf[MapType])) =>
+ val struct = CreateNamedStruct(fields.zipWithIndex.flatMap { case (f, i) =>
+ Seq(Literal(f.name), insertMapSortRecursively(
+ GetStructField(e, i, Some(f.name))))
+ }.toImmutableArraySeq)
+ if (struct.valExprs.forall(_.isInstanceOf[GetStructField])) {
+ // No field needs MapSort processing, just return the original expression.
+ e
+ } else if (e.nullable) {
+ If(IsNull(e), Literal(null, struct.dataType), struct)
+ } else {
+ struct
+ }
+
+ case ArrayType(et, containsNull) if et.existsRecursively(_.isInstanceOf[MapType]) =>
+ val param = NamedLambdaVariable("x", et, containsNull)
+ val funcBody = insertMapSortRecursively(param)
+
+ ArrayTransform(e, LambdaFunction(funcBody, Seq(param)))
+
+ case _ => e
+ }
+ }
+
}
diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/group-by.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/group-by.sql.out
index 46b40e451526..8849aa445225 100644
--- a/sql/core/src/test/resources/sql-tests/analyzer-results/group-by.sql.out
+++ b/sql/core/src/test/resources/sql-tests/analyzer-results/group-by.sql.out
@@ -1190,3 +1190,133 @@ ORDER BY col1
Sort [col1#x ASC NULLS FIRST], true
+- Aggregate [col1#x], [col1#x, count(1) AS cnt#xL]
+- LocalRelation [col1#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (ARRAY(MAP(1, 2, 2, 3), MAP(1, 3))), (ARRAY(MAP(2, 3), MAP(1, 3))), (ARRAY(MAP(2, 3, 1, 2), MAP(1, 3))) as t(a)
+GROUP BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', map(1, 2, 2, 3))), (named_struct('b', map(1, 3))), (named_struct('b', map(2, 3, 1, 2))) as t(a)
+GROUP BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', map(1, 2, 2, 3))), (named_struct('b', map(1, 3))), (named_struct('b', map(2, 3, 1, 2))) as t(a)
+GROUP BY a.b
+-- !query analysis
+Aggregate [a#x.b], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', array(map(1, 2, 2, 3), map(1, 3)))), (named_struct('b', array(map(2, 3), map(1, 3)))), (named_struct('b', array(map(2, 3, 1, 2), map(1, 3)))) as t(a)
+GROUP BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (ARRAY(named_struct('b', map(1, 2, 2, 3)), named_struct('b', map(1, 3)))), (ARRAY(named_struct('b', map(2, 3)), named_struct('b', map(1, 3)))), (ARRAY(named_struct('b', map(2, 3, 1, 2)), named_struct('b', map(1, 3)))) as t(a)
+group BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (map(array(1, 2), 2, array(2, 3), 3)), (map(array(1, 3), 3)), (map(array(2, 3), 3, array(1, 2), 2)) as t(a)
+group BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (map(array(1, 2, 3), 3)), (map(array(3, 2, 1), 3)) as t(a)
+group BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (ARRAY(map(array(1, 2), 2, array(2, 3), 3))), (ARRAY(MAP(ARRAY(1, 3), 3))), (ARRAY(map(array(2, 3), 3, array(1, 2), 2))) as t(a)
+group BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (map(named_struct('b', 1), 2, named_struct('b', 2), 3)), (map(named_struct('b', 1), 3)), (map(named_struct('b', 2), 3, named_struct('b', 1), 2)) as t(a)
+group BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', map(named_struct('c', 1), 2, named_struct('c', 2), 3))), (named_struct('b', map(named_struct('c', 1), 3))), (named_struct('b', map(named_struct('c', 2), 3, named_struct('c', 1), 2))) as t(a)
+group BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', map(named_struct('c', 1), 2, named_struct('c', 2), 3))), (named_struct('b', map(named_struct('c', 1), 3))), (named_struct('b', map(named_struct('c', 2), 3, named_struct('c', 1), 2))) as t(a)
+group BY a.b
+-- !query analysis
+Aggregate [a#x.b], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (Map(1, Map(1,2), 2, Map(2, 3, 1, 2))), (Map(2, Map(1, 2, 2,3), 1, Map(1, 2))), (Map(1, Map(1,2), 2, Map(2, 4))) as t(a)
+GROUP BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (Map(1, Array(Map(1,2)), 2, Array(Map(2, 3, 1, 2)))), (Map(2, Array(Map(1, 2, 2,3)), 1, Array(Map(1, 2)))), (Map(1, Array(Map(1,2)), 2, Array(Map(2, 4)))) as t(a)
+GROUP BY a
+-- !query analysis
+Aggregate [a#x], [count(1) AS count(1)#xL]
++- SubqueryAlias t
+ +- LocalRelation [a#x]
diff --git a/sql/core/src/test/resources/sql-tests/inputs/group-by.sql b/sql/core/src/test/resources/sql-tests/inputs/group-by.sql
index 7d6116ac1e3f..6dd0adbc8722 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/group-by.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/group-by.sql
@@ -274,3 +274,67 @@ FROM VALUES
GROUP BY col1
ORDER BY col1
;
+
+-- SC-170296: Verify that group by works when MapType is inside complex type for column type
+-- ARRAY<MAP<INT,INT>>
+SELECT count(*)
+FROM VALUES (ARRAY(MAP(1, 2, 2, 3), MAP(1, 3))), (ARRAY(MAP(2, 3), MAP(1, 3))), (ARRAY(MAP(2, 3, 1, 2), MAP(1, 3))) as t(a)
+GROUP BY a;
+
+-- STRUCT<b:MAP<INT,INT>>
+SELECT count(*)
+FROM VALUES (named_struct('b', map(1, 2, 2, 3))), (named_struct('b', map(1, 3))), (named_struct('b', map(2, 3, 1, 2))) as t(a)
+GROUP BY a;
+
+SELECT count(*)
+FROM VALUES (named_struct('b', map(1, 2, 2, 3))), (named_struct('b', map(1, 3))), (named_struct('b', map(2, 3, 1, 2))) as t(a)
+GROUP BY a.b;
+
+-- STRUCT<b:ARRAY<MAP<INT,INT>>>
+SELECT count(*)
+FROM VALUES (named_struct('b', array(map(1, 2, 2, 3), map(1, 3)))), (named_struct('b', array(map(2, 3), map(1, 3)))), (named_struct('b', array(map(2, 3, 1, 2), map(1, 3)))) as t(a)
+GROUP BY a;
+
+-- ARRAY<STRUCT<b:MAP<INT,INT>>>
+SELECT count(*)
+FROM VALUES (ARRAY(named_struct('b', map(1, 2, 2, 3)), named_struct('b', map(1, 3)))), (ARRAY(named_struct('b', map(2, 3)), named_struct('b', map(1, 3)))), (ARRAY(named_struct('b', map(2, 3, 1, 2)), named_struct('b', map(1, 3)))) as t(a)
+group BY a;
+
+-- MAP<ARRAY<INT>,INT>
+SELECT count(*)
+FROM VALUES (map(array(1, 2), 2, array(2, 3), 3)), (map(array(1, 3), 3)), (map(array(2, 3), 3, array(1, 2), 2)) as t(a)
+group BY a;
+
+SELECT count(*)
+FROM VALUES (map(array(1, 2, 3), 3)), (map(array(3, 2, 1), 3)) as t(a)
+group BY a;
+
+-- ARRAY<MAP<ARRAY<INT>,INT>>
+SELECT count(*)
+FROM VALUES (ARRAY(map(array(1, 2), 2, array(2, 3), 3))), (ARRAY(MAP(ARRAY(1, 3), 3))), (ARRAY(map(array(2, 3), 3, array(1, 2), 2))) as t(a)
+group BY a;
+
+-- MAP<STRUCT<b:INT>,INT>
+SELECT count(*)
+FROM VALUES (map(named_struct('b', 1), 2, named_struct('b', 2), 3)), (map(named_struct('b', 1), 3)), (map(named_struct('b', 2), 3, named_struct('b', 1), 2)) as t(a)
+group BY a;
+
+-- STRUCT<b:MAP<STRUCT<c:INT>,INT>>
+SELECT count(*)
+FROM VALUES (named_struct('b', map(named_struct('c', 1), 2, named_struct('c', 2), 3))), (named_struct('b', map(named_struct('c', 1), 3))), (named_struct('b', map(named_struct('c', 2), 3, named_struct('c', 1), 2))) as t(a)
+group BY a;
+
+SELECT count(*)
+FROM VALUES (named_struct('b', map(named_struct('c', 1), 2, named_struct('c', 2), 3))), (named_struct('b', map(named_struct('c', 1), 3))), (named_struct('b', map(named_struct('c', 2), 3, named_struct('c', 1), 2))) as t(a)
+group BY a.b;
+
+-- Map valueType contains MapType (possibly nested)
+-- MAP<INT, MAP<INT,INT>>
+SELECT count(*)
+FROM VALUES (Map(1, Map(1,2), 2, Map(2, 3, 1, 2))), (Map(2, Map(1, 2, 2,3), 1, Map(1, 2))), (Map(1, Map(1,2), 2, Map(2, 4))) as t(a)
+GROUP BY a;
+
+-- MAP<INT, ARRAY<MAP<INT,INT>>>
+SELECT count(*)
+FROM VALUES (Map(1, Array(Map(1,2)), 2, Array(Map(2, 3, 1, 2)))), (Map(2, Array(Map(1, 2, 2,3)), 1, Array(Map(1, 2)))), (Map(1, Array(Map(1,2)), 2, Array(Map(2, 4)))) as t(a)
+GROUP BY a;
\ No newline at end of file
diff --git a/sql/core/src/test/resources/sql-tests/results/group-by.sql.out b/sql/core/src/test/resources/sql-tests/results/group-by.sql.out
index 0b7bdfd85223..d8a9f4c2e11f 100644
--- a/sql/core/src/test/resources/sql-tests/results/group-by.sql.out
+++ b/sql/core/src/test/resources/sql-tests/results/group-by.sql.out
@@ -1124,3 +1124,146 @@ struct<col1:double,cnt:bigint>
0.0 2
Infinity 2
NaN 2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (ARRAY(MAP(1, 2, 2, 3), MAP(1, 3))), (ARRAY(MAP(2, 3), MAP(1, 3))), (ARRAY(MAP(2, 3, 1, 2), MAP(1, 3))) as t(a)
+GROUP BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', map(1, 2, 2, 3))), (named_struct('b', map(1, 3))), (named_struct('b', map(2, 3, 1, 2))) as t(a)
+GROUP BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', map(1, 2, 2, 3))), (named_struct('b', map(1, 3))), (named_struct('b', map(2, 3, 1, 2))) as t(a)
+GROUP BY a.b
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', array(map(1, 2, 2, 3), map(1, 3)))), (named_struct('b', array(map(2, 3), map(1, 3)))), (named_struct('b', array(map(2, 3, 1, 2), map(1, 3)))) as t(a)
+GROUP BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (ARRAY(named_struct('b', map(1, 2, 2, 3)), named_struct('b', map(1, 3)))), (ARRAY(named_struct('b', map(2, 3)), named_struct('b', map(1, 3)))), (ARRAY(named_struct('b', map(2, 3, 1, 2)), named_struct('b', map(1, 3)))) as t(a)
+group BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (map(array(1, 2), 2, array(2, 3), 3)), (map(array(1, 3), 3)), (map(array(2, 3), 3, array(1, 2), 2)) as t(a)
+group BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (map(array(1, 2, 3), 3)), (map(array(3, 2, 1), 3)) as t(a)
+group BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+1
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (ARRAY(map(array(1, 2), 2, array(2, 3), 3))), (ARRAY(MAP(ARRAY(1, 3), 3))), (ARRAY(map(array(2, 3), 3, array(1, 2), 2))) as t(a)
+group BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (map(named_struct('b', 1), 2, named_struct('b', 2), 3)), (map(named_struct('b', 1), 3)), (map(named_struct('b', 2), 3, named_struct('b', 1), 2)) as t(a)
+group BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', map(named_struct('c', 1), 2, named_struct('c', 2), 3))), (named_struct('b', map(named_struct('c', 1), 3))), (named_struct('b', map(named_struct('c', 2), 3, named_struct('c', 1), 2))) as t(a)
+group BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (named_struct('b', map(named_struct('c', 1), 2, named_struct('c', 2), 3))), (named_struct('b', map(named_struct('c', 1), 3))), (named_struct('b', map(named_struct('c', 2), 3, named_struct('c', 1), 2))) as t(a)
+group BY a.b
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (Map(1, Map(1,2), 2, Map(2, 3, 1, 2))), (Map(2, Map(1, 2, 2,3), 1, Map(1, 2))), (Map(1, Map(1,2), 2, Map(2, 4))) as t(a)
+GROUP BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
+
+
+-- !query
+SELECT count(*)
+FROM VALUES (Map(1, Array(Map(1,2)), 2, Array(Map(2, 3, 1, 2)))), (Map(2, Array(Map(1, 2, 2,3)), 1, Array(Map(1, 2)))), (Map(1, Array(Map(1,2)), 2, Array(Map(2, 4)))) as t(a)
+GROUP BY a
+-- !query schema
+struct<count(1):bigint>
+-- !query output
+1
+2
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org