group-by: transient-vector buckets + fix nil keys in transient maps (#155)

The map build already used a transient map, but each bucket was rebuilt with
a persistent (conj (get ret k []) x) per element — an O(log n) trie path
rebuild + alloc each. A coarse grouping (few large buckets) was bound on that
conj, not the map build. Buckets are now native arrays (transient vectors,
O(1) push) frozen once; distinct keys are tracked in a side vector so the
buckets freeze in place with no second map rebuild. A bucket's first element
stays a cheap persistent [x] and only promotes to a transient on the second,
so an all-singletons grouping pays no transient alloc.

  coarse (10/100 buckets, 50k): ~313ms -> ~125ms (~2.5x)
  2 buckets (50k):              ~322ms -> ~129ms (~2.5x)
  all-unique (50k):             ~949ms -> ~892ms (no regression)

Surfaced a latent bug: canon-key returns nil for a nil key and Janet tables
drop a nil key, so the canon-keyed transient map silently lost a nil-key
entry — group-by/frequencies/assoc!/into{} dropped the whole nil bucket
((group-by identity [nil nil 1]) gave {1 [1]}, not {nil [nil nil], 1 [1]}).
Route nil through a sentinel (tbl-key) at the transient-map keying sites;
persistent!/count/dissoc! work unchanged since the real [nil v] pair is kept
as the stored value, and phm already has its own has-nil slot. The transient
set has the analogous bug (needs phs nil support) — filed separately.

Co-authored-by: Yogthos <yogthos@gmail.com>
This commit is contained in:
Dmitri Sotnikov 2026-06-16 23:10:08 +00:00 committed by GitHub
parent 1873736045
commit 2668a76837
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 56 additions and 9 deletions

View file

@ -73,9 +73,32 @@
(persistent!
(reduce (fn [counts x] (assoc! counts x (inc (get counts x 0)))) (transient {}) coll)))
;; Buckets are transient vectors, not persistent ones: the JVM form rebuilds the
;; bucket's persistent vector per element (conj (get ret k []) x), an O(log n)
;; trie path-rebuild + alloc per element — so a coarse grouping (few large
;; buckets) is bound on that conj, not the map build. Push onto a per-bucket
;; native array (O(1)) instead, then bulk-build the persistent map ONCE.
;; Distinct keys are recorded in a side vector so the buckets can be frozen in
;; place (no second map rebuild). A bucket's FIRST element is stored as a cheap
;; persistent [x]; only the second element promotes it to a transient — so an
;; all-singletons grouping pays no transient alloc and matches the old cost,
;; while any bucket that actually grows rides the O(1) push.
(defn group-by [f coll]
(persistent!
(reduce (fn [ret x] (let [k (f x)] (assoc! ret k (conj (get ret k []) x)))) (transient {}) coll)))
(let [tm (transient {})
ks (reduce (fn [ks x]
(let [k (f x)
b (get tm k)]
(if (nil? b)
(do (assoc! tm k [x]) (conj! ks k))
(if (vector? b)
(do (assoc! tm k (conj! (transient b) x)) ks)
(do (conj! b x) ks)))))
(transient []) coll)]
(reduce (fn [_ k]
(let [b (get tm k)]
(if (vector? b) nil (assoc! tm k (persistent! b)))))
nil (persistent! ks))
(persistent! tm)))
(defn not-empty [coll]
(if (or (nil? coll) (zero? (count coll))) nil coll))