Polars Groups Columns So One Pass Feeds Many Aggregations

TL;DR: A group-by in Polars is two stages, not one. The first stage turns key columns into group ids: the keys are encoded (four different ways, depending on their dtype and count), hashed, and looked up in a hash table that hands out a group index per distinct key. The second stage is every aggregation expression, each of which consumes the same group structure. One grouping pass, many aggregations.
The streaming engine pushes this further: a fixed-size hot table keeps the frequent keys in memory, the rare ones are spilled, and the hot table is allowed to evict.

A naive group-by implementation computes the groups once per aggregation: df.group_by("k").agg(sum("a")) and df.group_by("k").agg(mean("b")) would each build their own hash table over the keys. That is a lot of duplicated work, and it scales badly with the number of aggregations.

Polars separates the two concerns. The grouping structure is computed once, and the aggregations are expressions that read it.

Two grouping engines, not one

Before going further: the type you may know as GroupsProxy is not in this tree. It has been split in two, and the split follows the execution engine:

the eager engine (polars-core) groups a materialized DataFrame into group tuples — GroupsType::Idx (a list of row indices per group) or GroupsType::Slice (a list of (offset, len) ranges, for sorted input);
the lazy/streaming engine (polars-expr) groups hash keys into group indices through the Grouper trait, because in a streaming query the keys arrive in morsels and the group table has to grow incrementally.

The eager one is the older design and is worth a look, because it is where the “sorted input” fast path lives:

rust
// crates/polars-core/src/frame/group_by/into_groups.rs:20
/// Used to create the tuples for a group_by operation.
pub trait IntoGroupsType {
    /// Create the tuples need for a group_by operation.
    ///     * The first value in the tuple is the first index of the group.
    ///     * The second value in the tuple is the indexes of the groups including the first value.
    fn group_tuples(&self, _multithreaded: bool, _sorted: bool) -> PolarsResult<GroupsType> {
        unimplemented!()
    }
}

Note the sorted flag: it is a promise from the caller that the key column is already sorted, and it unlocks a path that never builds a hash table at all:

rust
// crates/polars-core/src/frame/group_by/into_groups.rs:65
    fn create_groups_from_sorted(&self, multithreaded: bool) -> GroupsSlice {
        if verbose() {
            eprintln!("group_by keys are sorted; running sorted key fast path");
        }

Sorted keys mean the groups are contiguous ranges of rows, so the whole grouping reduces to a scan for run boundaries, which can be done in parallel over partitions of the sorted array.

For everything else, the eager path hashes:

rust
// crates/polars-core/src/frame/group_by/into_groups.rs:30
fn num_groups_proxy<T>(ca: &ChunkedArray<T>, multithreaded: bool, sorted: bool) -> GroupsType
where
    T: PolarsNumericType,
    T::Native: TotalHash + TotalEq + DirtyHash + ToTotalOrd,

TotalHash/TotalEq are the important detail: they give floats a total order, so NaN == NaN and -0.0 == 0.0 behave sanely inside a hash table.

The `Grouper` trait

The lazy/streaming side has one trait, and it is small:

rust
// crates/polars-expr/src/groups/mod.rs:45
/// A Grouper maps keys to groups, such that duplicate keys map to the same group.
pub trait Grouper: Any + Send + Sync {
    /// Creates a new empty Grouper similar to this one.
    fn new_empty(&self) -> Box<dyn Grouper>;

    /// Reserves space for the given number additional groups.
    fn reserve(&mut self, additional: usize);

    /// Returns the number of groups in this Grouper.
    fn num_groups(&self) -> IdxSize;

    /// Inserts the given subset of keys into this Grouper. If groups_idxs is
    /// passed it is extended such with the group index of keys[subset[i]].
    /// New groups get consecutive indices in the order they first occur.
    ///
    /// # Safety
    /// The subset indexes must be in-bounds.
    unsafe fn insert_keys_subset(
        &mut self,
        keys: &HashKeys,
        subset: &[IdxSize],
        group_idxs: Option<&mut Vec<IdxSize>>,
    );

    /// Returns the keys in this Grouper in group order, that is the key for
    /// group i is returned in row i.
    fn get_keys_in_group_order(&self, schema: &Schema) -> DataFrame;

That is the whole contract: insert keys, get group indices out, hand back the distinct keys in group order. The trait does not mention aggregations at all, which is exactly the point — the grouping structure is reusable by any number of consumers.

The other three methods on the trait (probe_partitioned_groupers, contains_key_partitioned_groupers, mark_groups_partitioned_groupers) exist for partitioned hash joins: they probe other groupers by hash partition, so a join can look up a key in the partition that owns it without searching all of them.

Four encodings, one interface

The interesting decision is how to represent a key so that it can be hashed and compared. Polars picks between four encodings based on the key schema:

rust
// crates/polars-expr/src/groups/mod.rs:113
pub fn new_hash_grouper(key_schema: Arc<Schema>) -> Box<dyn Grouper> {
    if let Some(layout) = KeyRowLayout::new(key_schema.iter_values()) {
        Box::new(key_rows::KeyRowHashGrouper::new(Arc::new(layout)))
    } else if key_schema.len() > 1 {
        Box::new(row_encoded::RowEncodedHashGrouper::new())
    } else {
        let (_name, dt) = key_schema.get_at_index(0).unwrap();
        match dt {
            dt if dt.is_primitive_numeric() | dt.is_temporal() => {
                with_match_physical_numeric_polars_type!(dt.to_physical(), |$T| {
                    Box::new(single_key::SingleKeyHashGrouper::<$T>::new())
                })
            },

            #[cfg(feature = "dtype-decimal")]
            DataType::Decimal(_, _) => {
                Box::new(single_key::SingleKeyHashGrouper::<Int128Type>::new())
            },
            #[cfg(feature = "dtype-categorical")]
            dt @ (DataType::Enum(_, _) | DataType::Categorical(_, _)) => {
                with_match_categorical_physical_type!(dt.cat_physical().unwrap(), |$C| {
                    Box::new(single_key::SingleKeyHashGrouper::<<$C as PolarsCategoricalType>::PolarsPhysical>::new())
                })
            },

            DataType::String | DataType::Binary => Box::new(binview::BinviewHashGrouper::new()),

            _ => Box::new(row_encoded::RowEncodedHashGrouper::new()),
        }
    }
}

Reading that top to bottom: if the key columns can be laid out as fixed-width key rows, use that; otherwise if there is more than one key column, use generic row encoding; otherwise, for a single key, use the most specialized grouper the dtype allows — a direct hash table over the physical values for numbers and dates, a view-aware one for strings, and generic row encoding for anything exotic.

Single key: the values themselves

For one numeric or temporal key, no encoding is needed at all — the physical value is the key:

rust
// crates/polars-expr/src/groups/single_key.rs:13
pub struct SingleKeyHashGrouper<T: PolarsDataType> {
    /// Each key with its group.
    table: HashTable<(T::Physical<'static>, IdxSize)>,
    num_groups: IdxSize,
    /// The group of the null key, IdxSize::MAX if there is none.
    null_idx: IdxSize,
    random_state: PlRandomState,
}

and the insert path is a textbook open-addressing hash table, with two hardware-friendly touches. First, the null key gets its own group index rather than being hashed, so nulls are never confused with a real value. Second, the work is done in blocks: the hashes of a whole block are computed (and prefetched) before the lookups happen, so the CPU has several independent memory accesses in flight:

rust
// crates/polars-expr/src/groups/single_key.rs:64
    unsafe fn insert_impl(
        &mut self,
        arr: &T::Array,
        null_is_valid: bool,
        subset: &[IdxSize],
        mut on_group: impl FnMut(IdxSize),
    ) {

rust
// crates/polars-expr/src/groups/single_key.rs:82
        let mut hashes = [0; BLOCK_SIZE];
        for block in subset.chunks(BLOCK_SIZE) {
            // Reserve up front so the table doesn't move after prefetching.
            self.table.reserve(block.len(), hasher);
            let prefetch = self.should_prefetch();
            for (hash, idx) in hashes.iter_mut().zip(block) {
                *hash = random_state.tot_hash_one(arr.value_unchecked(*idx as usize));
                if prefetch {
                    self.table.prefetch_insert(*hash);
                }
            }

That “reserve up front so the table doesn’t move after prefetching” comment is the kind of detail that only shows up after a profiler: reallocating the table after issuing prefetches would invalidate the addresses you just prefetched.

Key rows: a hand-rolled row format

The newest and most interesting encoding is key_rows. Instead of using the generic polars-row encoder, Polars lays out the key columns itself into a row format with a fixed stride, and the layout is computed once per schema:

rust
// crates/polars-expr/src/key_rows/layout.rs:72
/// The row format of keys of one schema. It depends only on the key dtypes: the null
/// bits are always present, also when nulls are not keys, so keys that disagree on
/// whether nulls match (the two sides of an anti join) share one layout.
///
/// Rows are written into zeroed words and null values are left zero, so equal keys
/// have equal words, including padding. The one exception is the second word of a
/// long view, which holds where its bytes are stored: equality skips that word and
/// compares the bytes instead.
#[derive(Debug, PartialEq)]
pub(crate) struct KeyRowLayout {
    pub(super) cols: Vec<ColLayout>,
    /// Byte offset of the null bits, one per column.
    pub(super) null_offset: usize,
    null_width: usize,
    pub(super) stride_words: usize,
    /// The first word of each view column, in column order.
    pub(super) view_words: Vec<usize>,
    /// Every word except the second half of each view.
    pub(super) plain_words: Vec<usize>,
}

Three design decisions are worth spelling out:

Nulls live in the row. One null bit per key column, always present even when nulls are not part of the key. That means the two sides of an anti join (where one side treats null as a key and the other does not) can share one layout and one hash table.
Columns are sorted by width, descending. The layout puts the widest columns first so that the whole row is a power-of-two number of 64-bit words, and so that comparisons can bail out early on the most discriminating word:

rust
// crates/polars-expr/src/key_rows/layout.rs:90
        let null_width = cols.len().div_ceil(8).next_power_of_two();
        let mut fields: Vec<(usize, Option<usize>)> = widths
            .into_iter()
            .enumerate()
            .map(|(i, width)| (width, Some(i)))
            .chain([(null_width, None)])
            .collect();
        fields.sort_by_key(|(width, _)| Reverse(*width));

Equality is a word-wise XOR fold. Because rows are written into zeroed words, “equal” means “no differing bits”, and that is one instruction per word for the fixed-width part:

rust
// crates/polars-expr/src/key_rows/layout.rs:170
    pub(super) unsafe fn rows_eq(
        &self,
        a: &[u64],
        b: &[u64],
        mut long_eq: impl FnMut(View, View) -> bool,
    ) -> bool {
        if self.view_words.is_empty() {
            return a.iter().zip(b).fold(0, |acc, (x, y)| acc | (x ^ y)) == 0;
        }
        for w in &self.plain_words {
            if a.get_unchecked(*w) != b.get_unchecked(*w) {
                return false;
            }
        }

Strings (Arrow “views”) are the awkward case: a view is a 16-byte struct where short strings are inlined and long strings point into a buffer. The second half of a view word holds a location, not data, so it is skipped by the word comparison and the bytes are compared instead.

Row encoding and binview: the fallbacks

Generic row encoding uses the polars-row crate and stores the encoded bytes in a byte-indexed map:

rust
// crates/polars-expr/src/groups/row_encoded.rs:12
#[derive(Default)]
pub struct RowEncodedHashGrouper {
    idx_map: BytesIndexMap<()>,
}

rust
// crates/polars-expr/src/groups/row_encoded.rs:23
    fn insert_key(&mut self, hash: u64, key: &[u8]) -> IdxSize {
        match self.idx_map.entry(hash, key) {
            Entry::Occupied(o) => o.index(),
            Entry::Vacant(v) => {
                let index = v.index();
                v.insert(());
                index
            },
        }
    }

For a single string key, there is a dedicated grouper that compares the Arrow views directly, so a 100 MB string column is never copied into a row buffer:

rust
// crates/polars-expr/src/groups/binview.rs:33
    /// # Safety
    /// The view must be valid for the given buffer set.
    #[inline(always)]
    unsafe fn insert_key(
        &mut self,
        hash: u64,
        view: View,
        buffers: &Buffer<Buffer<u8>>,
    ) -> IdxSize {
        unsafe {
            match self.idx_map.entry_view(hash, view, buffers) {

The keys themselves

The Grouper does not take a DataFrame of keys; it takes a HashKeys, which is the encoded key material plus a precomputed hash per row:

rust
// crates/polars-expr/src/hash_keys.rs:116
/// Represents a DataFrame plus a hash per row, intended for keys in grouping
/// or joining. The hashes may or may not actually be physically pre-computed,
/// this depends per type.
#[derive(Clone, Debug)]
pub enum HashKeys {
    RowEncoded(RowEncodedKeys),
    KeyRows(KeyRowKeys),
    Binview(BinviewKeys),
    Single(SingleKeys),
}

The from_df constructor is where the encoding decision is actually made, and it prefers key rows over everything else:

rust
// crates/polars-expr/src/hash_keys.rs:130
        if !force_row_encoding
            && let Some(layout) = KeyRowLayout::new(df.columns().iter().map(|c| c.dtype()))
        {
            return Self::KeyRows(KeyRowKeys::from_columns(
                df.columns(),
                Arc::new(layout),
                &random_state,
                null_is_valid,
            ));
        }

        let first_col_variant = hash_keys_variant_for_dtype(df[0].dtype());
        let use_row_encoding = force_row_encoding
            || df.width() > 1
            || first_col_variant == HashKeysVariant::RowEncoded;

Precomputing the hash per row is what makes partitioned joins and the partitioned groupers possible: a key’s hash determines its partition, so a probe only has to look in one partition’s table.

One pass, many aggregations

Now the payoff. The eager engine groups once and then runs every aggregation against the same group structure:

rust
// crates/polars-mem-engine/src/executors/group_by.rs:68
pub(super) fn group_by_helper(
    mut df: DataFrame,
    keys: Vec<Column>,
    aggs: &[Arc<dyn PhysicalExpr>],
    apply: Option<PlanCallback<DataFrame, DataFrame>>,
    state: &ExecutionState,
    maintain_order: bool,
    output_schema: &SchemaRef,
    slice: Option<(i64, usize)>,
) -> PolarsResult<DataFrame> {
    df.rechunk_mut_par();
    let gb = df.group_by_with_series(keys, true, maintain_order)?;

rust
// crates/polars-mem-engine/src/executors/group_by.rs:15
pub(super) fn evaluate_aggs(
    df: &DataFrame,
    aggs: &[Arc<dyn PhysicalExpr>],
    groups: &GroupPositions,
    state: &ExecutionState,
) -> PolarsResult<Vec<Column>> {
    RAYON.install(|| {
        aggs.par_iter()
            .map(|expr| {
                let agg = expr.evaluate_on_groups(df, groups, state)?.finalize();
                polars_ensure!(agg.len() == groups.len(), agg_len = agg.len(), groups.len());
                Ok(agg)
            })
            .collect::<PolarsResult<Vec<_>>>()
    })
}

groups is computed once and then shared, read-only, by every aggregation expression, which are themselves free to run in parallel. The lazy engine does the same thing, one level down: an aggregation expression receives the shared GroupPositions and reduces its own input column against it:

rust
// crates/polars-expr/src/expressions/aggregation.rs:232
    fn evaluate_on_groups_impl<'a>(
        &self,
        df: &DataFrame,
        groups: &'a GroupPositions,
        state: &ExecutionState,
    ) -> PolarsResult<AggregationContext<'a>> {
        let mut ac = self.input.evaluate_on_groups(df, groups, state)?;

The GroupedReduction trait is the streaming-engine equivalent, and it is designed so that many reductions can share one grouping structure:

rust
// crates/polars-expr/src/reduce/mod.rs:163
    /// Updates the specified group with the given values.
    ///
    /// For order-sensitive grouped reductions, seq_id can be used to resolve
    /// order between calls/multiple reductions.
    fn update_group(
        &mut self,
        values: &[&Column],
        group_idx: IdxSize,
        seq_id: u64,
    ) -> PolarsResult<()>;

and the partition table of the streaming group-by shows the pattern in full: one grouper, N reductions, one pass over the rows:

rust
// crates/polars-stream/src/nodes/group_by/table.rs:143
    pub fn add_rows(
        &mut self,
        input: usize,
        frame: &FrameWithKeys,
        rows: Option<&[IdxSize]>,
    ) -> PolarsResult<()> {
        let rows = rows_or_all(rows, &mut self.all_rows, frame.keys.len());
        self.group_idxs.clear();
        unsafe {
            self.grouper
                .insert_keys_subset(&frame.keys, rows, Some(&mut self.group_idxs));
        }
        let num_groups = self.grouper.num_groups();
        let params = &*self.params;
        let mut in_cols = Vec::new();
        for r_idx in &params.reductions_per_input[input] {
            let r = &mut self.reductions[*r_idx];
            r.resize(num_groups);
            in_cols.clear();
            in_cols.extend(
                params.grouped_reduction_cols[*r_idx]
                    .iter()
                    .map(|col| frame.frame.column(col).unwrap()),
            );
            // Cold rows only exist without order-sensitive reductions.
            unsafe { r.update_groups_subset(&in_cols, rows, &self.group_idxs, 0)? };
        }
        Ok(())
    }

The group indices computed by the grouper are handed to every reduction. r.resize(num_groups) grows the reduction’s state to the number of groups found so far; update_groups_subset then adds the values of the given rows to the right groups. Adding a tenth aggregation to a query costs a tenth reduction state and a tenth pass over the group’s rows — but not a tenth grouping pass.

The hot table

There is a second, more aggressive version of this idea in the streaming engine. Because a group-by over a large table has to hold the group table in memory, Polars splits the keys into hot (frequent) and cold (rare) ones:

rust
// crates/polars-stream/src/nodes/group_by/mod.rs:1
//! The streaming group-by.
//!
//! While receiving input, each pipeline aggregates rows with frequent keys in small hot tables
//! and stores the other (cold) rows with their hashed keys in spillable frames. Groups evicted
//! from a hot table, and the hot tables at the end of the input, become pre-aggregates. Once all
//! input is received, the cold rows and pre-aggregates are split into hash partitions, which are
//! aggregated independently. With at most one partition per pipeline, all partitions are built
//! before sending. Otherwise the cold rows are first stored per partition, and the partitions are
//! built a few at a time while the results are sent.

The hot table is a HotGrouper, a sibling of Grouper with one crucial difference — it has a fixed size and evicts rather than growing:

rust
// crates/polars-expr/src/hot_groups/mod.rs:21
/// A HotGrouper maps keys to groups, such that duplicate keys map to the same
/// group. Unlike a Grouper it has a fixed size and will cause evictions rather
/// than growing.
pub trait HotGrouper: Any + Send + Sync {
    /// Creates a new empty HotGrouper similar to this one, with the given size.
    fn new_empty(&self, groups: usize) -> Box<dyn HotGrouper>;

    /// Returns the number of groups in this HotGrouper.
    fn num_groups(&self) -> IdxSize;

    /// Returns the number of slots in the hot table.
    fn num_slots(&self) -> usize;

    /// Doubles the number of slots in the hot table. Group indices are unchanged.
    fn double(&mut self);

    /// Inserts the given keys into this Grouper, extending groups_idxs with
    /// the group index of keys[i].
    ///
    /// A missed key that is followed by a key with the same hash is inserted
    /// even if that evicts another key, as is every missed key with
    /// `force_hot`. This keeps sorted keys from first going cold.
    fn insert_keys(
        &mut self,
        keys: &HashKeys,
        hot_idxs: &mut Vec<IdxSize>,
        hot_group_idxs: &mut Vec<EvictIdx>,
        cold_idxs: &mut Vec<IdxSize>,
        force_hot: bool,
    );

The EvictIdx type is the trick that makes this work: a group index whose top bit is set means “this group was evicted, reset it before updating”. The reduction state carries the same flag, so an evicted group’s partial sum is discarded and rebuilt from the cold rows later. The note about “a missed key that is followed by a key with the same hash” is about sorted input: if the keys are sorted, the same key appears in a run, and evicting the first occurrence and re-inserting the second would defeat the purpose, so the grouper is allowed to evict on a hash collision instead.

The sink keeps both halves:

rust
// crates/polars-stream/src/nodes/group_by/sink.rs:60
/// The state of one pipeline of the sink.
struct LocalSink {
    // Created when the input first sends this pipeline a non-empty morsel.
    hot_grouper_per_input: Vec<Option<Box<dyn HotGrouper>>>,
    hot_grouped_reductions: Vec<Box<dyn GroupedReduction>>,

    // The number of slots of each hot grouper, and the rows missed by them.
    hot_table_size: usize,
    miss_f2: F2Sketch,

    // The keys of the cold rows and the pre-aggregates.
    key_sketch: CardinalitySketch,

    // Per input, the cold rows not yet in a frame.
    cold_row_builders: Vec<Option<FrameWithKeysBuilder>>,
    // Each frame with the input it came from.
    cold_frames: Vec<(usize, SpillFrameWithKeys)>,
    // The estimated bytes of `cold_frames`.
    cold_bytes: u64,
    // In creation order.
    pre_aggs: Vec<PreAgg>,
}

The hot table grows only when the evidence says it should — the F2 sketch counts how often keys are missed, and the table is enlarged only if the missed keys repeat enough and dominate the misses:

rust
// crates/polars-stream/src/nodes/group_by/sink.rs:47
/// The hot tables only grow if the frequently missed keys were missed at least this
/// many times each.
const HOT_TABLE_GROW_MIN_REPEAT: f64 = 4.0;

/// The hot tables only grow if the frequently missed keys make up at least this
/// fraction of the missed rows.
const HOT_TABLE_GROW_MIN_HEAVY_SHARE: f64 = 0.5;

The bottom line

The design in one sentence: grouping is a data structure, not a step. The group structure is built once, by a Grouper that picks an encoding suited to the key schema, and then it is a read-only input to every aggregation.

The trade-offs that come with it:

Encoding choice is per-query, not per-user. A single numeric key gets a direct hash table over the values; a multi-column key gets fixed-width key rows; strings get view comparisons; anything else falls back to generic row encoding. The cost is four implementations of one trait, and the benefit is that the common cases avoid copying key data at all.
The group structure is shared, so it is immutable during aggregation. Every aggregation reads it; none of them may reorder it. That is why evaluate_aggs can run the aggregations in parallel over one groups.
The hot table trades exactness of memory prediction for speed. It is a fixed-size cache of frequent keys, and it evicts. The evicted groups are not lost, they are re-derived from the spilled cold rows, which means the grouping work is done twice for evicted keys — a deliberate exchange of CPU for a bounded memory footprint.