code wiki / (root) / sketch_count_sketch.nx

sketch_count_sketch.nx source

↩ module page · 230 lines · 8067 B

1// sketch_count_sketch.nx -- CountSketch (Charikar-Chen-Farach-Colton 2002). 2// 3// Frequency-counting variant that gives UNBIASED estimates via ±1 4// sign hashing. Each (key, count) update adds s_j(key) * count to 5// row j column h_j(key), where s_j ∈ {-1, +1} is a sign hash. 6// Query: median of (s_j(key) * counter[j][h_j(key)]) across rows. 7// 8// COMPLEMENTS CMS (runtime/sketch_cms.nx): 9// - CMS: always OVERESTIMATES (false positives via positive collisions). 10// Use when "upper bound" matters. 11// - CountSketch: UNBIASED. Use when expectation matters (statistical 12// summaries, F_2 estimation, sparse-approximation literature). 13// 14// CAPABILITY STOMP: shipping both means callers pick by error semantics. 15// DataSketches ships CMS as "FrequentLongs"; CountSketch is queued in 16// their docs but not implemented. 17// 18// ERROR BOUND: 19// |estimate(x) - f(x)| <= ε · ||f||_2 / sqrt(d) 20// where ||f||_2 = sqrt(Σ f_i²) is the L2 norm of the frequency vector. 21// For w = O(1/ε²) and d = O(log(1/δ)) rows: confidence 1-δ. 22// 23// Default params: d=5 rows, w=2048 columns -> conf >99%, ε ~ 0.022. 24 25import "syscalls.nx" 26import "sketch_types.nx" 27 28const NX_CS_MIN_D: i64 = 3 29const NX_CS_MAX_D: i64 = 32 30const NX_CS_MIN_W: i64 = 64 31const NX_CS_MAX_W: i64 = 65536 32 33struct CountSketch { 34 counters: *i64, // d * w grid, in i64 (signed) 35 d: i64, 36 w: i64, 37 seed: i64, 38 total: i64, 39 scratch: *i64, // d-sized scratch for median (bits-up: hoisted 40 // from per-call sys_mmap to alloc-time). Saves 41 // one syscall per nx_cs_estimate call. 42} 43 44// === construction ================================================= 45 46func nx_cs_alloc(d: i64, w: i64, seed: i64) -> *CountSketch { 47 if d < NX_CS_MIN_D { return 0 as *CountSketch } 48 if d > NX_CS_MAX_D { return 0 as *CountSketch } 49 if w < NX_CS_MIN_W { return 0 as *CountSketch } 50 if w > NX_CS_MAX_W { return 0 as *CountSketch } 51 if (w & (w - 1)) != 0 { return 0 as *CountSketch } // power of 2 52 let raw: *u8 = sys_mmap(56) 53 let c: *CountSketch = raw as *CountSketch 54 let cells: i64 = d * w 55 let cells_raw: *u8 = sys_mmap(cells * 8) 56 c.counters = cells_raw as *i64 57 var i: i64 = 0 58 while i < cells { 59 c.counters[i] = 0 60 i = i + 1 61 } 62 c.d = d 63 c.w = w 64 c.seed = seed 65 c.total = 0 66 let scratch_raw: *u8 = sys_mmap(d * 8) 67 c.scratch = scratch_raw as *i64 68 return c 69} 70 71// === hash + sign helpers ========================================= 72// 73// Kirsch-Mitzenmacher double-hashing: h_j(x) = (h1 + j*h2) & mask. 74// Sign: top bit of a separately-derived hash, mapped to ±1. 75 76func nx_cs_h1(c: *CountSketch, key: i64) -> i64 { 77 let mixed: i64 = (key * 0x9E3779B97F4A7C15 + c.seed) & 0xFFFFFFFFFFFFFFFF 78 return mixed & 0xFFFFFFFF 79} 80 81func nx_cs_h2(c: *CountSketch, key: i64) -> i64 { 82 let mixed: i64 = (key * 0xBF58476D1CE4E5B9 + c.seed) & 0xFFFFFFFFFFFFFFFF 83 return mixed & 0xFFFFFFFF 84} 85 86func nx_cs_column(c: *CountSketch, key: i64, j: i64) -> i64 { 87 let combined: i64 = (nx_cs_h1(c, key) + j * nx_cs_h2(c, key)) & 0xFFFFFFFF 88 return combined & (c.w - 1) 89} 90 91// Sign for row j: hash key+j and check the top bit. 92func nx_cs_sign(c: *CountSketch, key: i64, j: i64) -> i64 { 93 let mixed: i64 = ((key + j * 0x9E3779B9) * 0xC2B2AE3D27D4EB4F + c.seed) & 0xFFFFFFFFFFFFFFFF 94 if (mixed & (1 << 63)) == 0 { return 1 } 95 return -1 96} 97 98// === cell access ================================================= 99 100func nx_cs_cell_idx(c: *CountSketch, j: i64, col: i64) -> i64 { 101 return j * c.w + col 102} 103 104// === add ========================================================== 105 106func nx_cs_add(c: *CountSketch, key: i64, count: i64) -> i64 { 107 if count == 0 { return 0 } 108 c.total = c.total + count 109 // Bits-up: compute h1 and h2 ONCE per add; were recomputed per row. 110 // Saves (d-1)*2 multiplies per add. 111 let h1: i64 = nx_cs_h1(c, key) 112 let h2: i64 = nx_cs_h2(c, key) 113 let w_mask: i64 = c.w - 1 114 var j: i64 = 0 115 while j < c.d { 116 let col: i64 = ((h1 + j * h2) & 0xFFFFFFFF) & w_mask 117 let sign: i64 = nx_cs_sign(c, key, j) 118 let idx: i64 = j * c.w + col 119 c.counters[idx] = c.counters[idx] + sign * count 120 j = j + 1 121 } 122 return 0 123} 124 125// === query (median of signed reads) ============================== 126// 127// For each row j: signed_read = sign_j(key) * counter[j][h_j(key)] 128// Return median of d signed reads. Median is approximate median via 129// insertion-sort of small d. 130 131func nx_cs_estimate(c: *CountSketch, key: i64) -> i64 { 132 // Bits-up: scratch hoisted to struct; h1+h2 cached. 133 let scratch: *i64 = c.scratch 134 let h1: i64 = nx_cs_h1(c, key) 135 let h2: i64 = nx_cs_h2(c, key) 136 let w_mask: i64 = c.w - 1 137 var j: i64 = 0 138 while j < c.d { 139 let col: i64 = ((h1 + j * h2) & 0xFFFFFFFF) & w_mask 140 let sign: i64 = nx_cs_sign(c, key, j) 141 let idx: i64 = j * c.w + col 142 scratch[j] = sign * c.counters[idx] 143 j = j + 1 144 } 145 // Insertion sort. 146 var i: i64 = 1 147 while i < c.d { 148 let cur: i64 = scratch[i] 149 var k: i64 = i - 1 150 var done: i64 = 0 151 while done == 0 { 152 if k < 0 { done = 1 } 153 if done == 0 { 154 if scratch[k] <= cur { done = 1 } 155 if done == 0 { 156 scratch[k + 1] = scratch[k] 157 k = k - 1 158 } 159 } 160 } 161 scratch[k + 1] = cur 162 i = i + 1 163 } 164 return scratch[c.d / 2] 165} 166 167// === typed envelope =============================================== 168// 169// Error bound is |est - f| <= ||f||_2 * sqrt(2/w) at confidence 170// (3/4)^(d/2). We declare a CONSERVATIVE absolute bound = total/sqrt(w) 171// (this is the worst case when L2 = total, i.e. one massive item). 172// param_a holds the bound; conf = (3/4)^(d/2) tabulated. 173 174func nx_cs_isqrt(x: i64) -> i64 { 175 if x < 0 { return 0 } 176 if x == 0 { return 0 } 177 if x < 4 { return 1 } 178 var g: i64 = (x >> 1) + 1 179 var iter: i64 = 0 180 while iter < 64 { 181 let next_g: i64 = (g + x / g) / 2 182 if next_g >= g { iter = 64 } 183 if next_g < g { 184 g = next_g 185 iter = iter + 1 186 } 187 } 188 return g 189} 190 191// Confidence depends on d: (3/4)^(d/2). Approximate by tabulating. 192func nx_cs_conf_ppb(d: i64) -> i64 { 193 if d <= 3 { return 562500000 } // (3/4)^1.5 = 0.65 194 if d <= 5 { return 421900000 } // (3/4)^2.5 195 if d <= 7 { return 316400000 } // (3/4)^3.5 196 if d <= 11 { return 177900000 } // (3/4)^5.5 197 return 100000000 // higher d converges 198} 199 200func nx_cs_query(c: *CountSketch, key: i64) -> *ApproxI64 { 201 let est: i64 = nx_cs_estimate(c, key) 202 let bound: i64 = c.total / nx_cs_isqrt(c.w) 203 // Returned envelope: 1e9 - conf_ppb confidence the bound holds; 204 // i.e. there's a small chance the estimate is further than `bound`. 205 return nx_approx_new(est, NX_ENV_ABS, bound, 206 1000000000 - nx_cs_conf_ppb(c.d), 207 NX_MATURITY_REFERENCE_IMPL, 208 NX_ADV_HONEST) 209} 210 211// === merge ======================================================== 212 213func nx_cs_merge(a: *CountSketch, b: *CountSketch) -> *CountSketch { 214 if a.d != b.d { return 0 as *CountSketch } 215 if a.w != b.w { return 0 as *CountSketch } 216 if a.seed != b.seed { return 0 as *CountSketch } 217 let out: *CountSketch = nx_cs_alloc(a.d, a.w, a.seed) 218 let cells: i64 = a.d * a.w 219 var i: i64 = 0 220 while i < cells { 221 out.counters[i] = a.counters[i] + b.counters[i] 222 i = i + 1 223 } 224 out.total = a.total + b.total 225 return out 226} 227 228func nx_cs_memory_bytes(c: *CountSketch) -> i64 { 229 return 48 + c.d * c.w * 8 230}