Add an array_filter_nulls method

mrpowers-io · Oct 15, 2019 · 9482818 · 9482818
1 parent 993b6df
commit 9482818
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/src/main/scala/com/github/mrpowers/spark/daria/sql/functions.scala b/src/main/scala/com/github/mrpowers/spark/daria/sql/functions.scala
@@ -555,6 +555,14 @@ object functions {
     }
   }
 
+  def array_filter_nulls[T: TypeTag]() = udf[Option[Seq[T]], Seq[T]] { arr: Seq[T] =>
+    if (arr == null) {
+      null
+    } else {
+      Some(arr.filter(_ != null))
+    }
+  }
+
   def array_map_ex_null[T: TypeTag](f: T => T) = udf[Option[Seq[T]], Seq[T]] { arr: Seq[T] =>
     if (arr == null) {
       null

diff --git a/src/test/scala/com/github/mrpowers/spark/daria/sql/FunctionsTest.scala b/src/test/scala/com/github/mrpowers/spark/daria/sql/FunctionsTest.scala
@@ -960,6 +960,38 @@ object FunctionsTest extends TestSuite with DataFrameComparer with ColumnCompare
 
     }
 
+    'array_filter_nulls - {
+
+      "removes all nulls from arrays" - {
+
+        val df = spark
+          .createDF(
+            List(
+              (Array("snake", null, "cool"), Array("snake", "cool")),
+              (Array("cat", null, null), Array("cat")),
+              (Array(null, null), Array[String]()),
+              (null, null)
+            ),
+            List(
+              ("words", ArrayType(StringType, true), true),
+              ("expected", ArrayType(StringType, true), true)
+            )
+          )
+          .withColumn(
+            "words_filtered",
+            functions.array_filter_nulls[String]().apply(col("words"))
+          )
+
+        assertColumnEquality(
+          df,
+          "words_filtered",
+          "expected"
+        )
+
+      }
+
+    }
+
     'array_map_ex_null - {
 
       "works like the Scala map method on Strings" - {