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

@ -182,7 +182,7 @@
(if (core-transient? m)
(case (m :kind)
:vector (if (and (number? k) (>= k 0) (< k (length (m :arr)))) (in (m :arr) k) default)
:map (let [p (get (m :tbl) (canon-key k))] (if p (in p 1) default))
:map (let [p (get (m :tbl) (tbl-key k))] (if p (in p 1) default))
:set (if (nil? (get (m :tbl) (canon-key k))) default k))
(if (set? m) (phs-get m k default)
(if (phm? m) (phm-get m k default)
@ -251,7 +251,7 @@
(if (core-transient? coll)
(case (coll :kind)
:vector (and (number? key) (>= key 0) (< key (length (coll :arr))))
(not (nil? (get (coll :tbl) (canon-key key)))))
(not (nil? (get (coll :tbl) (if (= :map (coll :kind)) (tbl-key key) (canon-key key))))))
(if (set? coll) (phs-contains? coll key)
(if (phm? coll) (phm-contains? coll key)
(if (pvec? coll) (and (number? key) (>= key 0) (< key (pv-count coll)))

View file

@ -343,7 +343,7 @@
(or (phm? coll) (and (struct? coll) (nil? (get coll :jolt/type))))
(let [t @{}]
(each pair (realize-for-iteration coll)
(put t (canon-key (in pair 0)) @[(in pair 0) (in pair 1)]))
(put t (tbl-key (in pair 0)) @[(in pair 0) (in pair 1)]))
@{:jolt/type :jolt/transient :kind :map :tbl t})
# mutable-build arrays (vectors/lists) — copy into a transient vector
(array? coll) @{:jolt/type :jolt/transient :kind :vector :arr (array/slice coll)}
@ -365,11 +365,11 @@
# a [k v] pair (map-entry / 2-vector)
(and (or (pvec? x) (tuple? x) (array? x))
(= 2 (if (pvec? x) (pv-count x) (length x))))
(put (t :tbl) (canon-key (vnth x 0)) @[(vnth x 0) (vnth x 1)])
(put (t :tbl) (tbl-key (vnth x 0)) @[(vnth x 0) (vnth x 1)])
# a map: merge all its entries
(or (phm? x) (and (struct? x) (nil? (get x :jolt/type))))
(each e (map-entries-of x)
(put (t :tbl) (canon-key (in e 0)) @[(in e 0) (in e 1)]))
(put (t :tbl) (tbl-key (in e 0)) @[(in e 0) (in e 1)]))
(error "conj! on a transient map requires a [key value] pair or a map")))
t)
@ -380,7 +380,7 @@
(when (not (and (number? k) (= k (math/floor k)) (>= k 0) (<= k (length a))))
(error (string "Index " k " out of bounds for assoc! on a transient vector of length " (length a))))
(if (= k (length a)) (array/push a v) (put a k v)))
:map (put (t :tbl) (canon-key k) @[k v])
:map (put (t :tbl) (tbl-key k) @[k v])
(error "assoc! expects a transient vector or map"))
t)
@ -405,7 +405,7 @@
(defn core-dissoc! [t & ks]
(if (and (core-transient? t) (= :map (t :kind)))
(do (tr-check-active! t) (each k ks (put (t :tbl) (canon-key k) nil)) t)
(do (tr-check-active! t) (each k ks (put (t :tbl) (tbl-key k) nil)) t)
(error "dissoc! requires a transient map")))
(defn core-disj! [t & xs]

View file

@ -96,6 +96,14 @@
k)))
(set-canonicalize-key! canon-key)
# Janet tables silently drop a nil key (put/get with nil is a no-op), but Clojure
# maps allow a nil key. The transient map keys its native table by canon-key, and
# canon-key returns nil only for nil input — so route nil to a unique sentinel.
# The sentinel is a fresh mutable table; canon-key never produces one, so it can't
# collide with the canon-key of any real key. (phm keeps its own has-nil slot.)
(def tbl-nil-key @{})
(defn tbl-key [k] (if (nil? k) tbl-nil-key (canon-key k)))
# All [k v] entries of a map (struct or phm), nil-valued keys included. Use this
# instead of (keys (phm-to-struct m)) — phm-to-struct drops keys whose value is
# nil, which is exactly what Clojure maps must keep.