[SPARK-27905] [SQL] Add higher order function 'forall'

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -449,6 +449,7 @@ object FunctionRegistry {
     expression[MapFilter]("map_filter"),
     expression[ArrayFilter]("filter"),
     expression[ArrayExists]("exists"),
+    expression[ArrayForAll]("forall"),
     expression[ArrayAggregate]("aggregate"),
     expression[TransformValues]("transform_values"),
     expression[TransformKeys]("transform_keys"),

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

@@ -445,8 +445,73 @@ case class ArrayExists(
       false
     }
   }
+}
+
+/**
+ * Tests whether a predicate holds for all elements in the array.
+ */
+@ExpressionDescription(usage =
+  "_FUNC_(expr, pred) - Tests whether a predicate holds for all elements in the array.",
+  examples = """
+    Examples:
+      > SELECT _FUNC_(array(1, 2, 3), x -> x % 2 == 0);
+       false
+      > SELECT _FUNC_(array(2, 4, 8), x -> x % 2 == 0);
+       true
+      > SELECT _FUNC_(array(1, null, 3), x -> x % 2 == 0);
+       false
+      > SELECT _FUNC_(array(2, null, 8), x -> x % 2 == 0);
+       null
+  """,
+  since = "3.0.0")
+case class ArrayForAll(
+    argument: Expression,
+    function: Expression)
+  extends ArrayBasedSimpleHigherOrderFunction with CodegenFallback {
+
+  override def nullable: Boolean =
+      super.nullable || function.nullable
+
+  override def dataType: DataType = BooleanType
+
+  override def functionType: AbstractDataType = BooleanType
 
-  override def prettyName: String = "exists"
+  override def bind(f: (Expression, Seq[(DataType, Boolean)]) => LambdaFunction): ArrayForAll = {
+    val ArrayType(elementType, containsNull) = argument.dataType
+    copy(function = f(function, (elementType, containsNull) :: Nil))
+  }
+
+  @transient lazy val LambdaFunction(_, Seq(elementVar: NamedLambdaVariable), _) = function
+
+  /*
+   * true for all non null elements foundNull      result
+   *    F                              F             F
+   *    F                              T             F
+   *    T                              F             T
+   *    T                              T             N
+   */
+  override def nullSafeEval(inputRow: InternalRow, argumentValue: Any): Any = {
+    val arr = argumentValue.asInstanceOf[ArrayData]
+    val f = functionForEval
+    var forall = true
+    var foundNull = false
+    var i = 0
+    while (i < arr.numElements && forall) {
+      elementVar.value.set(arr.get(i, elementVar.dataType))
+      val ret = f.eval(inputRow)
+      if (ret == null) {
+        foundNull = true
+      } else if (!ret.asInstanceOf[Boolean]) {
+        forall = false
+      }
+      i += 1
+    }
+    if (foundNull && forall) {
+      null
+    } else {
+      forall
+    }
+  }
 }
 
 /**

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

@@ -304,6 +304,56 @@ class HigherOrderFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper
       Seq(true, null, true))
   }
 
+  test("ArrayForAll") {
+    def forall(expr: Expression, f: Expression => Expression): Expression = {
+      val ArrayType(et, cn) = expr.dataType
+      ArrayForAll(expr, createLambda(et, cn, f)).bind(validateBinding)
+    }
+
+    val ai0 = Literal.create(Seq(2, 4, 8), ArrayType(IntegerType, containsNull = false))
+    val ai1 = Literal.create(Seq[Integer](1, null, 3), ArrayType(IntegerType, containsNull = true))
+    val ai2 = Literal.create(Seq[Integer](2, null, 8), ArrayType(IntegerType, containsNull = true))
+    val ain = Literal.create(null, ArrayType(IntegerType, containsNull = false))
+
+    val isEven: Expression => Expression = x => x % 2 === 0
+    val isNullOrOdd: Expression => Expression = x => x.isNull || x % 2 === 1
+    val alwaysFalse: Expression => Expression = _ => Literal.FalseLiteral
+    val alwaysNull: Expression => Expression = _ => Literal(null, BooleanType)
+
+    checkEvaluation(forall(ai0, isEven), true)
+    checkEvaluation(forall(ai0, isNullOrOdd), false)
+    checkEvaluation(forall(ai0, alwaysFalse), false)
+    checkEvaluation(forall(ai0, alwaysNull), null)
+    checkEvaluation(forall(ai1, isEven), false)
+    checkEvaluation(forall(ai1, isNullOrOdd), true)
+    checkEvaluation(forall(ai1, alwaysFalse), false)
+    checkEvaluation(forall(ai1, alwaysNull), null)
+    checkEvaluation(forall(ai2, isEven), null)
+    checkEvaluation(forall(ai2, isNullOrOdd), false)
+    checkEvaluation(forall(ai2, alwaysFalse), false)
+    checkEvaluation(forall(ai2, alwaysNull), null)
+    checkEvaluation(forall(ain, isEven), null)
+    checkEvaluation(forall(ain, isNullOrOdd), null)
+    checkEvaluation(forall(ain, alwaysFalse), null)
+    checkEvaluation(forall(ain, alwaysNull), null)
+
+    val as0 =
+      Literal.create(Seq("a0", "a1", "a2", "a3"), ArrayType(StringType, containsNull = false))
+    val as1 = Literal.create(Seq(null, "b", "c"), ArrayType(StringType, containsNull = true))
+    val asn = Literal.create(null, ArrayType(StringType, containsNull = false))
+
+    val startsWithA: Expression => Expression = x => x.startsWith("a")
+
+    checkEvaluation(forall(as0, startsWithA), true)
+    checkEvaluation(forall(as1, startsWithA), false)
+    checkEvaluation(forall(asn, startsWithA), null)
+
+    val aai = Literal.create(Seq(Seq(1, 3, null), null, Seq(4, 5)),
+      ArrayType(ArrayType(IntegerType, containsNull = true), containsNull = true))
+    checkEvaluation(transform(aai, ix => forall(ix, isNullOrOdd)),
+      Seq(true, null, false))
+  }
+
   test("ArrayAggregate") {
     val ai0 = Literal.create(Seq(1, 2, 3), ArrayType(IntegerType, containsNull = false))
     val ai1 = Literal.create(Seq[Integer](1, null, 3), ArrayType(IntegerType, containsNull = true))

sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala

@@ -2322,6 +2322,116 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
     assert(ex4.getMessage.contains("cannot resolve '`a`'"))
   }
 
+  test("forall function - array for primitive type not containing null") {
+    val df = Seq(
+      Seq(1, 9, 8, 7),
+      Seq(2, 4, 6),
+      Seq.empty,
+      null
+    ).toDF("i")
+
+    def testArrayOfPrimitiveTypeNotContainsNull(): Unit = {
+      checkAnswer(df.selectExpr("forall(i, x -> x % 2 == 0)"),
+        Seq(
+          Row(false),
+          Row(true),
+          Row(true),
+          Row(null)))
+    }
+
+    // Test with local relation, the Project will be evaluated without codegen
+    testArrayOfPrimitiveTypeNotContainsNull()
+    // Test with cached relation, the Project will be evaluated with codegen
+    df.cache()
+    testArrayOfPrimitiveTypeNotContainsNull()
+  }
+
+  test("forall function - array for primitive type containing null") {
+    val df = Seq[Seq[Integer]](
+      Seq(1, 9, 8, null, 7),
+      Seq(2, null, null, 4, 6, null),
+      Seq(2, 4, 6, 8),
+      Seq.empty,
+      null
+    ).toDF("i")
+
+    def testArrayOfPrimitiveTypeContainsNull(): Unit = {
+      checkAnswer(df.selectExpr("forall(i, x -> x % 2 == 0 or x is null)"),
+        Seq(
+          Row(false),
+          Row(true),
+          Row(true),
+          Row(true),
+          Row(null)))
+      checkAnswer(df.selectExpr("forall(i, x -> x % 2 == 0)"),
+        Seq(
+          Row(false),
+          Row(null),
+          Row(true),
+          Row(true),
+          Row(null)))
+    }
+
+    // Test with local relation, the Project will be evaluated without codegen
+    testArrayOfPrimitiveTypeContainsNull()
+    // Test with cached relation, the Project will be evaluated with codegen
+    df.cache()
+    testArrayOfPrimitiveTypeContainsNull()
+  }
+
+  test("forall function - array for non-primitive type") {
+    val df = Seq(
+      Seq("c", "a", "b"),
+      Seq[String](null, null, null, null),
+      Seq.empty,
+      null
+    ).toDF("s")
+
+    def testNonPrimitiveType(): Unit = {
+      checkAnswer(df.selectExpr("forall(s, x -> x is null)"),
+        Seq(
+          Row(false),
+          Row(true),
+          Row(true),
+          Row(null)))
+    }
+
+    // Test with local relation, the Project will be evaluated without codegen
+    testNonPrimitiveType()
+    // Test with cached relation, the Project will be evaluated with codegen
+    df.cache()
+    testNonPrimitiveType()
+  }
+
+  test("forall function - invalid") {
+    val df = Seq(
+      (Seq("c", "a", "b"), 1),
+      (Seq("b", null, "c", null), 2),
+      (Seq.empty, 3),
+      (null, 4)
+    ).toDF("s", "i")
+
+    val ex1 = intercept[AnalysisException] {
+      df.selectExpr("forall(s, (x, y) -> x + y)")
+    }
+    assert(ex1.getMessage.contains("The number of lambda function arguments '2' does not match"))
+
+    val ex2 = intercept[AnalysisException] {
+      df.selectExpr("forall(i, x -> x)")
+    }
+    assert(ex2.getMessage.contains("data type mismatch: argument 1 requires array type"))
+
+    val ex3 = intercept[AnalysisException] {
+      df.selectExpr("forall(s, x -> x)")
+    }
+    assert(ex3.getMessage.contains("data type mismatch: argument 2 requires boolean type"))
+
+    val ex4 = intercept[AnalysisException] {
+      df.selectExpr("forall(a, x -> x)")
+    }
+    assert(ex4.getMessage.contains("cannot resolve '`a`'"))
+  }
+
   test("aggregate function - array for primitive type not containing null") {
     val df = Seq(
       Seq(1, 9, 8, 7),