Github user pepinoflo commented on a diff in the pull request:
https://github.com/apache/spark/pull/21208#discussion_r185958014
--- Diff:
sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala ---
@@ -798,6 +798,111 @@ class DataFrameFunctionsSuite extends QueryTest with
SharedSQLContext {
}
}
+ test("array_repeat function") {
+ val strDF = Seq(
+ ("hi", 1),
+ (null, 2)
+ ).toDF("a", "b")
+
+ checkAnswer(
+ strDF.select(array_repeat(strDF("a"), 0)),
+ Seq(
+ Row(Seq[String]()),
+ Row(Seq[String]())
+ ))
+
+ checkAnswer(
+ strDF.select(array_repeat(strDF("a"), 1)),
+ Seq(
+ Row(Seq("hi")),
+ Row(Seq(null))
+ ))
+
+ checkAnswer(
+ strDF.select(array_repeat(strDF("a"), 2)),
+ Seq(
+ Row(Seq("hi", "hi")),
+ Row(Seq(null, null))
+ ))
+
+ checkAnswer(
+ strDF.select(array_repeat(strDF("a"), strDF("b"))),
+ Seq(
+ Row(Seq("hi")),
+ Row(Seq(null, null))
+ ))
+
+ checkAnswer(
+ strDF.selectExpr("array_repeat(a, 2)"),
+ Seq(
+ Row(Seq("hi", "hi")),
+ Row(Seq(null, null))
+ ))
+
+ checkAnswer(
+ strDF.selectExpr("array_repeat(a, b)"),
+ Seq(
+ Row(Seq("hi")),
+ Row(Seq(null, null))
+ ))
+
+ val intDF = Seq(
+ (1, 1),
+ (3, 2)
+ ).toDF("a", "b")
+
+ checkAnswer(
+ intDF.select(array_repeat(intDF("a"), 0)),
+ Seq(
+ Row(Seq[Int]()),
+ Row(Seq[Int]())
+ ))
+
+ checkAnswer(
+ intDF.select(array_repeat(intDF("a"), 1)),
+ Seq(
+ Row(Seq(1)),
+ Row(Seq(3))
+ ))
+
+ checkAnswer(
+ intDF.select(array_repeat(intDF("a"), 2)),
+ Seq(
+ Row(Seq(1, 1)),
+ Row(Seq(3, 3))
+ ))
+
+ checkAnswer(
+ intDF.select(array_repeat(intDF("a"), intDF("b"))),
+ Seq(
+ Row(Seq(1)),
+ Row(Seq(3, 3))
+ ))
+
+ checkAnswer(
+ intDF.selectExpr("array_repeat(a, 2)"),
+ Seq(
+ Row(Seq(1, 1)),
+ Row(Seq(3, 3))
+ ))
+
+ checkAnswer(
+ intDF.selectExpr("array_repeat(a, b)"),
+ Seq(
+ Row(Seq(1)),
+ Row(Seq(3, 3))
+ ))
+
+ val nullDF = Seq(
+ ("hi", null)
+ ).toDF("a", "b")
+
+ intercept[AnalysisException] {
--- End diff --
Thanks, this is actually the test that confused me for null handling.
Fixed in commit 596a54f
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
