segment: add columnar (transposed) datom segments

ober

4b628cc316281042a04a865b9ac2bad756a30c9f

diff --git a/lib/jerboa-db/segment.ss b/lib/jerboa-db/segment.ss
new file mode 100644
index 0000000..09ea428
--- /dev/null
+++ b/lib/jerboa-db/segment.ss
@@ -0,0 +1,275 @@
+#!chezscheme
+;;; (jerboa-db segment) — Columnar (transposed) datom segments
+;;;
+;;; A segment is an immutable batch of datoms stored column-by-column
+;;; (struct-of-arrays), after Datomic's datomic.index/TransposedData. Instead
+;;; of one boxed record per datom (array-of-structs), each field is its own
+;;; primitive array:
+;;;
+;;;   es  : fxvector   entity ids
+;;;   as  : fxvector   attribute ids
+;;;   txs : fxvector   transaction ids
+;;;   ops : bytevector 1=assertion 0=retraction
+;;;   vs  : a TYPE-SPECIALISED value column —
+;;;           'double -> flvector  (unboxed flonums)
+;;;           'long   -> fxvector  (unboxed fixnums; covers ref/instant too)
+;;;           'mixed  -> vector    (boxed; heterogeneous/strings/bools)
+;;;
+;;; Two payoffs over per-datom storage:
+;;;   1. Scans/aggregates over a homogeneous numeric column run as a tight
+;;;      primitive loop with no boxing (the Q4/Q8 analytic case).
+;;;   2. Serialisation is far smaller: columns delta+varint compress (sorted
+;;;      e/tx runs shrink to ~1 byte each) and there is no per-datom record
+;;;      overhead.
+;;;
+;;; This is the storage substrate; wiring it under a durable segment tree and a
+;;; native group-by operator are follow-ons.
+;;;
+;;; NOTE: es/as/txs assume fixnum-range ids (true for jerboa-db's counter-based
+;;; eids/txs). A 'long value column falls back to 'mixed if any value is a
+;;; bignum, so correctness is preserved for out-of-range values.
+
+(library (jerboa-db segment)
+  (export
+    make-segment segment? segment-count segment-vtype
+    segment-e segment-a segment-v segment-tx segment-added? segment-ref
+    segment-min-datom segment-max-datom
+    segment-fold segment-sum-v segment-min-v segment-max-v segment-avg-v
+    segment->bytevector bytevector->segment)
+
+  (import (except (chezscheme)
+                  make-hash-table hash-table?
+                  sort sort!
+                  printf fprintf
+                  path-extension path-absolute?
+                  with-input-from-string with-output-to-string
+                  iota 1+ 1-
+                  partition
+                  make-date make-time
+                atom? meta)
+          (jerboa prelude)
+          (jerboa-db datom))
+
+  (define-record-type seg
+    (fields cnt es as vs vtype txs ops))
+
+  (def (segment? x) (seg? x))
+  (def (segment-count s) (seg-cnt s))
+  (def (segment-vtype s) (seg-vtype s))
+
+  ;; ---- Build ----
+
+  ;; Infer the value column type and fill the matching primitive array.
+  (def (build-vcol v n)
+    (let loop ([i 0] [all-fx #t] [all-fl #t])
+      (cond
+        [(= n 0) (values 'mixed (vector))]
+        [(< i n)
+         (let ([val (datom-v (vector-ref v i))])
+           (loop (+ i 1) (and all-fx (fixnum? val)) (and all-fl (flonum? val))))]
+        [all-fl
+         (let ([col (make-flvector n 0.0)])
+           (do ([j 0 (+ j 1)]) ((= j n) (values 'double col))
+             (flvector-set! col j (datom-v (vector-ref v j)))))]
+        [all-fx
+         (let ([col (make-fxvector n 0)])
+           (do ([j 0 (+ j 1)]) ((= j n) (values 'long col))
+             (fxvector-set! col j (datom-v (vector-ref v j)))))]
+        [else
+         (let ([col (make-vector n)])
+           (do ([j 0 (+ j 1)]) ((= j n) (values 'mixed col))
+             (vector-set! col j (datom-v (vector-ref v j)))))])))
+
+  ;; Build a segment from a list of datoms (kept in the given order).
+  (def (make-segment datom-list)
+    (let* ([v (list->vector datom-list)]
+           [n (vector-length v)]
+           [es (make-fxvector n 0)]
+           [as (make-fxvector n 0)]
+           [txs (make-fxvector n 0)]
+           [ops (make-bytevector n 0)])
+      (do ([i 0 (+ i 1)]) ((= i n))
+        (let ([d (vector-ref v i)])
+          (fxvector-set! es i (datom-e d))
+          (fxvector-set! as i (datom-a d))
+          (fxvector-set! txs i (datom-tx d))
+          (bytevector-u8-set! ops i (if (datom-added? d) 1 0))))
+      (let-values ([(vtype vs) (build-vcol v n)])
+        (make-seg n es as vs vtype txs ops))))
+
+  ;; ---- Columnar accessors ----
+
+  (def (segment-e s i) (fxvector-ref (seg-es s) i))
+  (def (segment-a s i) (fxvector-ref (seg-as s) i))
+  (def (segment-tx s i) (fxvector-ref (seg-txs s) i))
+  (def (segment-added? s i) (= 1 (bytevector-u8-ref (seg-ops s) i)))
+
+  (def (segment-v s i)
+    (case (seg-vtype s)
+      [(double) (flvector-ref (seg-vs s) i)]
+      [(long)   (fxvector-ref (seg-vs s) i)]
+      [else     (vector-ref (seg-vs s) i)]))
+
+  (def (segment-ref s i)
+    (make-datom (segment-e s i) (segment-a s i) (segment-v s i)
+                (segment-tx s i) (segment-added? s i)))
+
+  (def (segment-min-datom s) (and (> (seg-cnt s) 0) (segment-ref s 0)))
+  (def (segment-max-datom s) (and (> (seg-cnt s) 0) (segment-ref s (- (seg-cnt s) 1))))
+
+  ;; ---- Iteration / aggregates ----
+
+  (def (segment-fold proc init s)
+    (let ([n (seg-cnt s)])
+      (let loop ([i 0] [acc init])
+        (if (= i n) acc (loop (+ i 1) (proc (segment-ref s i) acc))))))
+
+  ;; Vectorised numeric aggregates over the value column (no boxing).
+  (def (segment-sum-v s)
+    (let ([n (seg-cnt s)])
+      (case (seg-vtype s)
+        [(double) (let ([col (seg-vs s)])
+                    (let loop ([i 0] [acc 0.0])
+                      (if (= i n) acc (loop (+ i 1) (fl+ acc (flvector-ref col i))))))]
+        [(long)   (let ([col (seg-vs s)])
+                    (let loop ([i 0] [acc 0])
+                      (if (= i n) acc (loop (+ i 1) (+ acc (fxvector-ref col i))))))]
+        [else (error 'segment-sum-v "value column is not numeric")])))
+
+  (def (segment-avg-v s)
+    (let ([n (seg-cnt s)])
+      (if (= n 0) 0 (/ (segment-sum-v s) n))))
+
+  (def (segment-min-v s)
+    (let ([n (seg-cnt s)])
+      (if (= n 0)
+          #f
+          (let loop ([i 1] [m (segment-v s 0)])
+            (if (= i n) m (loop (+ i 1) (let ([x (segment-v s i)]) (if (< x m) x m))))))))
+
+  (def (segment-max-v s)
+    (let ([n (seg-cnt s)])
+      (if (= n 0)
+          #f
+          (let loop ([i 1] [m (segment-v s 0)])
+            (if (= i n) m (loop (+ i 1) (let ([x (segment-v s i)]) (if (> x m) x m))))))))
+
+  ;; ---- Compact serialisation ----
+  ;; Columns are delta+zigzag+varint encoded (sorted runs shrink to ~1 byte),
+  ;; ops are bit-packed, double values are raw 8-byte little-endian, mixed
+  ;; values are length-prefixed FASL.
+
+  (def (zigzag n)   (if (< n 0) (- (* n -2) 1) (* n 2)))
+  (def (unzigzag z) (if (even? z) (quotient z 2) (- (quotient (+ z 1) 2))))
+
+  (def (put-uvarint port n)
+    (let loop ([n n])
+      (if (< n 128)
+          (put-u8 port n)
+          (begin
+            (put-u8 port (bitwise-ior (bitwise-and n 127) 128))
+            (loop (bitwise-arithmetic-shift-right n 7))))))
+
+  (def (get-uvarint port)
+    (let loop ([shift 0] [acc 0])
+      (let ([b (get-u8 port)])
+        (let ([acc (bitwise-ior acc (bitwise-arithmetic-shift-left (bitwise-and b 127) shift))])
+          (if (< b 128) acc (loop (+ shift 7) acc))))))
+
+  (def (vtype-code t) (case t [(mixed) 0] [(long) 1] [(double) 2]))
+  (def (vtype-decode c) (case c [(0) 'mixed] [(1) 'long] [(2) 'double]))
+
+  (def (put-delta-col port n ref)
+    (let loop ([i 0] [prev 0])
+      (when (< i n)
+        (let ([x (ref i)])
+          (put-uvarint port (zigzag (- x prev)))
+          (loop (+ i 1) x)))))
+
+  (def (get-delta-col port n set!)
+    (let loop ([i 0] [prev 0])
+      (when (< i n)
+        (let ([x (+ prev (unzigzag (get-uvarint port)))])
+          (set! i x)
+          (loop (+ i 1) x)))))
+
+  (def (segment->bytevector s)
+    (let-values ([(port extract) (open-bytevector-output-port)])
+      (let ([n (seg-cnt s)])
+        (put-u8 port 1)                       ;; format version
+        (put-uvarint port n)
+        (put-u8 port (vtype-code (seg-vtype s)))
+        (put-delta-col port n (lambda (i) (fxvector-ref (seg-es s) i)))
+        (put-delta-col port n (lambda (i) (fxvector-ref (seg-as s) i)))
+        (put-delta-col port n (lambda (i) (fxvector-ref (seg-txs s) i)))
+        ;; ops: bit-packed, 8 per byte
+        (let loop ([i 0])
+          (when (< i n)
+            (let bit ([j i] [k 0] [b 0])
+              (if (or (= k 8) (= j n))
+                  (put-u8 port b)
+                  (bit (+ j 1) (+ k 1)
+                       (bitwise-ior b (bitwise-arithmetic-shift-left
+                                        (bytevector-u8-ref (seg-ops s) j) k)))))
+            (loop (+ i 8))))
+        ;; value column
+        (case (seg-vtype s)
+          [(long) (put-delta-col port n (lambda (i) (fxvector-ref (seg-vs s) i)))]
+          [(double)
+           (let ([col (seg-vs s)] [tmp (make-bytevector 8)])
+             (do ([i 0 (+ i 1)]) ((= i n))
+               (bytevector-ieee-double-set! tmp 0 (flvector-ref col i) (endianness little))
+               (put-bytevector port tmp)))]
+          [else
+           (let ([col (seg-vs s)])
+             (do ([i 0 (+ i 1)]) ((= i n))
+               (let-values ([(p g) (open-bytevector-output-port)])
+                 (fasl-write (vector-ref col i) p)
+                 (let ([b (g)])
+                   (put-uvarint port (bytevector-length b))
+                   (put-bytevector port b)))))]))
+      (extract)))
+
+  (def (bytevector->segment bv)
+    (let ([port (open-bytevector-input-port bv)])
+      (get-u8 port)                           ;; version
+      (let* ([n (get-uvarint port)]
+             [vt (vtype-decode (get-u8 port))]
+             [es (make-fxvector n 0)]
+             [as (make-fxvector n 0)]
+             [txs (make-fxvector n 0)]
+             [ops (make-bytevector n 0)])
+        (get-delta-col port n (lambda (i x) (fxvector-set! es i x)))
+        (get-delta-col port n (lambda (i x) (fxvector-set! as i x)))
+        (get-delta-col port n (lambda (i x) (fxvector-set! txs i x)))
+        (let loop ([i 0])
+          (when (< i n)
+            (let ([b (get-u8 port)])
+              (let bit ([j i] [k 0])
+                (when (and (< j n) (< k 8))
+                  (bytevector-u8-set! ops j (bitwise-and (bitwise-arithmetic-shift-right b k) 1))
+                  (bit (+ j 1) (+ k 1)))))
+            (loop (+ i 8))))
+        (let ([vs (case vt
+                    [(long)
+                     (let ([col (make-fxvector n 0)])
+                       (let loop ([i 0] [prev 0])
+                         (if (= i n)
+                             col
+                             (let ([x (+ prev (unzigzag (get-uvarint port)))])
+                               (fxvector-set! col i x)
+                               (loop (+ i 1) x)))))]
+                    [(double)
+                     (let ([col (make-flvector n 0.0)])
+                       (do ([i 0 (+ i 1)]) ((= i n) col)
+                         (let ([b (get-bytevector-n port 8)])
+                           (flvector-set! col i (bytevector-ieee-double-ref b 0 (endianness little))))))]
+                    [else
+                     (let ([col (make-vector n)])
+                       (do ([i 0 (+ i 1)]) ((= i n) col)
+                         (let* ([len (get-uvarint port)]
+                                [b (get-bytevector-n port len)])
+                           (vector-set! col i (fasl-read (open-bytevector-input-port b))))))])])
+          (make-seg n es as vs vt txs ops)))))
+
+) ;; end library
diff --git a/tests/test-core.ss b/tests/test-core.ss
index cec0c59..7f6ac92 100644
--- a/tests/test-core.ss
+++ b/tests/test-core.ss
@@ -7,7 +7,8 @@
         (jerboa-db spec)
         (jerboa-db value-store)
         (jerboa-db index protocol)
-        (jerboa-db index memory))
+        (jerboa-db index memory)
+        (jerboa-db segment))
 
 ;; ---- Test harness ----
 
@@ -702,6 +703,39 @@
             (list (dbi-cursor a lo hi) (dbi-cursor b lo hi)))))
       '(1 3 4 5 6))))
 
+(test "columnar segment round-trips through serialization"
+  (let* ([ds (list (make-datom 1 7 100 1 #t)
+                   (make-datom 2 7 250 1 #t)
+                   (make-datom 5 7 175 2 #f))]
+         [seg (make-segment ds)]
+         [seg2 (bytevector->segment (segment->bytevector seg))])
+    (assert-equal (segment-vtype seg) 'long)
+    (assert-equal (segment-count seg2) 3)
+    (assert-equal (map (lambda (i) (datom->list (segment-ref seg2 i))) '(0 1 2))
+                  (map datom->list ds))))
+
+(test "columnar segment vectorized aggregate (double column)"
+  (let* ([ds (list (make-datom 1 9 1.0 1 #t)
+                   (make-datom 2 9 2.0 1 #t)
+                   (make-datom 3 9 3.0 1 #t)
+                   (make-datom 4 9 4.0 1 #t))]
+         [seg (make-segment ds)])
+    (assert-equal (segment-vtype seg) 'double)
+    (assert-equal (segment-sum-v seg) 10.0)
+    (assert-equal (segment-avg-v seg) 2.5)
+    (assert-equal (segment-min-v seg) 1.0)
+    (assert-equal (segment-max-v seg) 4.0)))
+
+(test "columnar segment mixed (string) column round-trips"
+  (let* ([ds (list (make-datom 1 3 "alpha" 1 #t)
+                   (make-datom 2 3 "beta" 2 #f))]
+         [seg (make-segment ds)]
+         [seg2 (bytevector->segment (segment->bytevector seg))])
+    (assert-equal (segment-vtype seg) 'mixed)
+    (assert-equal (segment-v seg2 0) "alpha")
+    (assert-equal (segment-added? seg2 1) #f)
+    (assert-equal (segment-v seg2 1) "beta")))
+
 ;; ============================================================
 ;; Report
 ;; ============================================================