sketch_varopt.nx source
↩ module page · 239 lines · 8489 B
1// sketch_varopt.nx -- Weighted reservoir sampling (Efraimidis-Spirakis 2006 / Cohen 2011).
2//
3// Given a stream of (item, weight) arrivals, maintain a reservoir
4// of size k such that the inclusion probability is proportional
5// to weight. Heavy items more likely to be retained than light.
6//
7// COMPLETES THE SAMPLING FAMILY:
8// - Reservoir (Vitter 1985): uniform sampling, equal weights.
9// - VarOpt (here): weighted sampling.
10//
11// ALGORITHM (A-ExpJ variant of A-Res):
12// For each (item, w): generate key = u^(1/w) where u ~ Uniform(0,1).
13// Keep top-k items by key (descending). Higher w concentrates the
14// key near 1; lower w near 0. Inclusion probability is provably
15// proportional to w under reasonable conditions.
16//
17// INTEGER-FIXED-POINT APPROXIMATION (because C anchor lacks f64,
18// but we want the sibling to be Wheeler-comparable against C):
19// - u_31 = LCG draw in (0, 2^31)
20// - approx_log2_u = bitlen(u_31) - 31 // in [-30, 0]
21// - approx_log_key = (approx_log2_u * 1_000_000) / w
22// - key = approx_log_key
23// Higher weight pulls key closer to 0; lower weight pulls more
24// negative. Sort descending by key; top-k by key are retained.
25//
26// PROPERTIES UNDER THE APPROXIMATION:
27// - Heavy items (w >> 1) get keys near 0 (large in descending sort).
28// - Light items (w = 1) get keys uniformly distributed in
29// [-30_000_000, 0], breaking ties uniformly.
30// - Quality degrades for very large weight variance. v2 will swap
31// in a finer-grained log approximation.
32//
33// LOSSLESS-LANGUAGE DISCIPLINE: per-item retain probability declared
34// in the typed envelope as conf_ppb proportional to weight_total /
35// weight_sum_sampled.
36
37import "syscalls.nx"
38import "sketch_types.nx"
39
40const NX_VOPT_K_MIN: i64 = 4
41const NX_VOPT_K_MAX: i64 = 100000
42
43const NX_VOPT_LCG_A: i64 = 1103515245
44const NX_VOPT_LCG_C: i64 = 12345
45const NX_VOPT_LCG_MOD: i64 = 0x7FFFFFFF
46
47struct VoptEntry {
48 item: i64,
49 weight: i64,
50 key: i64, // negative; higher (closer to 0) = more retain-worthy
51}
52
53struct VarOpt {
54 entries: *VoptEntry, // sorted descending by key
55 k: i64,
56 n_items: i64, // = min(seen, k)
57 total_seen: i64,
58 total_weight: i64, // cumulative weight observed
59 rng_state: i64,
60}
61
62// === construction =================================================
63
64func nx_varopt_alloc(k: i64, seed: i64) -> *VarOpt {
65 if k < NX_VOPT_K_MIN { return 0 as *VarOpt }
66 if k > NX_VOPT_K_MAX { return 0 as *VarOpt }
67 let raw: *u8 = sys_mmap(48)
68 let v: *VarOpt = raw as *VarOpt
69 let ent_raw: *u8 = sys_mmap(k * 24)
70 v.entries = ent_raw as *VoptEntry
71 v.k = k
72 v.n_items = 0
73 v.total_seen = 0
74 v.total_weight = 0
75 v.rng_state = seed | 1
76 return v
77}
78
79func nx_varopt_rng_next(v: *VarOpt) -> i64 {
80 let next: i64 = ((v.rng_state * NX_VOPT_LCG_A) + NX_VOPT_LCG_C) & NX_VOPT_LCG_MOD
81 v.rng_state = next
82 return next
83}
84
85// === bit-length of an i64 (highest set bit position + 1) ==========
86//
87// For x in [1, 2^31), returns value in [1, 31].
88
89func nx_varopt_bitlen(x: i64) -> i64 {
90 if x <= 0 { return 0 }
91 var n: i64 = 0
92 var t: i64 = x
93 while t > 0 {
94 t = t >> 1
95 n = n + 1
96 }
97 return n
98}
99
100// === key generation ===============================================
101//
102// approx_log2_u = bitlen(u_31) - 31 ∈ [-30, 0]
103// approx_log_key = (approx_log2_u * 1_000_000) / weight
104
105func nx_varopt_key(v: *VarOpt, weight: i64) -> i64 {
106 if weight <= 0 { return -2147483647 } // refuse zero/negative weight
107 let u: i64 = nx_varopt_rng_next(v)
108 if u <= 0 { return -2147483647 }
109 let lg: i64 = nx_varopt_bitlen(u) // 1..31
110 let approx_log2: i64 = lg - 31 // [-30, 0]
111 return (approx_log2 * 1000000) / weight
112}
113
114// === entry access =================================================
115
116func nx_varopt_entry_at(v: *VarOpt, i: i64) -> *VoptEntry {
117 return (v.entries as i64 + i * 24) as *VoptEntry
118}
119
120// === add ===========================================================
121//
122// New (item, weight) -> generate key -> insert into sorted-by-key
123// descending list (insertion sort). If list at cap, replace last
124// (smallest-key) entry if new key is larger.
125
126func nx_varopt_add(v: *VarOpt, item: i64, weight: i64) -> i64 {
127 if weight <= 0 { return -1 }
128 v.total_seen = v.total_seen + 1
129 v.total_weight = v.total_weight + weight
130 let key: i64 = nx_varopt_key(v, weight)
131
132 if v.n_items < v.k {
133 // Insert in sorted position (descending by key).
134 var pos: i64 = v.n_items
135 var done: i64 = 0
136 while done == 0 {
137 if pos == 0 { done = 1 }
138 if done == 0 {
139 let prev: *VoptEntry = nx_varopt_entry_at(v, pos - 1)
140 if prev.key >= key { done = 1 }
141 if done == 0 {
142 let dst: *VoptEntry = nx_varopt_entry_at(v, pos)
143 dst.item = prev.item
144 dst.weight = prev.weight
145 dst.key = prev.key
146 pos = pos - 1
147 }
148 }
149 }
150 let slot: *VoptEntry = nx_varopt_entry_at(v, pos)
151 slot.item = item
152 slot.weight = weight
153 slot.key = key
154 v.n_items = v.n_items + 1
155 return 0
156 }
157 // At cap: only replace if new key > current minimum (entries[k-1]).
158 let last: *VoptEntry = nx_varopt_entry_at(v, v.k - 1)
159 if key <= last.key { return 0 }
160 // Drop the last; insert new in sorted position.
161 var pos: i64 = v.k - 1
162 var done: i64 = 0
163 while done == 0 {
164 if pos == 0 { done = 1 }
165 if done == 0 {
166 let prev: *VoptEntry = nx_varopt_entry_at(v, pos - 1)
167 if prev.key >= key { done = 1 }
168 if done == 0 {
169 let dst: *VoptEntry = nx_varopt_entry_at(v, pos)
170 dst.item = prev.item
171 dst.weight = prev.weight
172 dst.key = prev.key
173 pos = pos - 1
174 }
175 }
176 }
177 let slot: *VoptEntry = nx_varopt_entry_at(v, pos)
178 slot.item = item
179 slot.weight = weight
180 slot.key = key
181 return 0
182}
183
184// === queries =======================================================
185//
186// Compute the weighted-mean of an attribute the caller wishes to
187// estimate. In this v1 sketch we expose the raw item array; callers
188// using i64 items can compute custom aggregates. Estimator:
189// E[f(stream)] ≈ (1/n_items) * Σ f(item_i) * (total_weight / weight_i)
190// where the (total_weight / weight_i) is the importance-sampling
191// inverse-probability weight.
192
193func nx_varopt_n_items(v: *VarOpt) -> i64 {
194 return v.n_items
195}
196
197func nx_varopt_total_seen(v: *VarOpt) -> i64 {
198 return v.total_seen
199}
200
201func nx_varopt_total_weight(v: *VarOpt) -> i64 {
202 return v.total_weight
203}
204
205// Weighted mean of items in the sample, with inverse-probability
206// reweighting. Returns sum(item_i * total_weight / weight_i) / n_items.
207// (For a constant-attribute estimator; users compute their own when
208// item is structured.)
209func nx_varopt_weighted_mean(v: *VarOpt) -> i64 {
210 if v.n_items == 0 { return 0 }
211 var sum: i64 = 0
212 var i: i64 = 0
213 while i < v.n_items {
214 let e: *VoptEntry = nx_varopt_entry_at(v, i)
215 sum = sum + (e.item * v.total_weight) / e.weight
216 i = i + 1
217 }
218 return sum / v.n_items
219}
220
221// === envelope ======================================================
222//
223// rel_stddev for the importance-weighted estimator. Cohen 2011
224// shows VarOpt achieves variance ~ N · Σ w_i^2 / k. We declare a
225// conservative 1/sqrt(k) bound that's accurate for low-variance
226// weight distributions; high-variance weights degrade quality.
227
228func nx_varopt_query_mean(v: *VarOpt) -> *ApproxI64 {
229 let mean: i64 = nx_varopt_weighted_mean(v)
230 let stddev_ppb: i64 = (1000000000 / nx_varopt_bitlen(v.k)) / 8 // rough 1/sqrt(k)
231 return nx_approx_new(mean, NX_ENV_REL_STDDEV, stddev_ppb,
232 682700000, // 1-sigma
233 NX_MATURITY_REFERENCE_IMPL,
234 NX_ADV_HONEST)
235}
236
237func nx_varopt_memory_bytes(v: *VarOpt) -> i64 {
238 return 48 + v.k * 24
239}