explode support array and map

apache · Dec 12, 2023 · 2127eb0 · 2127eb0
1 parent 8c6f164
commit 2127eb0
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 19 deletions.
diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/ValidatorApiImpl.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/ValidatorApiImpl.scala
@@ -23,10 +23,10 @@ import io.glutenproject.substrait.plan.PlanNode
 import io.glutenproject.validate.NativePlanValidationInfo
 import io.glutenproject.vectorized.NativePlanEvaluator
 
-import org.apache.spark.sql.catalyst.expressions.{CreateMap, Explode, Expression, Generator, JsonTuple, Literal, PosExplode}
+import org.apache.spark.sql.catalyst.expressions.{Expression, Generator, JsonTuple, PosExplode}
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
 import org.apache.spark.sql.execution.SparkPlan
-import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampType}
+import org.apache.spark.sql.types._
 
 class ValidatorApiImpl extends ValidatorApi {
 
@@ -95,24 +95,11 @@ class ValidatorApiImpl extends ValidatorApi {
       return ValidationResult.notOk(s"Velox backend does not support outer")
     }
     generator match {
-      case generator: JsonTuple =>
+      case _: JsonTuple =>
         ValidationResult.notOk(s"Velox backend does not support this json_tuple")
-      case generator: PosExplode =>
+      case _: PosExplode =>
         // TODO(yuan): support posexplode and remove this check
         ValidationResult.notOk(s"Velox backend does not support this posexplode")
-      case explode: Explode if (explode.child.isInstanceOf[CreateMap]) =>
-        // explode(MAP(col1, col2))
-        ValidationResult.notOk(s"Velox backend does not support MAP datatype")
-      case explode: Explode if (explode.child.isInstanceOf[Literal]) =>
-        // explode(ARRAY(1, 2, 3))
-        ValidationResult.notOk(s"Velox backend does not support literal Array datatype")
-      case explode: Explode =>
-        explode.child.dataType match {
-          case _: MapType =>
-            ValidationResult.notOk(s"Velox backend does not support MAP datatype")
-          case _ =>
-            ValidationResult.ok
-        }
       case _ =>
         ValidationResult.ok
     }

diff --git a/backends-velox/src/test/scala/io/glutenproject/execution/TestOperator.scala b/backends-velox/src/test/scala/io/glutenproject/execution/TestOperator.scala
@@ -675,4 +675,17 @@ class TestOperator extends VeloxWholeStageTransformerSuite with AdaptiveSparkPla
       }
     }
   }
+
+  test("test explode function") {
+    runQueryAndCompare("""
+                         |SELECT explode(array(1, 2, 3));
+                         |""".stripMargin) {
+      checkOperatorMatch[GenerateExecTransformer]
+    }
+    runQueryAndCompare("""
+                         |SELECT explode(map(1, "a", 2, "b"));
+                         |""".stripMargin) {
+      checkOperatorMatch[GenerateExecTransformer]
+    }
+  }
 }
diff --git a/gluten-core/src/main/scala/io/glutenproject/execution/GenerateExecTransformer.scala b/gluten-core/src/main/scala/io/glutenproject/execution/GenerateExecTransformer.scala
@@ -18,7 +18,7 @@ package io.glutenproject.execution
 
 import io.glutenproject.backendsapi.BackendsApiManager
 import io.glutenproject.exception.GlutenException
-import io.glutenproject.expression.{ConverterUtils, ExpressionConverter, ExpressionTransformer}
+import io.glutenproject.expression.{ConverterUtils, ExpressionConverter}
 import io.glutenproject.extension.ValidationResult
 import io.glutenproject.metrics.MetricsUpdater
 import io.glutenproject.substrait.`type`.TypeBuilder
@@ -69,7 +69,7 @@ case class GenerateExecTransformer(
     val operatorId = context.nextOperatorId(this.nodeName)
     val generatorExpr =
       ExpressionConverter.replaceWithExpressionTransformer(generator, child.output)
-    val generatorNode = generatorExpr.asInstanceOf[ExpressionTransformer].doTransform(args)
+    val generatorNode = generatorExpr.doTransform(args)
     val childOutputNodes = new java.util.ArrayList[ExpressionNode]
     for (target <- requiredChildOutput) {
       val found = child.output.zipWithIndex.filter(_._1.name == target.name)