Skip to content

Add a stable counting-sort / group-ordering primitive for dense factorize codes #215

Description

@flexatone

Summary

Add a primitive that turns the dense integer codes from factorize into a stable
grouping permutation
— the row order that groups equal codes contiguously, in code order,
preserving original order within each group — computed in O(n) via counting sort rather
than a comparison sort.

This is the natural companion to factorize (#212). factorize already gives O(n) dense
codes; the remaining bottleneck for group-by is turning those codes into a group ordering,
which today must be done with np.argsort(codes, kind='stable') — an O(n log n) comparison
sort that dominates the runtime.

Motivation / benchmark

In static-frame, group-by and pivot now factorize the key and then need a stable permutation
that orders rows by their code (to walk contiguous groups). Decomposition on a 1M-row,
5k-group string key:

step time
factorize(key, sort=True) 12 ms
np.bincount(codes) 1.4 ms
np.argsort(codes, kind='stable') (the group ordering) ~100 ms ← dominates

The counting-sort inputs are already cheap (bincount is 1.4 ms), but numpy has no
vectorized stable counting sort
— the stable permutation can only be obtained via argsort,
which for int64 uses timsort (O(n log n)), not radix. So the O(n) win that dense codes should
enable is lost to the sort. A C-level counting sort would take the ~100 ms step to a few ms,
turning static-frame's current ~2.3x string group-by speedup into roughly ~10x.

Proposed API

Any of these shapes would work; the essential output is the stable permutation and the group
boundaries:

# Option 1: standalone, operating on dense codes + group count
ordering, offsets = group_ordering(codes: np.ndarray[int64], size: int)
#   ordering: intp permutation; rows grouped by code, code order, stable within group
#   offsets:  size+1 int array; group g occupies ordering[offsets[g]:offsets[g+1]]

# Option 2: fold into factorize as an opt-in extra return
uniques, codes, ordering = factorize(array, sort=True, return_ordering=True)

Requirements:

  • Stable: within a group, original row order is preserved (needed to match a stable sort).
  • O(n): counting sort over the dense [0, size) code range (bincount + prefix sum +
    a single stable scatter pass).
  • ordering dtype intp; ideally the group offsets/counts returned too, since callers
    typically want the contiguous group boundaries (equivalent to a transition_slices over the
    reordered codes) without recomputing them.

Context

Consumer: static-frame's factorize_argsort helper (group-by dispatch + pivot). It currently
does factorize(...) then np.argsort(codes, kind='stable'); this primitive would replace the
argsort. The existing transition_slices_from_group pairs naturally with the returned offsets.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions