Skip to content

fix: restrict array_filter array_compact fast path to the lambda variable #4848

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andygrove wants to merge 1 commit into
base: main
Choose a base branch
from
+26 −5
Open

fix: restrict array_filter array_compact fast path to the lambda variable #4848

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions spark/src/main/scala/org/apache/comet/serde/arrays.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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._
Expand Down Expand Up @@ -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
Expand Down
19 changes: 19 additions & 0 deletions spark/src/test/resources/sql-tests/expressions/array/array_filter.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>, 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
Loading