diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 8eda097ce6..8241051cf2 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -22,7 +22,7 @@ package org.apache.comet.serde import scala.annotation.tailrec import scala.jdk.CollectionConverters._ -import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArrayRepeat, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, IsNotNull, Literal, Reverse, Sequence, Size, Slice, SortArray, ZipWith} +import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArrayRepeat, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, IsNotNull, LambdaFunction, Literal, NamedLambdaVariable, Reverse, Sequence, Size, Slice, SortArray, ZipWith} import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -730,10 +730,12 @@ object CometArrayFilter extends CometExpressionSerde[ArrayFilter] { expr: ArrayFilter, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - expr.function.children.headOption match { - case Some(_: IsNotNull) => - // Fast path: `array_compact` lowers to `filter(arr, x -> x is not null)`. Use the native - // array_compact serde to avoid the per-batch JNI cost of the codegen dispatcher. + expr.function match { + case LambdaFunction(IsNotNull(v: NamedLambdaVariable), Seq(lambdaVar), _) + if v.exprId == lambdaVar.exprId => + // Fast path: Catalyst desugars `array_compact` to `filter(arr, x -> x IS NOT NULL)`, so + // restore the native serde here (avoids per-batch JNI). Guard requires the IsNotNull + // operand to be the lambda variable itself, not a captured column (#4830). CometArrayCompact.convert(expr, inputs, binding) case _ => // General lambda: run Spark's own evaluation through the codegen dispatcher so the result diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_filter.sql b/spark/src/test/resources/sql-tests/expressions/array/array_filter.sql index c4511e28e4..f7375744f3 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_filter.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_filter.sql @@ -29,3 +29,22 @@ SELECT filter(arr, x -> x >= 0) FROM test_array_filter query SELECT filter(arr, (x, i) -> i > 0) FROM test_array_filter + +statement +CREATE TABLE test_array_filter_captured(arr array, c int) USING parquet + +statement +INSERT INTO test_array_filter_captured VALUES + (array(1, NULL, 2), 5), + (array(3, NULL, 4), NULL), + (array(NULL, NULL), 7), + (NULL, 1) + +-- Genuine array_compact fast path: IsNotNull on the lambda variable drops the null elements. +query +SELECT filter(arr, x -> x IS NOT NULL) FROM test_array_filter_captured + +-- Regression for #4830: IsNotNull on captured column `c` is not array_compact and must not drop +-- the null elements of `arr`. +query +SELECT filter(arr, x -> c IS NOT NULL) FROM test_array_filter_captured