Skip to content

Share one Spark-equality match kernel between array_contains, array_position and array_remove #6614

Open

Description

@andygrove

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_contains uses datafusion-spark's array_contains, which is DataFusion's array_has plus a null fix-up. On main, 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 adds spark_array_contains for them.
  • array_position uses spark_array_position. It has typed per-row loops, compare_floats for flat floats, and for nested elements normalize_nested_floats followed by one Arrow make_comparator. A constant value is expanded to a full column first.
  • array_remove uses DataFusion's array_remove_all, except for float leaves, which go to spark_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 uses spark_equality for 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_before already 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 from set_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:

Additional context

Follow-up to #6520, part of #6385.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions