diff --git a/src/ops/agg_engine.c b/src/ops/agg_engine.c index e030b13c..8f414c61 100644 --- a/src/ops/agg_engine.c +++ b/src/ops/agg_engine.c @@ -3724,20 +3724,91 @@ agg_radix_select_topn(ray_pool_t* pool, const agg_radix_part_t* parts, uint32_t int64_t kept = 0; for (uint32_t p = 0; p < n_parts; p++) kept += c.kept[p]; ray_free_raw(cand); ray_free_raw(cand_n); - agg_radix_order_t* sel = ray_alloc_raw((size_t)(kept > 0 ? kept : 1) * sizeof(agg_radix_order_t)); - int64_t* fr = ray_alloc_raw((size_t)(kept > 0 ? kept : 1) * sizeof(int64_t)); + /* The kept set is every group at or beyond the threshold. Groups + * strictly beyond it number fewer than N by construction, but the groups + * AT the threshold can be nearly all of them — a count-per-group top-10 + * over a near-unique key has a threshold of 1 and every group ties it — + * and sorting that set is a comparison sort over the whole grouping. + * The emitted prefix is the tied groups' first-seen order, so only the + * N tied groups with the smallest first_row can ever be taken: keep + * exactly those, through a bounded max-heap, and the sort below runs + * over at most 2N entries. Without a threshold every group is kept + * (the take exceeds the group count) and the set is small anyway. */ + int64_t take = c.have_thr ? ef->top_count_take : 0; + int64_t sel_cap = c.have_thr ? 2 * take : kept; + if (sel_cap <= 0) sel_cap = 1; + agg_radix_order_t* sel = ray_alloc_raw((size_t)sel_cap * sizeof(agg_radix_order_t)); + int64_t* fr = ray_alloc_raw((size_t)sel_cap * sizeof(int64_t)); if (!sel || !fr) { ray_free_raw(sel); ray_free_raw(fr); ray_free_raw(base); ray_free_raw(vals); ray_free_raw(keep); *rc = 1; return NULL; } int64_t k = 0; - for (uint32_t p = 0; p < n_parts; p++) - for (int64_t gg = 0; gg < parts[p].ng; gg++) - if (keep[base[p] + gg]) { - sel[k].idx = ((int64_t)p << 32) | (uint32_t)gg; - fr[k] = parts[p].first_row[gg]; - k++; + if (!c.have_thr) { + for (uint32_t p = 0; p < n_parts; p++) + for (int64_t gg = 0; gg < parts[p].ng; gg++) + if (keep[base[p] + gg]) { + sel[k].idx = ((int64_t)p << 32) | (uint32_t)gg; + fr[k] = parts[p].first_row[gg]; + k++; + } + } else { + /* strict groups fill sel[0..k); the tied heap lives in sel[take..) */ + agg_radix_order_t* hp = sel + take; + int64_t* hk = fr + take; + int64_t hn = 0; + bool overflow = false; + for (uint32_t p = 0; p < n_parts && !overflow; p++) { + const double* pv = vals + base[p]; + const uint8_t* pk = keep + base[p]; + const int64_t* pfr = parts[p].first_row; + for (int64_t gg = 0; gg < parts[p].ng; gg++) { + if (!pk[gg]) continue; + int64_t payload = ((int64_t)p << 32) | (uint32_t)gg; + if (pv[gg] != c.thr) { + if (k >= take) { overflow = true; break; } + sel[k].idx = payload; fr[k] = pfr[gg]; k++; + continue; + } + int64_t first = pfr[gg]; + if (hn < take) { + int64_t i = hn++; + hk[i] = first; hp[i].idx = payload; + while (i > 0) { /* sift up */ + int64_t par = (i - 1) / 2; + if (hk[par] >= hk[i]) break; + int64_t tk = hk[par]; hk[par] = hk[i]; hk[i] = tk; + int64_t tp = hp[par].idx; hp[par].idx = hp[i].idx; hp[i].idx = tp; + i = par; + } + } else if (first < hk[0]) { + hk[0] = first; hp[0].idx = payload; + int64_t i = 0; + for (;;) { /* sift down */ + int64_t l = 2 * i + 1, r = l + 1, m = i; + if (l < take && hk[l] > hk[m]) m = l; + if (r < take && hk[r] > hk[m]) m = r; + if (m == i) break; + int64_t tk = hk[m]; hk[m] = hk[i]; hk[i] = tk; + int64_t tp = hp[m].idx; hp[m].idx = hp[i].idx; hp[i].idx = tp; + i = m; + } + } } + } + if (overflow) { + /* more than N groups beyond the threshold: not a threshold this + * selection understands — let the caller take the full path */ + ray_free_raw(sel); ray_free_raw(fr); + ray_free_raw(base); ray_free_raw(vals); ray_free_raw(keep); *rc = 2; return NULL; + } + /* close the gap between the strict prefix and the tied heap */ + if (k < take) { + memmove(sel + k, hp, (size_t)hn * sizeof(agg_radix_order_t)); + memmove(fr + k, hk, (size_t)hn * sizeof(int64_t)); + } + k += hn; + } agg_sort_pairs_by_key(sel, fr, k); ray_free_raw(fr); ray_free_raw(base); ray_free_raw(vals); ray_free_raw(keep); *n_emit = k; diff --git a/test/rfl/group/emit_filter_v2_route.rfl b/test/rfl/group/emit_filter_v2_route.rfl index 54447a99..41a1ba01 100644 --- a/test/rfl/group/emit_filter_v2_route.rfl +++ b/test/rfl/group/emit_filter_v2_route.rfl @@ -109,3 +109,28 @@ (set nested2 (select {from: (select {from: (select {from: tz by: k c: (count v)}) by: c n: (count c)}) where: (> n 100)})) (count nested2) -- 1 (at (at nested2 'n) 0) -- 1000 + +;; a threshold that every group ties: a near-unique two-key grouping where +;; almost every count is 1, so the N-th largest count is 1 and the kept set +;; would be the whole grouping. The result is still the first N tied groups +;; in first-seen order, and the selection must not scale with the group count. +(set n9 300000) +(set i9 (til n9)) +(set a9 (as 'I64 (+ (* i9 2654435761) (% (* i9 7) 3)))) +(set b9 (as 'I64 (% (* i9 40503) 150000))) +(set t9 (table [a b v] (list a9 b9 (% i9 5)))) +(set r9 (select {from: t9 by: [a b] c: (count v) s: (sum v) desc: c take: 10})) +(set f9 (select {from: t9 by: [a b] c: (count v) s: (sum v)})) +(count r9) -- 10 +(all (>= (at r9 'c) (at (at (xdesc f9 'c) 'c) 9))) -- true +(== (sum (at r9 'c)) (sum (take (at (xdesc f9 'c) 'c) 10))) -- true +(all (== (at r9 'a) (take (at (select {from: f9 where: (== c (max c))}) 'a) 10))) -- true +(all (== (at r9 's) (take (at (select {from: f9 where: (== c (max c))}) 's) 10))) -- true +;; the same shape under a row selection +(set r9w (select {from: t9 by: [a b] c: (count v) desc: c take: 10 where: (!= v 0)})) +(set f9w (select {from: t9 by: [a b] c: (count v) where: (!= v 0)})) +(count r9w) -- 10 +(all (== (at r9w 'a) (take (at (select {from: f9w where: (== c (max c))}) 'a) 10))) -- true +;; ascending: the smallest counts, again first-seen among the ties +(set r9a (select {from: t9 by: [a b] c: (count v) asc: c take: 10})) +(all (== (at r9a 'a) (take (at (select {from: f9 where: (== c (min c))}) 'a) 10))) -- true diff --git a/test/test_agg_contract.c b/test/test_agg_contract.c index c05e53ea..2134755c 100644 --- a/test/test_agg_contract.c +++ b/test/test_agg_contract.c @@ -1815,9 +1815,11 @@ static test_result_t test_radix_native_topn(void) { TEST_ASSERT_EQ_I(stats.routes[AGG_ROUTE_LEGACY], 0); TEST_ASSERT_EQ_I(stats.routes[AGG_ROUTE_V2_RADIX], 1); TEST_ASSERT_TRUE(stats.topn_native); - /* counts are 1 or 2: the top-10 superset is every count-2 group (500,000 - * of 1,500,000), not the full group set */ - TEST_ASSERT_EQ_I(stats.topn_kept, 500000); + /* counts are 1 or 2, so the threshold is 2 and every count-2 group + * (500,000 of 1,500,000) ties it: the selection keeps the ten of them + * with the smallest first row — the ones the emitted first-seen order + * would take — not the whole tie set */ + TEST_ASSERT_EQ_I(stats.topn_kept, 10); TEST_ASSERT_EQ_I(ray_table_nrows(r), 10); ray_t* check = ray_eval_str( "(== (at (at (select {from:rt_t by:[k j] c:(count v) desc:c take:10}) 'c) 0) "