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.
Summary
Add a primitive that turns the dense integer
codesfromfactorizeinto a stablegrouping 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).factorizealready gives O(n) densecodes; 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) comparisonsort 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:
factorize(key, sort=True)np.bincount(codes)np.argsort(codes, kind='stable')(the group ordering)The counting-sort inputs are already cheap (
bincountis 1.4 ms), but numpy has novectorized 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:
Requirements:
[0, size)code range (bincount+ prefix sum +a single stable scatter pass).
orderingdtypeintp; ideally the groupoffsets/countsreturned too, since callerstypically want the contiguous group boundaries (equivalent to a
transition_slicesover thereordered codes) without recomputing them.
Context
Consumer: static-frame's
factorize_argsorthelper (group-by dispatch + pivot). It currentlydoes
factorize(...)thennp.argsort(codes, kind='stable'); this primitive would replace theargsort. The existing
transition_slices_from_grouppairs naturally with the returned offsets.