diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala index 0041ccf775573..456c2967b91d4 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala +++ b/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"), diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala index b326e1c4c6af4..87b35b10d29fa 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala +++ b/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 + } + } } /** diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HigherOrderFunctionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HigherOrderFunctionsSuite.scala index 1411be8007deb..b83d03025d21c 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HigherOrderFunctionsSuite.scala +++ b/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)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala index 3f16f64f4b900..69d442a86c29e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala +++ b/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),