andygrove opened a new pull request, #6716: URL: https://github.com/apache/datafusion-comet/pull/6716
## Which issue does this PR close? Closes #6613. ## Rationale for this change Spark's array lookups are null-intolerant `BinaryExpression`s or `TernaryExpression`s. For a NULL array, Spark returns NULL without evaluating the other arguments. A native scalar function evaluates every argument over the whole batch first, so an argument that raises, such as an ANSI cast of a malformed string, raised on rows Spark never evaluates it for, and the query failed. Besides the six shapes in the issue, a probe on Spark 4.1 found the same failure for `m[k]`, `slice` and `a[i]`. The same eager evaluation hands a nondeterministic argument rows that Spark never gives it. On main, `array_position(ai, CAST(rand(7L) * 3 AS INT))` returns different values from Spark in either ANSI mode, because `rand` also advances on the rows whose array is NULL. The issue points at `CometElementAt`'s `CASE WHEN <array> IS NOT NULL` guard and `CometArrayJoin`'s dispatcher route. The guard evaluates the array twice and filters every array element on each batch that has a NULL array, and it would need the nondeterministic decline that `CometElementAt` has. The dispatcher route would take common shapes such as `array_contains(ids, CAST(i AS BIGINT))` off the native kernels. `CometElementAt`'s doc comment already names the fix that avoids both: evaluate the operand once and mask the later arguments with it. ## What changes are included in this PR? - `NullShortCircuit` (`native/spark-expr/src/conditional_funcs/null_short_circuit.rs`) evaluates a function's arguments the way Spark does: the first over the whole batch, and each later one only for the rows where no earlier argument is NULL, over a batch projected down to the columns that argument reads. The function still runs over every row, with NULL for each skipped argument. Native `ArrayInsert` already evaluates its own arguments this way. - `ScalarFunc` and `ListExtract` get a `null_short_circuit` flag, which the native planner turns into the wrapper. The planner leaves the function unwrapped when every later argument passes `is_infallible` from `case_when.rs`, so a column, a literal or a widening cast plans exactly as before. - `QueryPlanSerde.withNullShortCircuit` sets the flag. The serdes for `array_contains` (and so `map_contains_key`), `array_position`, `array_remove`, `arrays_overlap`, `array_union`, `slice`, `GetMapValue` and `GetArrayItem` use it. - The new-expression guide says when to use it. `element_at` already has its `CASE WHEN` guard under ANSI, and #5867 changes that serde, so I left it alone. It could move to the flag later and drop its nondeterministic decline. One case moves the other way. When the same failing subexpression appears twice in one projection, as in `array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT))`, Spark's codegen subexpression elimination evaluates the shared cast for every row and raises. Comet now returns NULL there, as Spark itself does with `spark.sql.subexpressionElimination.enabled=false` or with interpreted evaluation; I checked both on 4.1. On main, Comet raised in that case too. ## How are these changes tested? - Two SQL fixtures, `array/array_null_short_circuit.sql` and `map/map_null_short_circuit.sql`, read a single file in which NULL and non-NULL arrays share a batch. They pin native evaluation with `expect_native`, cover a failure that is not ANSI-only (a `slice` start of 0) and the `rand` argument, and check that a malformed value on a row whose array is not NULL still raises. With the flag turned off in the serde, both fixtures fail, and the `rand` query fails with mismatched results rather than an error. - Rust unit tests cover the masking, a ternary function, a NULL literal array, `ListExtract`, a CASE that rewrites the wrapper's column indices, and the skip when every later argument is infallible. Evaluating the later arguments over the whole batch fails five of them, and filtering every column instead of the ones an argument reads fails another. - The `array/` and `map/` SQL fixtures, `CometArrayExpressionSuite` and `CometMapExpressionSuite` pass on Spark 3.4, 3.5, 4.0, 4.1 and 4.2. `CometFloatSemanticsSuite`, `CometDecimalPromotionSuite` and `CometGenerateExecSuite` pass on 4.1. I haven't run Spark's SQL suite locally. Cost, measured with a release build over 8M rows: with 20% NULL arrays and a cast argument in non-ANSI mode, where main does not fail, the expression takes 5 to 12% longer. With no NULL arrays the difference is within noise. The guard applies there because `Cast::is_infallible` counts every string-to-integral cast as one that can fail, including the LEGACY and TRY ones, which return NULL for a malformed string. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
