You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Add a stable counting-sort / group-ordering primitive for dense factorize codes #215
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 countordering, 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 returnuniques, 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.
Activity