Improve the performance of `Aggregator`, grouping, aggregation

[alamb]

Is your feature request related to a problem or challenge? Please describe what you are trying to do.
The performance of DataFusion is a very important feature, and aggregation is one of the core algorithms.

For example it features prominently in the ClickBench queries #5404

This ticket covers making it even faster.

Background

Prior Art

After #4924 we have a single GroupByHash implementation that uses the Arrow row format for group keys (good!) -- thank you @ozankabak and @mustafasrepo

However, the aggregator (e.g. the thing storing partial sums, counts, etc) is non ideal. Depending on the type of aggregate, it is either:

1. Box<dyn RowAccumulator> per
2. Box<dyn Accumulator>

The GroupByHash operator manages the mapping from group key to state and passes a RowAccessor to the RowAccumulator if needed.

Groups are calculated once per batch, then a take kernel reshuffles (aka copies) everything and slices are passed to the RowAccumulator / Accumulator.

Proposal

Ditch the word-aligned row format for the state management and change the aggregator to:

trait Aggregator: MemoryConsumer {
    /// Update aggregator state for given groups.
    fn update_batch(&mut self, keys: Rows, batch: &RecordBatch)) -> Result<()>;
    ...
}

Hence the aggregator will be dyn-dispatched ONCE per record batch and will keep its own internal state. This moves the key->state map from the [row_]hash.rs to the aggregators. We will provide tooling (macros or generics) to simplify the implementation and to avoid boilerplate code as much as possible.

Note that this also removes the take kernel since we think that it doesn't provide any performance improvements over iterating over the hashable rows and perform the aggregation row-by-row. We may bring the take handling back (as a "perform take if at least one aggregator wants that) if we can proof (via benchmarks) that this is desirable for certain aggregators, but we leave this out for now.

This also moves the memory consumer handling into the aggregator since the aggregator knows best how to split states and spill data.

Implementation Plan

TBD

Addtional Context

Note this was broken out of #2723 and is based on the wonderful writeup from @crepererum and @tustvold on #2723 (comment)

• group by high cardinality column in datafusion 10 times slower than low cardinality column #1246
• Support StreamAggregation / streaming group by #5133
• Improve grouping performance by special casing small / fixed size keys #846
• Make aggregate accumulators storage column-based  #956

[yjshen]

I'm curious about the proposed new Aggregator API. Do you know if it needs a hash table in each aggregator? I'm wondering because I'm a bit concerned about memory usage, especially for high cardinality aggregations.

Suppose keys are duplicated n times during execution (where n is the number of aggregators in the query). In that case, this could potentially lead to a significant increase in memory consumption, which might not be acceptable. What do you think about this?

[crepererum]

Do you know if it needs a hash table in each aggregator? I'm wondering because I'm a bit concerned about memory usage, especially for high cardinality aggregations.

Yes, see #2723 (comment)

We're trading speed for slightly more memory usage here. We could probably avoid storing the actual key in the per-aggregator hash tables by either interning the key using another hash table (key -> ID) or by hashing the keys to a "wide enough" ID (like 128bit, also assuming that the hash function is actually proper).

[yjshen]

So there will also be n times hash probe, and rehash, besides n memory overhead. Won't this hinder performance? How is the prototype benchmarked?

[crepererum]

So there will also be n times hash probe, and rehash, besides n memory overhead. Won't this hinder performance?

These are good questions.

Hashing an u128 ID should be very cheap since we don't need all this secure stuff anymore. So it's probably close to upper_word(id) ^ lower_word(id). The random access pattern could be a problem, but I think this will be hard to avoid in general, at least if you wanna have a maintainable, safe solution (in contrast to the untyped aligned row format that is currently used).

How is the prototype benchmarked?

There is no prototype yet, but a PR must run the groupby benchmarks.

---

That said, I am open to other suggestions. The issue is that the current system needs a redesign:

• It should be easy to add types and aggregations
• It should ideally use a single aggregator type
• It should work w/ memory accounting
• Data accesses should be properly typed (in contrast to the aligned row format)
• The solution should consider variable-length types. These cannot be added as an afterthought.

[yjshen]

While I agree that we may need to revisit the current Box<dyn Accumulator> approach when creating a fallback or default strategy, I do not believe we should combine all hash aggregate methods into one universal approach. Given that we are constructing an essential operator for a database execution engine, it is important to consider specialized implementations that offer performance benefits for handling different use cases.

I view row hash as a fast lane for certain queries rather than a dead end. By efficiently executing queries through row hash, we can avoid slower or more general approaches. Furthermore, I am open to exploring other specialized aggregate implementations that offer performance benefits to handle more use cases quickly.

For example, Spark provides multiple hash aggregate implementations, including vectorized hash, hash, and object hash. Similarly, ClickHouse offers different implementation options for its hash table, based on key types. Utilizing these specialized implementations can optimize query performance without relying on a one-size-fits-all approach.

It should be easy to add types and aggregations

Regarding the proposed criteria for a fallback strategy, I believe that while it should be easy to add types and aggregations, not all internal implementations should be treated equally, or else we risk the entire engine running at the speed of UDAFs.

It should work w/ memory accounting

While it would be beneficial for the fallback strategy to work with memory accounting, it is not strictly necessary. Users can take risks when defining UDAFs, or we could even risk the less-used built-in functions.

Data accesses should be properly typed (in contrast to the aligned row format)

I'm not convinced that data accesses need to be typed for aggregate implementations, which is entirely internal to database engines.

The solution should consider variable-length types. These cannot be added as an afterthought.

I understand that you were referring to the aligned row format, but it is possible to handle variable length fields using 'pointer swizzling.' Additionally, I don't believe that median is a performance-critical aggregate function that needs to be considered a first-class citizen for all aggregate strategies.