[GLUTEN-4039][VL] Add flatten function support (#5551)

apache · May 11, 2024 · d931006 · d931006
1 parent d0f611e
commit d931006
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 2 deletions.
diff --git a/backends-clickhouse/src/main/scala/org/apache/gluten/utils/CHExpressionUtil.scala b/backends-clickhouse/src/main/scala/org/apache/gluten/utils/CHExpressionUtil.scala
@@ -198,6 +198,7 @@ object CHExpressionUtil {
     UNIX_MILLIS -> DefaultValidator(),
     UNIX_MICROS -> DefaultValidator(),
     TIMESTAMP_MILLIS -> DefaultValidator(),
-    TIMESTAMP_MICROS -> DefaultValidator()
+    TIMESTAMP_MICROS -> DefaultValidator(),
+    FLATTEN -> DefaultValidator()
   )
 }
diff --git a/backends-velox/src/test/scala/org/apache/gluten/execution/ScalarFunctionsValidateSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/execution/ScalarFunctionsValidateSuite.scala
@@ -941,4 +941,24 @@ class ScalarFunctionsValidateSuite extends FunctionsValidateTest {
     }
   }
 
+  test("test flatten nested array") {
+    withTempPath {
+      path =>
+        Seq[Seq[Seq[Integer]]](
+          Seq(Seq(1, 2), Seq(4, 5)),
+          null,
+          Seq(null, Seq(1, 2)),
+          Seq(null, null),
+          Seq(Seq(1, 2, null), Seq(null, null), Seq(3, 4), Seq.empty))
+          .toDF("arrays")
+          .write
+          .parquet(path.getCanonicalPath)
+
+        spark.read.parquet(path.getCanonicalPath).createOrReplaceTempView("array_tbl")
+
+        runQueryAndCompare("select flatten(arrays) as res from array_tbl;") {
+          checkGlutenOperatorMatch[ProjectExecTransformer]
+        }
+    }
+  }
 }
diff --git a/docs/velox-backend-support-progress.md b/docs/velox-backend-support-progress.md
@@ -281,7 +281,7 @@ Gluten supports 199 functions. (Drag to right to see all data types)
 | explode_outer, explode        |                        |                       |        |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
 | filter                        | filter                 | filter                | S      | Lambda with index argument not supported |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
 | forall                        | all_match              |                       | S      |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
-| flatten                       | flatten                |                       |        |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
+| flatten                       | flatten                | flatten               | S      |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
 | map                           | map                    | map                   | S      |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
 | map_concat                    | map_concat             |                       |        |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
 | map_entries                   | map_entries            |                       |        |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |

diff --git a/gluten-core/src/main/scala/org/apache/gluten/expression/ExpressionMappings.scala b/gluten-core/src/main/scala/org/apache/gluten/expression/ExpressionMappings.scala
@@ -241,6 +241,7 @@ object ExpressionMappings {
     Sig[ArrayExists](EXISTS),
     Sig[Shuffle](SHUFFLE),
     Sig[ZipWith](ZIP_WITH),
+    Sig[Flatten](FLATTEN),
     // Map functions
     Sig[CreateMap](CREATE_MAP),
     Sig[GetMapValue](GET_MAP_VALUE),

diff --git a/shims/common/src/main/scala/org/apache/gluten/expression/ExpressionNames.scala b/shims/common/src/main/scala/org/apache/gluten/expression/ExpressionNames.scala
@@ -258,6 +258,7 @@ object ExpressionNames {
   final val TRANSFORM = "transform"
   final val SHUFFLE = "shuffle"
   final val ZIP_WITH = "zip_with"
+  final val FLATTEN = "flatten"
 
   // Map functions
   final val CREATE_MAP = "map"