Repository navigation
Description
What is the problem the feature request solves?
Spark's ArrayContains, ArrayPosition and ArrayRemove compare elements to the value with the same rule, ordering.equiv in eval and ctx.genEqual in generated code, in every Spark version Comet supports (3.4 through 4.2). Under it -0.0 equals 0.0, all NaNs are equal, and nested arrays and structs compare their floats the same way. The three only differ in what they return from the matches.
Comet finds those matches in four different ways:
array_containsuses datafusion-spark'sarray_contains, which is DataFusion'sarray_hasplus a null fix-up. Onmain, arrays with a float leaf go through the codegen dispatcher instead, and feat: run array_contains on float elements natively with Spark's equa… #6599 addsspark_array_containsfor them.array_positionusesspark_array_position. It has typed per-row loops,compare_floatsfor flat floats, and for nested elementsnormalize_nested_floatsfollowed by one Arrowmake_comparator. A constant value is expanded to a full column first.array_removeuses DataFusion'sarray_remove_all, except for float leaves, which go tospark_array_remove(fix: match Spark's float semantics in array_remove and sort_array #6518). That kernel has a one-pass bitmap for a constant value and usesspark_equalityfor nested elements.
So a fix to the equality, a fast path or an edge case lands in one function and not the others. A constant value gets a single vectorized pass in array_remove and in #6599's array_contains, but not in array_position. Nested elements are compared with spark_equality in two of them and with normalize-then-compare in the third. #6520 proposed one shared kernel, and #6599 leaves that for a follow-up, which this issue tracks.
Describe the potential solution
One kernel that takes the list and the value, either a scalar or one per row, and finds for each element of the flattened values whether it equals its row's value under Spark's rule. Each function derives its result from that:
array_contains: true on any match, otherwise null if the row holds a null element, otherwise false.set_bits_beforealready counts both in one pass.array_position: the 1-based index of the first match, or 0.array_remove: the elements that don't match, with offsets fromset_bits_before.
The kernel should keep the fast paths each function has today, such as the single pass for a constant flat-float value and the early exit per row for a value per row, and pick one approach for nested elements. In the #6599 review, spark_equality on struct elements cost about twice what the codegen dispatcher does, so nested elements need a benchmark against the dispatcher before they move. Arrow's row format, or normalizing floats once and then comparing, are candidates to measure.
A shared path is also the place to handle these once instead of three times:
- Evaluating the value only on rows whose array is not NULL, so ANSI errors match Spark (Array functions evaluate their second argument on rows where the array is NULL, raising errors that Spark skips #6613).
- Declining non-default string collations, which
ArraysBase.isTypeSupportedalready does forarray_positionandarray_remove. - Struct elements, which
isTypeSupporteddeclines forarray_positionandarray_remove(it still cites the closed array_remove with structs causes error at runtime #1307) whilearray_containsaccepts them.
Activity