Polars Aggregates Faster When the Keys Are Already Sorted

TL;DR: Every ChunkedArray in Polars carries an IsSorted flag (Ascending, Descending, Not) in its metadata. The eager group-by reads that flag before it builds anything: if the key column is sorted, the groups are contiguous row ranges, so they are produced by a linear run-length scan instead of a hash table. The scan can also be split across threads, because the partitions are cut at group boundaries.
The output of that path is not just a different set of groups: it is a different type of group structure, GroupsType::Slice, which downstream aggregations consume with slice arithmetic instead of gathers.

Group-by is usually a hashing problem. You have N rows and a key column, and you want to know which rows share a key. The standard solution is a hash table: hash the key, look it up, append the row index to the bucket. That is O(N) but with a random-access pattern and a data structure that has to be allocated, grown and probed.

But there is a much cheaper structure available when the keys happen to be sorted, and sorted keys are not exotic: they are what a sort produces, what a parquet file with a sorted column contains, and what a range-partitioned stream delivers. When the keys are sorted, all the rows with the same key are adjacent, so the groups are just run boundaries. Finding them is a single pass over the array, with no allocation beyond the output vector and no hashing at all.

Polars exploits this, and the interesting part is how it decides.

text

   sort("k")
      |
      v
   [ 1 1 1 2 2 7 7 9 ]      IsSorted::Ascending
      |
      v
   group_by("k")
      |
      +-- is_sorted_ascending_flag()?  --- no --->  hash table
      |                                                   |
      |                                                   v
      |                                         GroupsType::Idx
      |                                         (row indices per group)
      |
      +-- yes --->  run-length scan  --->  GroupsType::Slice
                                             (offset, len) per group
                                             [ (0,3) (3,2) (5,2) (7,1) ]

There is an obvious objection: sorting costs O(N log N), which is more than hashing costs. The flag only pays for itself when the sort was going to happen anyway, or when the data arrived sorted. In practice that is common: a sort before a group_by, a parquet file whose key column is physically ordered, a window function’s PARTITION BY output, a table that was already sorted by an upstream system. In those cases the sort is not a cost of the group-by, it is a fact about the input, and the group-by gets to use it.

The flag, not the check

Polars does not verify that the keys are sorted. It trusts a bit of metadata that is set when the column is sorted and cleared whenever a mutation could invalidate it:

rust
// crates/polars-core/src/series/series_trait.rs:15
pub enum IsSorted {
    Ascending,
    Descending,
    Not,
}

rust
// crates/polars-core/src/chunked_array/ops/sort/mod.rs:199
        let mut ca = ChunkedArray::from_vec(ca.name().clone(), vals);
        let s = if options.descending {
            IsSorted::Descending
        } else {
            IsSorted::Ascending
        };
        ca.set_sorted_flag(s);

The flag lives in the chunked array’s StatisticsFlags, next to things like “can fast explode list”, and it is propagated by operations that preserve order (slice, take with sorted indices, append of two compatible runs). Anything that writes a value in an arbitrary position resets it.

So the group-by’s job is not “is this sorted?” but “does the column claim to be sorted?”. That is a single bit test:

rust
// crates/polars-core/src/frame/group_by/into_groups.rs:137
    fn group_tuples(&self, multithreaded: bool, sorted: bool) -> PolarsResult<GroupsType> {
        // sorted path
        if self.is_sorted_ascending_flag() || self.is_sorted_descending_flag() {
            // don't have to pass `sorted` arg, GroupSlice is always sorted.
            let groups = self.rechunk().create_groups_from_sorted(multithreaded);
            return Ok(GroupsType::new_slice(groups, false, true));
        }

Two things to notice. The sorted: bool argument of group_tuples is not about sorted keys — it is the “maintain the order of groups by first index” flag that group_by_stable passes, and it selects how the hash path orders its output. Key sortedness is detected from the column metadata.

Second, the sorted path returns GroupsType::new_slice(groups, false, true) and returns early. Nothing downstream is told “we took the fast path”; the type of the group structure is the message.

The run-length scan

Here is the whole sorted path:

rust
// crates/polars-core/src/frame/group_by/into_groups.rs:96
    fn create_groups_from_sorted(&self, multithreaded: bool) -> GroupsSlice {
        if verbose() {
            eprintln!("group_by keys are sorted; running sorted key fast path");
        }
        let arr = self.downcast_iter().next().unwrap();
        if arr.is_empty() {
            return GroupsSlice::default();
        }
        let mut values = arr.values().as_slice();
        let null_count = arr.null_count();
        let length = values.len();

        // all nulls
        if null_count == length {
            return vec![[0, length as IdxSize]];
        }

        let mut nulls_first = false;
        if null_count > 0 {
            nulls_first = arr.get(0).is_none()
        }

        if nulls_first {
            values = &values[null_count..];
        } else {
            values = &values[..length - null_count];
        };

        let n_threads = RAYON.current_num_threads();
        if multithreaded && n_threads > 1 {
            let parts =
                create_clean_partitions(values, n_threads, self.is_sorted_descending_flag());
            let n_parts = parts.len();

            let first_ptr = &values[0] as *const T::Native as usize;
            let groups = parts.par_iter().enumerate().map(|(i, part)| {
                // we go via usize as *const is not send
                let first_ptr = first_ptr as *const T::Native;

                let part_first_ptr = &part[0] as *const T::Native;
                let mut offset = unsafe { part_first_ptr.offset_from(first_ptr) } as IdxSize;

                // nulls first: only add the nulls at the first partition
                if nulls_first && i == 0 {
                    partition_to_groups(part, null_count as IdxSize, true, offset)
                }
                // nulls last: only compute at the last partition
                else if !nulls_first && i == n_parts - 1 {
                    partition_to_groups(part, null_count as IdxSize, false, offset)
                }
                // other partitions
                else {
                    if nulls_first {
                        offset += null_count as IdxSize;
                    };

                    partition_to_groups(part, 0, false, offset)
                }
            });
            let groups = RAYON.install(|| groups.collect::<Vec<_>>());
            flatten_par(&groups)
        } else {
            partition_to_groups(values, null_count as IdxSize, nulls_first, 0)
        }
    }

The scan itself is a single loop over the values, comparing each element to the first element of the current run:

rust
// crates/polars-arrow/src/legacy/kernels/sort_partition.rs:91
pub fn partition_to_groups_amortized<T>(
    values: &[T],
    first_group_offset: IdxSize,
    nulls_first: bool,
    offset: IdxSize,
    out: &mut Vec<[IdxSize; 2]>,
) where
    T: Debug + TotalEq + Sized,
{
    if let Some(mut first) = values.first() {
        out.clear();
        if nulls_first && first_group_offset > 0 {
            out.push([0, first_group_offset])
        }

        let mut first_idx = if nulls_first { first_group_offset } else { 0 } + offset;

        for val in values {
            // new group reached
            if val.tot_ne(first) {
                let val_ptr = val as *const T;
                let first_ptr = first as *const T;

                // SAFETY:
                // all pointers suffice the invariants
                let len = unsafe { val_ptr.offset_from(first_ptr) } as IdxSize;
                out.push([first_idx, len]);
                first_idx += len;
                first = val;
            }
        }

The trick in that loop is that the length of a run is computed by pointer arithmetic — val_ptr.offset_from(first_ptr) — rather than by counting. The loop never increments a counter; it just watches the pointer move. The comparison is tot_ne, the total-order equality, so floats behave (NaN is never equal to itself, but it is equal to itself under total order, and -0.0 and 0.0 compare equal).

Nulls get special treatment because Polars sorts them to one end of the array: the null run is emitted as its own group at the front or the back, and the values slice is trimmed to exclude it.

Making the scan parallel

A naive “split the array in half and scan each half” would be wrong: a group could straddle the boundary. Polars cuts the array at group boundaries instead, and the search for those boundaries is itself cheap because the array is sorted:

rust
// crates/polars-arrow/src/legacy/kernels/sort_partition.rs:10
/// Find partition indexes such that every partition contains unique groups.
fn find_partition_points<T>(values: &[T], n: usize, descending: bool) -> Vec<usize>
where
    T: Debug + NativeType,
{
    let len = values.len();
    if n > len {
        return find_partition_points(values, len / 2, descending);
    }
    if n < 2 {
        return vec![];
    }
    let chunk_size = len / n;

    let mut partition_points = Vec::with_capacity(n + 1);

    let mut start_idx = 0;
    loop {
        let end_idx = start_idx + chunk_size;
        if end_idx >= len {
            break;
        }
        // first take that partition as a slice
        // and then find the location where the group of the latest value starts
        let part = &values[start_idx..end_idx];

        let latest_val = values[end_idx];
        let idx = if descending {
            part.partition_point(|v| v.tot_gt(&latest_val))
        } else {
            part.partition_point(|v| v.tot_lt(&latest_val))
        };

        if idx != 0 {
            partition_points.push(idx + start_idx)
        }

        start_idx += chunk_size;
    }
    partition_points
}

For each candidate boundary it looks at the value just past the boundary and binary-searches backwards within the chunk for the start of that value’s run. That gives “clean partitions”: every partition contains whole groups. create_clean_partitions then slices the array at those points, and each thread scans one slice independently. The offsets are fixed up by pointer arithmetic in create_groups_from_sorted so that the group offsets are global row indices, and the per-partition results are concatenated.

The null group is handled by only emitting it in one partition (the first if nulls come first, the last if nulls come last), with the other partitions compensating by an offset shift.

What the hash path looks like

For comparison, this is the same operation when the keys are not sorted:

rust
// crates/polars-core/src/frame/group_by/hashing.rs:75
pub(crate) fn group_by<K>(keys: impl Iterator<Item = K>, sorted: bool) -> GroupsType
where
    K: TotalHash + TotalEq,
{
    let init_size = get_init_size();
    let (mut first, mut groups);
    if sorted {
        groups = Vec::with_capacity(get_init_size());
        first = Vec::with_capacity(get_init_size());
        let mut hash_tbl = PlHashMap::with_capacity(init_size);
        for (idx, k) in keys.enumerate_idx() {
            match hash_tbl.entry(TotalOrdWrap(k)) {
                Entry::Vacant(entry) => {
                    let group_idx = groups.len() as IdxSize;
                    entry.insert(group_idx);
                    groups.push(unitvec![idx]);
                    first.push(idx);
                },
                Entry::Occupied(entry) => unsafe {
                    groups.get_unchecked_mut(*entry.get() as usize).push(idx)
                },
            }
        }
    } else {
        let mut hash_tbl = PlHashMap::with_capacity(init_size);
        for (idx, k) in keys.enumerate_idx() {
            match hash_tbl.entry(TotalOrdWrap(k)) {
                Entry::Vacant(entry) => {
                    entry.insert((idx, unitvec![idx]));
                },
                Entry::Occupied(mut entry) => entry.get_mut().1.push(idx),
            }
        }
        (first, groups) = hash_tbl.into_values().unzip();
    }
    GroupsType::Idx(GroupsIdx::new(first, groups, sorted))
}

The hash path allocates a hash table, a unitvec per group, and produces GroupsType::Idx — a list of row indices per group. The sorted path produces GroupsType::Slice — a list of (offset, len) pairs, where the rows of a group are a contiguous range.

The multi-threaded variant of the hash path does not partition the rows; it partitions the hash space, so every thread has to look at every key:

rust
// crates/polars-core/src/frame/group_by/hashing.rs:146
pub(crate) fn group_by_threaded_slice<T, IntoSlice>(
    keys: Vec<IntoSlice>,
    n_partitions: usize,
    sorted: bool,
) -> GroupsType
where
    T: ToTotalOrd,
    <T as ToTotalOrd>::TotalOrdItem: Send + Sync + Copy + DirtyHash,
    IntoSlice: AsRef<[T]> + Send + Sync,
{
    let init_size = get_init_size();

    // We will create a hashtable in every thread.
    // We use the hash to partition the keys to the matching hashtable.
    // Every thread traverses all keys/hashes and ignores the ones that doesn't fall in that partition.
    let out = RAYON.install(|| {
        (0..n_partitions)
            .into_par_iter()
            .map(|thread_no| {
                let mut hash_tbl = PlHashMap::with_capacity(init_size);

                let mut offset = 0;
                for keys in &keys {
                    let keys = keys.as_ref();
                    let len = keys.len() as IdxSize;

                    for (key_idx, k) in keys.iter().enumerate_idx() {
                        let k = k.to_total_ord();
                        let idx = key_idx + offset;

                        if thread_no == hash_to_partition(k.dirty_hash(), n_partitions) {
                            match hash_tbl.entry(k) {
                                Entry::Vacant(entry) => {
                                    entry.insert((idx, unitvec![idx]));
                                },
                                Entry::Occupied(mut entry) => {
                                    entry.get_mut().1.push(idx);
                                },
                            }
                        }
                    }
                    offset += len;
                }
                hash_tbl
                    .into_iter()
                    .map(|(_k, v)| v)
                    .collect_trusted::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    });
    finish_group_order(out, sorted)
}

Note the shape of the work: every thread reads every key, computes its hash, and throws away the keys that are not its own. The hash table is partitioned rather than the data, so there is no work partitioning to do — but each thread touches the whole array. The sorted path does the opposite: the data is partitioned, so a thread only reads its own slice, and the only coordination is choosing where to cut.

That difference propagates all the way into the aggregations. Every aggregation has a match arm per group type:

rust
// crates/polars-core/src/frame/group_by/aggregations/dispatch.rs:30
    pub unsafe fn agg_valid_count(&self, groups: &GroupsType) -> Series {
        // Prevent a rechunk for every individual group.
        let valid = self.rechunk_validity();

        match groups {
            GroupsType::Idx(groups) => agg_helper_idx_on_all::<IdxType, _>(groups, |idxs| {
                debug_assert!(idxs.len() <= self.len());
                if let Some(v) = &valid {
                    let mut count = 0;
                    for idx in idxs.iter() {
                        count += unsafe { v.get_bit_unchecked(*idx as usize) as IdxSize };
                    }
                    Some(count)
                } else {
                    Some(idxs.len() as IdxSize)
                }
            }),
            GroupsType::Slice { groups, .. } => {
                _agg_helper_slice::<IdxType, _>(groups, |[first, len]| {
                    debug_assert!(len <= self.len() as IdxSize);
                    if let Some(v) = &valid {
                        let m = BitMask::from_bitmap(v).sliced(first as usize, len as usize);
                        Some(m.set_bits() as IdxSize)
                    } else {
                        Some(len)
                    }
                })
            },
        }
    }

The Idx arm walks a list of indices and tests each one’s validity bit. The Slice arm takes a bitmap slice — a contiguous run of bits — and counts the set bits in it, which is a popcount over a machine word rather than a per-row branch. Same answer, very different instruction mix.

The slice representation also enables a second class of optimization: an aggregation over a contiguous range can use rolling kernels, since the ranges are themselves monotone. agg_quantile_generic has a branch that checks whether the group slices are monotonic and, if so, hands the whole thing to a rolling-window kernel instead of computing a quantile per group.

The metadata that travels with the groups

The sorted fast path also leaves a trace in the output. When the group-by produces its key columns, it copies the sorted flag of the input key onto the output:

rust
// crates/polars-core/src/frame/group_by/mod.rs:263
                        GroupsType::Slice {
                            groups,
                            overlapping,
                            monotonic: _,
                        } => {
                            if *overlapping && !groups.is_empty() {
                                // Groups can be sliced.
                                let offset = groups[0][0];
                                let [upper_offset, upper_len] = groups[groups.len() - 1];
                                return s.slice(
                                    offset as i64,
                                    ((upper_offset + upper_len) - offset) as usize,
                                );
                            }

                            let indices = groups
                                .iter()
                                .map(|&[first, _len]| first)
                                .collect_ca(PlSmallStr::EMPTY);
                            // SAFETY: groups are always in bounds.
                            let mut out = unsafe { s.take_unchecked(&indices) };
                            // Sliced groups are always in order of discovery.
                            out.set_sorted_flag(s.is_sorted_flag());
                            out
                        },

So df.sort("k").group_by("k").agg(...) produces a result whose k column is still flagged as sorted, and a subsequent group_by on that result takes the fast path again. The same reasoning is used for the Idx variant, but only when the groups are known to be ordered by first index.

The GroupsType::Slice variant carries two more pieces of metadata:

rust
// crates/polars-core/src/frame/group_by/position.rs:260
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GroupsType {
    Idx(GroupsIdx),
    Slice {
        // the groups slices
        groups: GroupsSlice,
        /// Indicates if the groups may overlap, i.e., at least one index MAY be
        /// included in more than one group slice.
        overlapping: bool,
        /// Indicates if the groups are rolling, i.e. for every consecutive group
        /// slice (offset, len), and given start = offset and end = offset + len,
        /// then both new_start >= start AND new_end >= end MUST be true.
        monotonic: bool,
    },
}

monotonic says the group ranges are themselves ordered and non-shrinking, which is exactly the precondition for a rolling kernel. overlapping says a row may appear in more than one group, which is false for a group-by but true for e.g. a window function. Both invariants are checked in debug builds:

rust
// crates/polars-core/src/frame/group_by/position.rs:308
    pub fn new_slice(groups: GroupsSlice, overlapping: bool, monotonic: bool) -> Self {
        #[cfg(debug_assertions)]
        {
            fn groups_overlap(groups: &GroupsSlice) -> bool {
                if groups.len() < 2 {
                    return false;
                }
                let mut groups = groups.clone();
                groups.sort();
                let mut prev_end = groups[0][1];

                for g in &groups[1..] {
                    let start = g[0];
                    let end = g[1];
                    if start < prev_end {
                        return true;
                    }
                    if end > prev_end {
                        prev_end = end;
                    }
                }
                false
            }

            assert!(overlapping || !groups_overlap(&groups));
            assert!(!monotonic || slice_groups_are_monotonic(&groups));
        }

The streaming version

The same idea exists in the streaming engine, where the input arrives in morsels and you cannot assume the whole column is in memory. There is a dedicated node for it, and it does the same run-length scan — but on a buffer that is flushed at group boundaries:

rust
// crates/polars-stream/src/nodes/sorted_group_by.rs:86
        let mut offset = 0;
        let groups = idxs
            .iter()
            .map(|i| {
                let start = offset;
                offset += i;
                [start, *i]
            })
            .collect();
        let groups = GroupsType::new_slice(groups, false, true).into_sliceable();

idxs comes from rle_lengths(column, idxs), which computes the length of each run of equal values — the same run-length encoding, but as a list of lengths instead of a list of ranges. The node then buffers morsels until the buffer holds more than one distinct key, and flushes at a point where a group cannot straddle the boundary:

rust
// crates/polars-stream/src/nodes/sorted_group_by.rs:247
                if fst == lst {
                    continue;
                }

                let mut buf_key_column = buf_key_column.as_materialized_series().clone();
                buf_key_column.set_sorted_flag(IsSorted::Ascending);

                let descending = fst > lst;
                let num_flushable = search_sorted(
                    &buf_key_column,
                    &buf_key_column.tail(Some(1)),
                    SearchSortedSide::Left,
                    descending,
                )
                .unwrap();

The bottom line

The design in one sentence: sortedness is metadata, and metadata is exploited at every stage that sees it.

The trade-offs:

The fast path trusts the flag. If the flag is wrong, the group-by produces wrong groups — a silent correctness bug. That is why the flag is cleared on every mutation that could break the order, and why the sorted path rechunks first (it needs one contiguous chunk) but never re-validates.
It only pays off for a single key column. The multi-key path in group_by_with_series encodes all keys into row bytes first, and then hashes those bytes; the row bytes have no sortedness flag, so the sorted fast path is not available there.
The output type changes. GroupsType::Slice instead of GroupsType::Idx is not just a different representation: it selects slice-based aggregation code, enables rolling kernels through the monotonic flag, and lets the sorted flag survive into the output so the next group-by can take the same path.
Nulls are a special case everywhere. Nulls must be at one end of the sorted array, and the scan has to know which end, and the parallel version has to emit the null group exactly once.