Github user HyukjinKwon commented on a diff in the pull request:
https://github.com/apache/spark/pull/20938#discussion_r179970154
--- Diff:
sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala ---
@@ -413,6 +413,86 @@ class DataFrameFunctionsSuite extends QueryTest with
SharedSQLContext {
)
}
+ test("flatten function") {
+ val oneRowDF = Seq((1, "a", Seq(1, 2, 3))).toDF("i", "s", "arr")
+
+ // Test cases with a primitive type
+ val intDF = Seq(
+ (Seq(Seq(1, 2, 3), Seq(4, 5), Seq(6))),
+ (Seq(Seq(1, 2))),
+ (Seq(Seq(1), Seq.empty)),
+ (Seq(Seq.empty, Seq(1))),
+ (Seq(Seq.empty, Seq.empty)),
+ (Seq(Seq(1), null)),
+ (Seq(null, Seq(1))),
+ (Seq(null, null))
+ ).toDF("i")
+
+ val intDFResult = Seq(
+ Row(Seq(1, 2, 3, 4, 5, 6)),
+ Row(Seq(1, 2)),
+ Row(Seq(1)),
+ Row(Seq(1)),
+ Row(Seq.empty),
+ Row(null),
+ Row(null),
+ Row(null)
+ )
+
+ checkAnswer(intDF.select(flatten($"i")), intDFResult)
+ checkAnswer(intDF.selectExpr("flatten(i)"), intDFResult)
+ checkAnswer(
+ oneRowDF.selectExpr("flatten(array(arr, array(null, 5), array(6,
null)))"),
+ Seq(Row(Seq(1, 2, 3, null, 5, 6, null)))
+ )
--- End diff --
Let's move it up.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
