support nested keys in GpuMapConcat (#6290)

Signed-off-by: remzi <[email protected]> Signed-off-by: remzi <[email protected]>
NVIDIA · Aug 26, 2022 · 8f4b4cb · 8f4b4cb
1 parent 3929002
commit 8f4b4cb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 21 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala
@@ -3202,10 +3202,6 @@ object GpuOverrides extends Logging {
       }),
     expr[MapConcat](
       "Returns the union of all the given maps",
-      // Currently, GpuMapConcat supports nested values but not nested keys.
-      // We will add the nested key support after
-      // cuDF can fully support nested types in lists::drop_list_duplicates.
-      // Issue link: https://github.com/rapidsai/cudf/issues/11093
       ExprChecks.projectOnly(TypeSig.MAP.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 +
           TypeSig.NULL + TypeSig.ARRAY + TypeSig.STRUCT + TypeSig.MAP),
         TypeSig.MAP.nested(TypeSig.all),
@@ -3214,13 +3210,6 @@ object GpuOverrides extends Logging {
           TypeSig.NULL + TypeSig.ARRAY + TypeSig.STRUCT + TypeSig.MAP),
           TypeSig.MAP.nested(TypeSig.all)))),
       (a, conf, p, r) => new ComplexTypeMergingExprMeta[MapConcat](a, conf, p, r) {
-        override def tagExprForGpu(): Unit = {
-          a.dataType.keyType match {
-            case MapType(_,_,_) | ArrayType(_,_) | StructType(_) => willNotWorkOnGpu(
-              s"GpuMapConcat does not currently support the key type ${a.dataType.keyType}.")
-            case _ =>
-          }
-        }
         override def convertToGpu(child: Seq[Expression]): GpuExpression = GpuMapConcat(child)
       }),
     expr[ConcatWs](

diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/CollectionOpSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/CollectionOpSuite.scala
@@ -19,22 +19,18 @@ package com.nvidia.spark.rapids
 import org.apache.spark.sql.functions.map_concat
 
 class CollectionOpSuite extends SparkQueryCompareTestSuite {
-  testGpuFallback(
-    "MapConcat with Array keys fall back",
-    "ProjectExec",
-    ArrayKeyMapDF, 
-    execsAllowedNonGpu = Seq("ProjectExec", "ShuffleExchangeExec")) {
+  testSparkResultsAreEqual(
+    "MapConcat with Array keys",
+    ArrayKeyMapDF) {
     frame => {
       import frame.sparkSession.implicits._
       frame.select(map_concat($"col1", $"col2"))
     }
   }
 
-  testGpuFallback(
-    "MapConcat with Struct keys fall back",
-    "ProjectExec",
-    StructKeyMapDF, 
-    execsAllowedNonGpu = Seq("ProjectExec", "ShuffleExchangeExec")) {
+   testSparkResultsAreEqual(
+    "MapConcat with Struct keys",
+    StructKeyMapDF) {
     frame => {
       import frame.sparkSession.implicits._
       frame.select(map_concat($"col1", $"col2"))