apache / apache/auron

Support native bit_and / bit_or / bit_xor aggregate

Open
#2,362 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
1.8k
Forks
241
Avg merge
2d 12h
Merged PRs (30d)
21

Description

## Background

Auron does not implement the bitwise aggregates `bit_and` / `bit_or` / `bit_xor` natively, so they fall back to the generic UDAF path (a JNI call back into the JVM) and lose vectorized acceleration. These aggregates are commonly used for flag/bitmask rollups (e.g. `bit_or` to union feature flags across rows, `bit_and` for common flags) and parity/checksum use cases.

## Proposal

Add native `bit_and` / `bit_or` / `bit_xor` aggregates:

1. Implement a generic `AggBitwise

` in `datafusion-ext-plans` (`agg/bitwise.rs`), parameterized by the bitwise operator, with type aliases `AggBitAnd` / `AggBitOr` / `AggBitXor`. The accumulator is a single column of the input type: the first non-null value initializes the slot and every subsequent value is folded in with the operator. The operators are associative and commutative, so the result is independent of the visiting/merge order, and null inputs are skipped (an all-null group yields null). Only integral inputs (Int8/Int16/Int32/Int64) are accepted.
2. Wire the three functions through the `AggFunction` enum, `create_agg`, the protobuf contract (`BIT_AND` / `BIT_OR` / `BIT_XOR`), the `protobuf::AggFunction -> AggFunction` conversion, and the window-aggregate mapping.
3. Add the `BitAndAgg` / `BitOrAgg` / `BitXorAgg` expression conversions in `NativeConverters.convertAggregateExpr`.
4. Declare the native aggregate buffer schema in `NativeAggBase.computeNativeAggBufferDataTypes` (`Seq(dataType)`, a single column) so the partial -> shuffle -> final buffer schema matches the native side.

## Scope / Non-goals

- Integral inputs only (byte / short / int / long), matching Spark's `BitAndAgg` / `BitOrAgg` / `BitXorAgg`.
- Window aggregate (`bit_*` over a window) reuses the same `AggFunction`.

## Tests

- Rust unit test `agg_exec::test::test_agg_bitwise`: partial -> final two-phase aggregation over a nullable integer column, verifying bit_and / bit_or / bit_xor (including null skipping).
- Scala end-to-end test in `AuronDataFrameAggregateSuite` ("native bit_and / bit_or / bit_xor aggregate", spark34 + spark35): a grouped aggregate exercising the full partial -> shuffle -> final native path (including an all-null group), asserting correct values and that the plan offloads to `NativeAggBase`.

Contributor guide

Open the contributing guide

Research direction

Start with datafusion-ext-plans/agg/bitwise.rs and the existing AggFunction/create_agg wiring, then inspect the protobuf and window-aggregate mappings. Read NativeConverters.convertAggregateExpr and NativeAggBase.computeNativeAggBufferDataTypes before running agg_exec::test::test_agg_bitwise and the AuronDataFrameAggregateSuite case. Done means nullable integral grouped and window aggregates pass the Rust and Scala tests, including all-null groups, and the plan uses NativeAggBase.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust, scala, spark
Domain
backend, data-engineering
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.