index/memory: replace red-black tree with persistent B+-tree
ober
a68aa9af6d4d61a9145c85579b3f34ead6d7b6a2
--- a/lib/jerboa-db/index/memory.ss +++ b/lib/jerboa-db/index/memory.ss @@ -1,9 +1,18 @@ #!chezscheme -;;; (jerboa-db index memory) — In-memory index backend using sorted maps +;;; (jerboa-db index memory) — In-memory index backend (persistent B+-tree) ;;; -;;; Each index is a mutable cell holding a sorted-map. Datoms are the keys; -;;; values are #t (we only need key presence). Range scans use sorted-map-range -;;; with probe datoms at the boundaries. +;;; Each index is a mutable cell holding a persistent sorted set of datoms, +;;; implemented as a B+-tree (btset-style, after Datomic's datomic.btset). +;;; Datoms are the elements; set semantics (a duplicate is a no-op). Range +;;; scans walk the tree in order; counts use order-statistic subtree sizes. +;;; +;;; Why a B+-tree instead of a red-black tree: width-32 array nodes give a +;;; depth of ~log32(n) (≈4 at 1M) versus ~40 for a 2-way RB-tree, each insert +;;; path-copies a handful of flat vectors (cache-friendly, low GC) rather than +;;; allocating many 5-slot cells plus rotations, and an append-to-full-leaf +;;; fast path makes in-order bulk insert near-linear. Branch nodes carry a +;;; subtree size, so `count` over a key range is O(log n) (order-statistic +;;; rank) rather than O(range). (library (jerboa-db index memory) (export make-mem-index-set) @@ -22,143 +31,274 @@ (jerboa-db datom) (jerboa-db index protocol)) - ;; ---- Sorted map (minimal in-library implementation) ---- - ;; We use a red-black tree to avoid external module dependencies in Phase 1. - ;; Keys are datoms compared with an index-specific comparator. + ;; ---- B+-tree (persistent sorted set of datoms) ---- + ;; + ;; Nodes are tagged vectors: + ;; leaf: #(leaf keys) keys = sorted vector of datoms, 1..ORDER + ;; branch: #(branch seps kids size) kids = child vector (>= 2), + ;; seps = routing keys, |seps| = |kids|-1, + ;; seps[i] = min datom of kids[i+1], + ;; size = total datoms in this subtree + ;; The empty tree is #f. A node never under/over-flows search correctness; + ;; insert returns either a node or a split token #(split left sep right). + + (def +bt-order+ 32) ;; max keys per leaf / max children per branch + + ;; -- small immutable-vector helpers -- + + (def (vec-slice v start end) ;; copy [start, end) + (let ([out (make-vector (- end start))]) + (do ([i start (+ i 1)] [j 0 (+ j 1)]) + ((= i end) out) + (vector-set! out j (vector-ref v i))))) + + (def (vec-set-copy v i x) ;; copy with index i replaced by x + (let ([out (vec-slice v 0 (vector-length v))]) + (vector-set! out i x) + out)) + + (def (vec-insert v i x) ;; copy with x inserted at index i + (let* ([n (vector-length v)] [out (make-vector (+ n 1))]) + (do ([k 0 (+ k 1)]) ((= k i)) (vector-set! out k (vector-ref v k))) + (vector-set! out i x) + (do ([k i (+ k 1)]) ((= k n)) (vector-set! out (+ k 1) (vector-ref v k))) + out)) + + (def (vec-remove v i) ;; copy with index i removed + (let* ([n (vector-length v)] [out (make-vector (- n 1))]) + (do ([k 0 (+ k 1)]) ((= k i)) (vector-set! out k (vector-ref v k))) + (do ([k (+ i 1) (+ k 1)]) ((= k n)) (vector-set! out (- k 1) (vector-ref v k))) + out)) + + ;; Binary search over a sorted vector. strict?=#f -> lower bound (first index + ;; with vec[i] >= key, i.e. count of elements < key); strict?=#t -> upper + ;; bound (first index with vec[i] > key, i.e. count of elements <= key). + (def (vec-bsearch vec key cmp strict?) + (let ([n (vector-length vec)]) + (let loop ([lo 0] [hi n]) + (if (= lo hi) + lo + (let* ([mid (quotient (+ lo hi) 2)] + [c (cmp (vector-ref vec mid) key)]) + (if (if strict? (> c 0) (>= c 0)) + (loop lo mid) + (loop (+ mid 1) hi))))))) + + ;; -- node constructors / accessors -- + + (def (mk-leaf keys) (vector 'leaf keys)) + (def (leaf? n) (eq? (vector-ref n 0) 'leaf)) + (def (leaf-keys n) (vector-ref n 1)) + + (def (mk-branch seps kids size) (vector 'branch seps kids size)) + (def (branch-seps n) (vector-ref n 1)) + (def (branch-kids n) (vector-ref n 2)) + (def (branch-size n) (vector-ref n 3)) + + (def (node-size n) + (cond [(not n) 0] + [(leaf? n) (vector-length (leaf-keys n))] + [else (branch-size n)])) - ;; RB-tree node: #f (empty) or (color left key value right) - ;; color: 'R or 'B + (def (sum-sizes kids) + (let ([n (vector-length kids)]) + (let loop ([i 0] [acc 0]) + (if (= i n) acc (loop (+ i 1) (+ acc (node-size (vector-ref kids i)))))))) - (def (rb-empty) #f) - (def (rb-empty? t) (not t)) - (def (rb-node color left key val right) (vector color left key val right)) - (def (rb-color n) (vector-ref n 0)) - (def (rb-left n) (vector-ref n 1)) - (def (rb-key n) (vector-ref n 2)) - (def (rb-val n) (vector-ref n 3)) - (def (rb-right n) (vector-ref n 4)) + ;; split token + (def (mk-split left sep right) (vector 'split left sep right)) + (def (split? x) (and (vector? x) (eq? (vector-ref x 0) 'split))) + (def (split-left x) (vector-ref x 1)) + (def (split-sep x) (vector-ref x 2)) + (def (split-right x) (vector-ref x 3)) - (def (rb-balance color left key val right) - ;; Okasaki balance for red-black trees + ;; -- insert (set semantics; duplicate -> same tree) -- + + (def (bt-insert tree key cmp) + (define (ins node) + (if (leaf? node) (ins-leaf node) (ins-branch node))) + (define (ins-leaf node) + (let* ([keys (leaf-keys node)] + [n (vector-length keys)] + [i (vec-bsearch keys key cmp #f)]) + (cond + ;; already present -> no change + [(and (< i n) (= (cmp (vector-ref keys i) key) 0)) node] + ;; append to a full leaf -> cheap right sibling, no array copy + [(and (= i n) (= n +bt-order+)) + (mk-split node key (mk-leaf (vector key)))] + [else + (let ([newkeys (vec-insert keys i key)]) + (if (<= (vector-length newkeys) +bt-order+) + (mk-leaf newkeys) + (split-leaf newkeys)))]))) + (define (split-leaf keys) + (let* ([n (vector-length keys)] + [mid (quotient n 2)] + [left (vec-slice keys 0 mid)] + [right (vec-slice keys mid n)]) + (mk-split (mk-leaf left) (vector-ref right 0) (mk-leaf right)))) + (define (ins-branch node) + (let* ([seps (branch-seps node)] + [kids (branch-kids node)] + [ci (vec-bsearch seps key cmp #t)] + [child (vector-ref kids ci)] + [res (ins child)]) + (cond + [(split? res) + (let* ([newseps (vec-insert seps ci (split-sep res))] + [k1 (vec-set-copy kids ci (split-left res))] + [newkids (vec-insert k1 (+ ci 1) (split-right res))]) + (if (<= (vector-length newkids) +bt-order+) + (mk-branch newseps newkids (+ (branch-size node) 1)) + (split-branch newseps newkids)))] + [(eq? res child) node] ;; key already present + [else (mk-branch seps (vec-set-copy kids ci res) + (+ (branch-size node) 1))]))) + (define (split-branch seps kids) + (let* ([nk (vector-length kids)] + [mid (quotient nk 2)] + [up (vector-ref seps (- mid 1))] ;; median moves up + [lkids (vec-slice kids 0 mid)] + [lseps (vec-slice seps 0 (- mid 1))] + [rkids (vec-slice kids mid nk)] + [rseps (vec-slice seps mid (vector-length seps))]) + (mk-split (mk-branch lseps lkids (sum-sizes lkids)) + up + (mk-branch rseps rkids (sum-sizes rkids))))) + (if (not tree) + (mk-leaf (vector key)) + (let ([res (ins tree)]) + (if (split? res) + (mk-branch (vector (split-sep res)) + (vector (split-left res) (split-right res)) + (+ (node-size (split-left res)) + (node-size (split-right res)))) + res)))) + + ;; -- delete (no strict rebalance: drops empty leaves, collapses 1-child + ;; branches; harmless under-full nodes are fine since the whole memory + ;; index is discarded on flush to the durable layer). Returns the same + ;; tree (eq?) if the key was absent; #f if the tree became empty. -- + + (def (bt-delete tree key cmp) + (define (del node) + (if (leaf? node) (del-leaf node) (del-branch node))) + (define (del-leaf node) + (let* ([keys (leaf-keys node)] + [n (vector-length keys)] + [i (vec-bsearch keys key cmp #f)]) + (if (and (< i n) (= (cmp (vector-ref keys i) key) 0)) + (if (= n 1) #f (mk-leaf (vec-remove keys i))) + node))) + (define (del-branch node) + (let* ([seps (branch-seps node)] + [kids (branch-kids node)] + [ci (vec-bsearch seps key cmp #t)] + [child (vector-ref kids ci)] + [res (del child)]) + (cond + [(eq? res child) node] ;; absent + [res (mk-branch seps (vec-set-copy kids ci res) + (- (branch-size node) 1))] + [else ;; child emptied + (let* ([newkids (vec-remove kids ci)] + [sep-idx (if (= ci 0) 0 (- ci 1))] + [newseps (if (= (vector-length seps) 0) + seps (vec-remove seps sep-idx))]) + (if (= (vector-length newkids) 1) + (vector-ref newkids 0) ;; collapse + (mk-branch newseps newkids (- (branch-size node) 1))))]))) + (if (not tree) + #f + (let ([res (del tree)]) + (cond [(eq? res tree) tree] + [res res] + [else #f])))) + + ;; -- in-order traversal -- + + (def (bt-fold proc init tree) (cond - ;; Case 1: left-left red - [(and (eq? color 'B) (not (rb-empty? left)) (eq? (rb-color left) 'R) - (not (rb-empty? (rb-left left))) (eq? (rb-color (rb-left left)) 'R)) - (rb-node 'R - (rb-node 'B (rb-left (rb-left left)) - (rb-key (rb-left left)) (rb-val (rb-left left)) - (rb-right (rb-left left))) - (rb-key left) (rb-val left) - (rb-node 'B (rb-right left) key val right))] - ;; Case 2: left-right red - [(and (eq? color 'B) (not (rb-empty? left)) (eq? (rb-color left) 'R) - (not (rb-empty? (rb-right left))) (eq? (rb-color (rb-right left)) 'R)) - (rb-node 'R - (rb-node 'B (rb-left left) (rb-key left) (rb-val left) - (rb-left (rb-right left))) - (rb-key (rb-right left)) (rb-val (rb-right left)) - (rb-node 'B (rb-right (rb-right left)) key val right))] - ;; Case 3: right-left red - [(and (eq? color 'B) (not (rb-empty? right)) (eq? (rb-color right) 'R) - (not (rb-empty? (rb-left right))) (eq? (rb-color (rb-left right)) 'R)) - (rb-node 'R - (rb-node 'B left key val (rb-left (rb-left right))) - (rb-key (rb-left right)) (rb-val (rb-left right)) - (rb-node 'B (rb-right (rb-left right)) - (rb-key right) (rb-val right) - (rb-right right)))] - ;; Case 4: right-right red - [(and (eq? color 'B) (not (rb-empty? right)) (eq? (rb-color right) 'R) - (not (rb-empty? (rb-right right))) (eq? (rb-color (rb-right right)) 'R)) - (rb-node 'R - (rb-node 'B left key val (rb-left right)) - (rb-key right) (rb-val right) - (rb-node 'B (rb-left (rb-right right)) - (rb-key (rb-right right)) (rb-val (rb-right right)) - (rb-right (rb-right right))))] - [else (rb-node color left key val right)])) - - (def (rb-insert tree key val cmp) - (define (ins t) - (if (rb-empty? t) - (rb-node 'R (rb-empty) key val (rb-empty)) - (let ([c (cmp key (rb-key t))]) - (cond - [(< c 0) (rb-balance (rb-color t) - (ins (rb-left t)) (rb-key t) (rb-val t) (rb-right t))] - [(> c 0) (rb-balance (rb-color t) - (rb-left t) (rb-key t) (rb-val t) (ins (rb-right t)))] - [else t])))) ;; duplicate — no change - (let ([r (ins tree)]) - (rb-node 'B (rb-left r) (rb-key r) (rb-val r) (rb-right r)))) - - (def (rb-delete tree key cmp) - ;; Simplified delete — rebuild without the key - ;; For Phase 1 correctness, this is O(n) but works reliably. - (let ([pairs '()]) - (rb-fold (lambda (k v acc) - (if (= (cmp k key) 0) acc (cons (cons k v) acc))) - '() tree) - (let ([filtered (rb-fold (lambda (k v acc) - (if (= (cmp k key) 0) acc (cons (cons k v) acc))) - '() tree)]) - (fold-left (lambda (t pair) (rb-insert t (car pair) (cdr pair) cmp)) - (rb-empty) filtered)))) - - ;; In-order fold: (proc key value accumulator) -> accumulator - (def (rb-fold proc init tree) - (if (rb-empty? tree) - init - (let* ([left-result (rb-fold proc init (rb-left tree))] - [mid-result (proc (rb-key tree) (rb-val tree) left-result)]) - (rb-fold proc mid-result (rb-right tree))))) - - ;; Range fold: fold over keys in [lo, hi] (inclusive) by comparator - (def (rb-range-fold proc init tree lo hi cmp) - (if (rb-empty? tree) - init - (let ([k (rb-key tree)] - [v (rb-val tree)]) - (let ([cmp-lo (cmp k lo)] - [cmp-hi (cmp k hi)]) - (let* ([acc (if (>= cmp-lo 0) - (rb-range-fold proc init (rb-left tree) lo hi cmp) - init)] - [acc (if (and (>= cmp-lo 0) (<= cmp-hi 0)) - (proc k v acc) - acc)] - [acc (if (<= cmp-hi 0) - (rb-range-fold proc acc (rb-right tree) lo hi cmp) - acc)]) - acc))))) - - ;; Collect all keys in-order - (def (rb-keys tree) - (reverse (rb-fold (lambda (k v acc) (cons k acc)) '() tree))) - - ;; Count nodes - (def (rb-size tree) - (rb-fold (lambda (k v acc) (+ acc 1)) 0 tree)) + [(not tree) init] + [(leaf? tree) + (let ([keys (leaf-keys tree)]) + (let loop ([i 0] [acc init]) + (if (= i (vector-length keys)) + acc + (loop (+ i 1) (proc (vector-ref keys i) acc)))))] + [else + (let ([kids (branch-kids tree)]) + (let loop ([i 0] [acc init]) + (if (= i (vector-length kids)) + acc + (loop (+ i 1) (bt-fold proc acc (vector-ref kids i))))))])) + + (def (bt->list tree) + (reverse (bt-fold (lambda (k acc) (cons k acc)) '() tree))) + + ;; Fold over keys k with lo <= k <= hi (inclusive), ascending. + (def (bt-range-fold proc init tree lo hi cmp) + (cond + [(not tree) init] + [(leaf? tree) + (let* ([keys (leaf-keys tree)] + [n (vector-length keys)] + [start (vec-bsearch keys lo cmp #f)]) + (let loop ([i start] [acc init]) + (if (or (= i n) (> (cmp (vector-ref keys i) hi) 0)) + acc + (loop (+ i 1) (proc (vector-ref keys i) acc)))))] + [else + (let* ([seps (branch-seps tree)] + [kids (branch-kids tree)] + [nk (vector-length kids)]) + (let loop ([i 0] [acc init]) + (if (= i nk) + acc + (let ([left-sep (if (= i 0) #f (vector-ref seps (- i 1)))] + [right-sep (if (< i (- nk 1)) (vector-ref seps i) #f)]) + (cond + ;; child's min > hi -> all remaining children are larger too + [(and left-sep (> (cmp left-sep hi) 0)) acc] + ;; child's upper bound <= lo -> entirely below the range + [(and right-sep (<= (cmp right-sep lo) 0)) (loop (+ i 1) acc)] + [else (loop (+ i 1) + (bt-range-fold proc acc (vector-ref kids i) + lo hi cmp))])))))])) + + ;; Order-statistic rank: count of keys < key (le?=#f) or <= key (le?=#t). + (def (bt-rank tree key cmp le?) + (cond + [(not tree) 0] + [(leaf? tree) (vec-bsearch (leaf-keys tree) key cmp le?)] + [else + (let* ([seps (branch-seps tree)] + [kids (branch-kids tree)] + [ci (vec-bsearch seps key cmp #t)]) + (let loop ([i 0] [acc 0]) + (if (= i ci) + (+ acc (bt-rank (vector-ref kids ci) key cmp le?)) + (loop (+ i 1) (+ acc (node-size (vector-ref kids i)))))))])) + + ;; Count of keys in [lo, hi] inclusive, O(log n). + (def (bt-range-count tree lo hi cmp) + (- (bt-rank tree hi cmp #t) (bt-rank tree lo cmp #f))) ;; ---- Memory index implementation ---- ;; - ;; LSM-style two-tier layout: a base RB-tree of flushed datoms plus a - ;; staging buffer of pending writes. `add!` is O(1) (cons); reads merge - ;; base + buffer. The buffer auto-flushes once it crosses - ;; `+staging-threshold+`, amortising the RB-tree insert cost across - ;; many datoms instead of paying it per call. - ;; - ;; This collapses the per-datom restructuring cost: previously, loading - ;; 147K datoms into 3 indices issued ~440K rb-insert calls, each doing - ;; a recursive structural copy + balance. With the buffer, the heavy - ;; work happens once per flush window rather than once per datom. + ;; Two-tier layout: a persistent B+-tree of flushed datoms plus a staging + ;; buffer of pending writes. `add!` is O(1) (cons); reads merge base + + ;; buffer. The buffer auto-flushes once it crosses `+staging-threshold+`, + ;; and on flush the pending datoms are sorted into this index's order before + ;; bulk insert so they hit the B+-tree's append-to-full-leaf fast path. (def +staging-threshold+ 4096) (def (make-mem-index name comparator) - (let ([tree-cell (list (rb-empty))] ;; sorted base - [buffer-cell (list '())] ;; LIFO list of pending datoms - [bufcnt-cell (list 0)]) ;; size of buffer-cell + (let ([tree-cell (list #f)] ;; B+-tree root (#f = empty) + [buffer-cell (list '())] ;; LIFO list of pending datoms + [bufcnt-cell (list 0)]) ;; size of buffer-cell (define (get-tree) (car tree-cell)) (define (set-tree! t) (set-car! tree-cell t)) (define (get-buf) (car buffer-cell)) @@ -169,14 +309,13 @@ (define (flush!) (let ([buf (get-buf)]) (unless (null? buf) - ;; Sort buffer descending by comparator so that rb-insert sees - ;; them in some order (tree balancing handles any order). We - ;; just need to fold them in. - (set-tree! - (fold-left - (lambda (t d) (rb-insert t d #t comparator)) - (get-tree) - buf)) + ;; Insert in this index's sort order so runs append to full leaves. + (let ([sorted (list-sort (lambda (a b) (< (comparator a b) 0)) buf)]) + (set-tree! + (fold-left + (lambda (t d) (bt-insert t d comparator)) + (get-tree) + sorted))) (set-buf! '()) (set-bufcnt! 0)))) @@ -189,7 +328,7 @@ (define (remove! datom) ;; Remove must see the unified state. Flush first. (flush!) - (set-tree! (rb-delete (get-tree) datom comparator))) + (set-tree! (bt-delete (get-tree) datom comparator))) (define (in-range? d lo hi) (and (>= (comparator d lo) 0) @@ -220,8 +359,8 @@ (define (range-query start end) (let* ([base-list (reverse - (rb-range-fold - (lambda (k v acc) (cons k acc)) + (bt-range-fold + (lambda (k acc) (cons k acc)) '() (get-tree) start end comparator))] [buf-list (buffer-in-range start end)]) (cond @@ -250,16 +389,20 @@ (range-query lo hi)))) (define (count-range start end) - ;; range-query already dedupes; safe to count its result. - (length (range-query start end))) + ;; Flush so the B+-tree holds everything, then use order-statistic + ;; rank for an O(log n) count instead of materialising the range. + (flush!) + (bt-range-count (get-tree) start end comparator)) (define (snapshot) + ;; The B+-tree is persistent/immutable, so the flushed root is a valid + ;; point-in-time snapshot with full structural sharing. (flush!) (get-tree)) (define (all-datoms) (flush!) - (rb-keys (get-tree))) + (bt->list (get-tree))) (make-dbi name add! remove! range-query seek count-range snapshot all-datoms)))