address comments

Author: mgaido91

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala

@@ -141,6 +141,22 @@ trait UnaryHigherOrderFunction extends HigherOrderFunction with ExpectsInputType
   def expectingFunctionType: AbstractDataType = AnyDataType
 
   @transient lazy val functionForEval: Expression = functionsForEval.head
+
+  /**
+   * Called by [[eval]]. If a subclass keeps the default nullability, it can override this method
+   * in order to save null-check code.
+   */
+  protected def nullSafeEval(inputRow: InternalRow, input: Any): Any =
+    sys.error(s"UnaryHigherOrderFunction must override either eval or nullSafeEval")
+
+  override def eval(inputRow: InternalRow): Any = {
+    val value = input.eval(inputRow)
+    if (value == null) {
+      null
+    } else {
+      nullSafeEval(inputRow, value)
+    }
+  }
 }
 
 trait ArrayBasedUnaryHigherOrderFunction extends UnaryHigherOrderFunction {
@@ -199,24 +215,20 @@ case class ArrayTransform(
     (elementVar, indexVar)
   }
 
-  override def eval(input: InternalRow): Any = {
-    val arr = this.input.eval(input).asInstanceOf[ArrayData]
-    if (arr == null) {
-      null
-    } else {
-      val f = functionForEval
-      val result = new GenericArrayData(new Array[Any](arr.numElements))
-      var i = 0
-      while (i < arr.numElements) {
-        elementVar.value.set(arr.get(i, elementVar.dataType))
-        if (indexVar.isDefined) {
-          indexVar.get.value.set(i)
-        }
-        result.update(i, f.eval(input))
-        i += 1
+  override def nullSafeEval(inputRow: InternalRow, inputValue: Any): Any = {
+    val arr = inputValue.asInstanceOf[ArrayData]
+    val f = functionForEval
+    val result = new GenericArrayData(new Array[Any](arr.numElements))
+    var i = 0
+    while (i < arr.numElements) {
+      elementVar.value.set(arr.get(i, elementVar.dataType))
+      if (indexVar.isDefined) {
+        indexVar.get.value.set(i)
       }
-      result
+      result.update(i, f.eval(inputRow))
+      i += 1
     }
+    result
   }
 
   override def prettyName: String = "transform"
@@ -259,23 +271,20 @@ case class MapFilter(
 
   override def nullable: Boolean = input.nullable
 
-  override def eval(input: InternalRow): Any = {
-    val m = this.input.eval(input).asInstanceOf[MapData]
-    if (m == null) {
-      null
-    } else {
-      val retKeys = new mutable.ListBuffer[Any]
-      val retValues = new mutable.ListBuffer[Any]
-      m.foreach(keyType, valueType, (k, v) => {
-        keyVar.value.set(k)
-        valueVar.value.set(v)
-        if (functionForEval.eval(input).asInstanceOf[Boolean]) {
-          retKeys += k
-          retValues += v
-        }
-      })
-      ArrayBasedMapData(retKeys.toArray, retValues.toArray)
-    }
+  override def nullSafeEval(inputRow: InternalRow, value: Any): Any = {
+    val m = value.asInstanceOf[MapData]
+    val f = functionForEval
+    val retKeys = new mutable.ListBuffer[Any]
+    val retValues = new mutable.ListBuffer[Any]
+    m.foreach(keyType, valueType, (k, v) => {
+      keyVar.value.set(k)
+      valueVar.value.set(v)
+      if (f.eval(inputRow).asInstanceOf[Boolean]) {
+        retKeys += k
+        retValues += v
+      }
+    })
+    ArrayBasedMapData(retKeys.toArray, retValues.toArray)
   }
 
   override def dataType: DataType = input.dataType

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HigherOrderFunctionsSuite.scala

@@ -112,11 +112,11 @@ class HigherOrderFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper
     checkEvaluation(mapFilter(mii1, kGreaterThanV), Map())
     checkEvaluation(mapFilter(miin, kGreaterThanV), null)
 
-    val valueNull: (Expression, Expression) => Expression = (_, v) => v.isNull
+    val valueIsNull: (Expression, Expression) => Expression = (_, v) => v.isNull
 
-    checkEvaluation(mapFilter(mii0, valueNull), Map())
-    checkEvaluation(mapFilter(mii1, valueNull), Map(1 -> null, 3 -> null))
-    checkEvaluation(mapFilter(miin, valueNull), null)
+    checkEvaluation(mapFilter(mii0, valueIsNull), Map())
+    checkEvaluation(mapFilter(mii1, valueIsNull), Map(1 -> null, 3 -> null))
+    checkEvaluation(mapFilter(miin, valueIsNull), null)
 
     val msi0 = Literal.create(Map("abcdf" -> 5, "abc" -> 10, "" -> 0),
       MapType(StringType, IntegerType, valueContainsNull = false))