#include "ops/internal.h"
#include "ops/hash.h"
#include "ops/rowsel.h"
#include "ops/hll.h"
#include "lang/internal.h"
#include "lang/format.h"
#include "table/domain.h"
#include "ops/agg_engine.h"
#include "vec/str.h"
#include "ops/idxop.h"
#include "core/runtime.h"
#include "core/profile.h"
#include <stdlib.h>
static inline bool group_fp_type(int8_t t) {
return t == RAY_F32 || t == RAY_F64;
}
static inline bool group_avg_needs_hi(int8_t t, int64_t nrows) {
return t == RAY_I64 || t == RAY_TIMESTAMP || t == RAY_SYM || t == 0 ||
nrows >= ((int64_t)1 << 31);
}
static inline int64_t group_key_f64_bits(const void* data, int64_t row) {
int64_t bits;
memcpy(&bits, (const double*)data + row, 8);
if (bits == INT64_MIN) bits = 0;
return bits;
}
static inline double group_fp_at(const void* data, int8_t t, int64_t row) {
return t == RAY_F64 ? ((const double*)data)[row]
: (double)((const float*)data)[row];
}
static inline double group_pack_i64_as_f64(ray_t* col, int64_t row) {
int64_t v = read_col_i64(ray_data(col), row, col->type, col->attrs);
ray_t* src = (col->attrs & RAY_ATTR_SLICE) ? col->slice_parent : col;
if (src && ray_vec_may_have_nulls(src)) {
int64_t sent = agg_int_null_sentinel_for(col->type);
if (sent != 0 && v == sent) return (double)NAN;
}
return (double)v;
}
static inline int64_t group_pack_y_i64(ray_t* col, int64_t row) {
int64_t v = read_col_i64(ray_data(col), row, col->type, col->attrs);
ray_t* src = (col->attrs & RAY_ATTR_SLICE) ? col->slice_parent : col;
if (src && ray_vec_may_have_nulls(src)) {
int64_t sent = agg_int_null_sentinel_for(col->type);
if (sent != 0 && v == sent) return NULL_I64;
}
return v;
}
static _Atomic(int64_t) ray_group_perpart_runs_ctr = 0;
int64_t ray_group_perpart_runs(void) {
return atomic_load_explicit(&ray_group_perpart_runs_ctr,
memory_order_relaxed);
}
typedef struct {
double sum_f, min_f, max_f, prod_f, first_f, last_f, sum_sq_f;
int64_t sum_i, min_i, max_i, prod_i, first_i, last_i, sum_sq_i;
double sum_d;
int64_t sum_hi;
int64_t cnt;
int64_t zero_count;
bool has_first;
} reduce_acc_t;
static void reduce_acc_init(reduce_acc_t* acc) {
acc->sum_f = 0; acc->min_f = DBL_MAX; acc->max_f = -DBL_MAX;
acc->prod_f = 1.0; acc->first_f = 0; acc->last_f = 0; acc->sum_sq_f = 0;
acc->sum_i = 0; acc->min_i = INT64_MAX; acc->max_i = INT64_MIN;
acc->prod_i = 1; acc->first_i = 0; acc->last_i = 0; acc->sum_sq_i = 0;
acc->sum_d = 0; acc->sum_hi = 0;
acc->cnt = 0; acc->zero_count = 0; acc->has_first = false;
}
static inline bool sym_lex_lt(struct ray_sym_domain_s* dom,
int64_t a, int64_t b) {
if (a == b) return false;
const char* pa; const char* pb;
size_t la, lb;
ray_sym_domain_raw_t raw;
if (ray_sym_domain_raw_pin(dom, &raw) && a >= 0 && b >= 0 &&
a < raw.count && b < raw.count) {
pa = ray_sym_domain_raw_str(&raw, a, &la);
pb = ray_sym_domain_raw_str(&raw, b, &lb);
} else {
ray_t* sa = ray_sym_domain_str(dom, a);
ray_t* sb = ray_sym_domain_str(dom, b);
if (!sa || !sb) return a < b;
pa = ray_str_ptr(sa); pb = ray_str_ptr(sb);
la = ray_str_len(sa); lb = ray_str_len(sb);
}
size_t m = la < lb ? la : lb;
int c = memcmp(pa, pb, m);
if (c != 0) return c < 0;
return la < lb;
}
static inline bool sym_lex_gt(struct ray_sym_domain_s* dom,
int64_t a, int64_t b) {
return sym_lex_lt(dom, b, a);
}
static inline bool agg_is_wide_type(int8_t t) {
return t == RAY_STR || t == RAY_GUID;
}
static inline bool agg_needs_row_gather(int8_t t, uint16_t op) {
if (t == RAY_LIST)
return op == OP_FIRST || op == OP_LAST;
return agg_is_wide_type(t) &&
(op == OP_MIN || op == OP_MAX || op == OP_FIRST || op == OP_LAST);
}
static inline void out_col_adopt_str_pool(ray_t* dst, const ray_t* src) {
if (!dst || RAY_IS_ERR(dst) || dst->type != RAY_STR || !src) return;
const ray_t* owner = (src->attrs & RAY_ATTR_SLICE) ? src->slice_parent : src;
if (owner && owner->str_pool && !RAY_IS_ERR(owner->str_pool)) {
ray_retain(owner->str_pool);
dst->str_pool = owner->str_pool;
}
}
static int64_t wide_winner_row(ray_t* input, uint16_t op,
const int64_t* sel, int64_t scan_n,
bool has_nulls) {
if (op == OP_FIRST) {
for (int64_t i = 0; i < scan_n; i++) {
if ((i & 65535) == 0 && ray_interrupted()) return -1;
int64_t row = sel ? sel[i] : i;
if (input->type == RAY_LIST) {
ray_t* e = ((ray_t**)ray_data(input))[row];
if (e && !RAY_ATOM_IS_NULL(e)) return row;
} else if (!has_nulls || !ray_vec_is_null(input, row)) {
return row;
}
}
return -1;
}
if (op == OP_LAST) {
for (int64_t i = scan_n - 1; i >= 0; i--) {
if ((i & 65535) == 0 && ray_interrupted()) return -1;
int64_t row = sel ? sel[i] : i;
if (input->type == RAY_LIST) {
ray_t* e = ((ray_t**)ray_data(input))[row];
if (e && !RAY_ATOM_IS_NULL(e)) return row;
} else if (!has_nulls || !ray_vec_is_null(input, row)) {
return row;
}
}
return -1;
}
const bool want_min = (op == OP_MIN);
int64_t best = -1;
if (input->type == RAY_GUID) {
const uint8_t* d = (const uint8_t*)ray_data(input);
for (int64_t i = 0; i < scan_n; i++) {
if ((i & 65535) == 0 && ray_interrupted()) return -1;
int64_t row = sel ? sel[i] : i;
if (has_nulls && ray_vec_is_null(input, row)) continue;
if (best < 0) { best = row; continue; }
int c = memcmp(d + (size_t)row * 16, d + (size_t)best * 16, 16);
if (want_min ? (c < 0) : (c > 0)) best = row;
}
} else {
for (int64_t i = 0; i < scan_n; i++) {
if ((i & 65535) == 0 && ray_interrupted()) return -1;
int64_t row = sel ? sel[i] : i;
if (has_nulls && ray_vec_is_null(input, row)) continue;
if (best < 0) { best = row; continue; }
size_t la = 0, lb = 0;
const char* sa = ray_str_vec_get(input, row, &la);
const char* sb = ray_str_vec_get(input, best, &lb);
size_t m = la < lb ? la : lb;
int c = (m && sa && sb) ? memcmp(sa, sb, m) : 0;
if (!c) c = (la > lb) - (la < lb);
if (want_min ? (c < 0) : (c > 0)) best = row;
}
}
return best;
}
static ray_t* agg_wide_reduce(ray_t* input, uint16_t op,
const int64_t* sel, int64_t scan_n,
bool has_nulls) {
int64_t best = wide_winner_row(input, op, sel, scan_n, has_nulls);
if (best < 0)
return input->type == RAY_LIST
? ray_typed_null(-RAY_I64)
: ray_typed_null(-input->type);
int alloc;
ray_t* result = collection_elem(input, best, &alloc);
if (result && !alloc) ray_retain(result);
return result;
}
#define RED_NEED_SUM (1u << 0)
#define RED_NEED_SUM_SQ (1u << 1)
#define RED_NEED_PROD (1u << 2)
#define RED_NEED_SUM_D (1u << 3)
#define RED_NEED_COUNT (1u << 4)
#define RED_NEED_ZERO (1u << 5)
#define RED_NEED_MIN (1u << 6)
#define RED_NEED_MAX (1u << 7)
#define RED_NEED_FIRST (1u << 8)
#define RED_NEED_LAST (1u << 9)
#define RED_NEED_SUM_128 (1u << 10)
#define RED_MASK_ANY (RED_NEED_COUNT | RED_NEED_ZERO)
#define RED_MASK_MIN (RED_NEED_COUNT | RED_NEED_MIN)
#define RED_MASK_MAX (RED_NEED_COUNT | RED_NEED_MAX)
#define RED_MASK_FIRST (RED_NEED_COUNT | RED_NEED_FIRST)
#define RED_MASK_LAST (RED_NEED_COUNT | RED_NEED_LAST)
#define RED_MASK_AVG_I (RED_NEED_SUM_128 | RED_NEED_COUNT)
#define RED_MASK_AVG_F (RED_NEED_SUM | RED_NEED_COUNT)
#define RED_MASK_STATS_I (RED_NEED_SUM_D | RED_NEED_SUM_SQ | RED_NEED_COUNT)
#define RED_MASK_STATS_F (RED_NEED_SUM | RED_NEED_SUM_SQ | RED_NEED_COUNT)
#define REDUCE_LOOP_I(T, NULL_SENT, base, start, end, acc, HAS_NULLS, HAS_IDX, idx, NEEDS) \
do { \
const T* d = (const T*)(base); \
for (int64_t i = start; i < end; i++) { \
int64_t row = (HAS_IDX) ? (idx)[i] : i; \
T raw = d[row]; \
if ((HAS_NULLS) && raw == (T)(NULL_SENT)) continue; \
int64_t v = (int64_t)raw; \
\
if ((NEEDS) & RED_NEED_SUM) \
(acc)->sum_i = (int64_t)((uint64_t)(acc)->sum_i + (uint64_t)v); \
if ((NEEDS) & RED_NEED_SUM_SQ) \
(acc)->sum_sq_i = (int64_t)((uint64_t)(acc)->sum_sq_i + (uint64_t)v * (uint64_t)v); \
if ((NEEDS) & RED_NEED_PROD) \
(acc)->prod_i = (int64_t)((uint64_t)(acc)->prod_i * (uint64_t)v); \
if ((NEEDS) & RED_NEED_SUM_D) (acc)->sum_d += (double)v; \
if ((NEEDS) & RED_NEED_SUM_128) { \
uint64_t lo_ = (uint64_t)(acc)->sum_i + (uint64_t)v; \
(acc)->sum_hi += (v < 0 ? -1 : 0) + (lo_ < (uint64_t)v ? 1 : 0); \
(acc)->sum_i = (int64_t)lo_; \
} \
if (((NEEDS) & RED_NEED_ZERO) && v == 0) (acc)->zero_count++; \
if (((NEEDS) & RED_NEED_MIN) && v < (acc)->min_i) (acc)->min_i = v; \
if (((NEEDS) & RED_NEED_MAX) && v > (acc)->max_i) (acc)->max_i = v; \
if (((NEEDS) & RED_NEED_FIRST) && !(acc)->has_first) { \
(acc)->first_i = v; (acc)->has_first = true; \
} \
if ((NEEDS) & RED_NEED_LAST) (acc)->last_i = v; \
if ((NEEDS) & RED_NEED_COUNT) (acc)->cnt++; \
} \
} while (0)
#define REDUCE_LOOP_BITMAP_I(T, input, base, start, end, acc, HAS_NULLS, HAS_IDX, idx, NEEDS) \
do { \
const T* d = (const T*)(base); \
for (int64_t i = start; i < end; i++) { \
int64_t row = (HAS_IDX) ? (idx)[i] : i; \
if ((HAS_NULLS) && ray_vec_is_null((input), row)) continue; \
int64_t v = (int64_t)d[row]; \
if ((NEEDS) & RED_NEED_SUM) \
(acc)->sum_i = (int64_t)((uint64_t)(acc)->sum_i + (uint64_t)v); \
if ((NEEDS) & RED_NEED_SUM_SQ) \
(acc)->sum_sq_i = (int64_t)((uint64_t)(acc)->sum_sq_i + (uint64_t)v * (uint64_t)v); \
if ((NEEDS) & RED_NEED_PROD) \
(acc)->prod_i = (int64_t)((uint64_t)(acc)->prod_i * (uint64_t)v); \
if ((NEEDS) & RED_NEED_SUM_D) (acc)->sum_d += (double)v; \
if ((NEEDS) & RED_NEED_SUM_128) { \
uint64_t lo_ = (uint64_t)(acc)->sum_i + (uint64_t)v; \
(acc)->sum_hi += (v < 0 ? -1 : 0) + (lo_ < (uint64_t)v ? 1 : 0); \
(acc)->sum_i = (int64_t)lo_; \
} \
if (((NEEDS) & RED_NEED_ZERO) && v == 0) (acc)->zero_count++; \
if (((NEEDS) & RED_NEED_MIN) && v < (acc)->min_i) (acc)->min_i = v; \
if (((NEEDS) & RED_NEED_MAX) && v > (acc)->max_i) (acc)->max_i = v; \
if (((NEEDS) & RED_NEED_FIRST) && !(acc)->has_first) { \
(acc)->first_i = v; (acc)->has_first = true; \
} \
if ((NEEDS) & RED_NEED_LAST) (acc)->last_i = v; \
if ((NEEDS) & RED_NEED_COUNT) (acc)->cnt++; \
} \
} while (0)
#define REDUCE_LOOP_F(T, base, start, end, acc, HAS_NULLS, HAS_IDX, idx, NEEDS) \
do { \
const T* d = (const T*)(base); \
for (int64_t i = start; i < end; i++) { \
int64_t row = (HAS_IDX) ? (idx)[i] : i; \
double v = (double)d[row]; \
if ((HAS_NULLS) && v != v) continue; \
if ((NEEDS) & RED_NEED_SUM) (acc)->sum_f += v; \
if ((NEEDS) & RED_NEED_SUM_SQ) (acc)->sum_sq_f += v * v; \
if ((NEEDS) & RED_NEED_PROD) (acc)->prod_f *= v; \
if (((NEEDS) & RED_NEED_ZERO) && v == 0.0) (acc)->zero_count++; \
if (((NEEDS) & RED_NEED_MIN) && v < (acc)->min_f) (acc)->min_f = v; \
if (((NEEDS) & RED_NEED_MAX) && v > (acc)->max_f) (acc)->max_f = v; \
if (((NEEDS) & RED_NEED_FIRST) && !(acc)->has_first) { \
(acc)->first_f = v; (acc)->has_first = true; \
} \
if ((NEEDS) & RED_NEED_LAST) (acc)->last_f = v; \
if ((NEEDS) & RED_NEED_COUNT) (acc)->cnt++; \
} \
} while (0)
#define DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, NEEDS) \
do { \
if (!(has_nulls) && !(idx)) \
REDUCE_LOOP_I(T, NULL_SENT, base, start, end, acc, 0, 0, idx, NEEDS); \
else if (!(has_nulls)) \
REDUCE_LOOP_I(T, NULL_SENT, base, start, end, acc, 0, 1, idx, NEEDS); \
else if (!(idx)) \
REDUCE_LOOP_I(T, NULL_SENT, base, start, end, acc, 1, 0, idx, NEEDS); \
else \
REDUCE_LOOP_I(T, NULL_SENT, base, start, end, acc, 1, 1, idx, NEEDS); \
} while (0)
#define DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, NEEDS) \
do { \
if (!(has_nulls) && !(idx)) \
REDUCE_LOOP_BITMAP_I(T, input, base, start, end, acc, 0, 0, idx, NEEDS); \
else if (!(has_nulls)) \
REDUCE_LOOP_BITMAP_I(T, input, base, start, end, acc, 0, 1, idx, NEEDS); \
else if (!(idx)) \
REDUCE_LOOP_BITMAP_I(T, input, base, start, end, acc, 1, 0, idx, NEEDS); \
else \
REDUCE_LOOP_BITMAP_I(T, input, base, start, end, acc, 1, 1, idx, NEEDS); \
} while (0)
#define DISPATCH_F(T, base, start, end, acc, has_nulls, idx, NEEDS) \
do { \
if (!(has_nulls) && !(idx)) \
REDUCE_LOOP_F(T, base, start, end, acc, 0, 0, idx, NEEDS); \
else if (!(has_nulls)) \
REDUCE_LOOP_F(T, base, start, end, acc, 0, 1, idx, NEEDS); \
else if (!(idx)) \
REDUCE_LOOP_F(T, base, start, end, acc, 1, 0, idx, NEEDS); \
else \
REDUCE_LOOP_F(T, base, start, end, acc, 1, 1, idx, NEEDS); \
} while (0)
#define DISPATCH_I_OPCODE(T, NULL_SENT, base, start, end, acc, has_nulls, idx, opcode) \
do { \
switch (opcode) { \
case OP_SUM: DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_NEED_SUM); break; \
case OP_PROD: DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_NEED_PROD); break; \
case OP_ALL: DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_NEED_ZERO); break; \
case OP_ANY: DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_MASK_ANY); break; \
case OP_MIN: DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_MASK_MIN); break; \
case OP_MAX: DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_MASK_MAX); break; \
case OP_AVG: DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_MASK_AVG_I); break; \
case OP_VAR: case OP_VAR_POP: case OP_STDDEV: case OP_STDDEV_POP: \
DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_MASK_STATS_I); break; \
case OP_FIRST: DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_MASK_FIRST); break; \
case OP_LAST: DISPATCH_I(T, NULL_SENT, base, start, end, acc, has_nulls, idx, RED_MASK_LAST); break; \
default: break; \
} \
} while (0)
#define DISPATCH_BITMAP_I_OPCODE(T, input, base, start, end, acc, has_nulls, idx, opcode) \
do { \
switch (opcode) { \
case OP_SUM: DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_NEED_SUM); break; \
case OP_PROD: DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_NEED_PROD); break; \
case OP_ALL: DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_NEED_ZERO); break; \
case OP_ANY: DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_MASK_ANY); break; \
case OP_MIN: DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_MASK_MIN); break; \
case OP_MAX: DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_MASK_MAX); break; \
case OP_AVG: DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_MASK_AVG_I); break; \
case OP_VAR: case OP_VAR_POP: case OP_STDDEV: case OP_STDDEV_POP: \
DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_MASK_STATS_I); break; \
case OP_FIRST: DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_MASK_FIRST); break; \
case OP_LAST: DISPATCH_BITMAP_I(T, input, base, start, end, acc, has_nulls, idx, RED_MASK_LAST); break; \
default: break; \
} \
} while (0)
#define DISPATCH_F_OPCODE(T, base, start, end, acc, has_nulls, idx, opcode) \
do { \
switch (opcode) { \
case OP_SUM: DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_NEED_SUM); break; \
case OP_PROD: DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_NEED_PROD); break; \
case OP_ALL: DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_NEED_ZERO); break; \
case OP_ANY: DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_MASK_ANY); break; \
case OP_MIN: DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_MASK_MIN); break; \
case OP_MAX: DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_MASK_MAX); break; \
case OP_AVG: DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_MASK_AVG_F); break; \
case OP_VAR: case OP_VAR_POP: case OP_STDDEV: case OP_STDDEV_POP: \
DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_MASK_STATS_F); break; \
case OP_FIRST: DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_MASK_FIRST); break; \
case OP_LAST: DISPATCH_F(T, base, start, end, acc, has_nulls, idx, RED_MASK_LAST); break; \
default: break; \
} \
} while (0)
#define REDUCE_SYM_EXTREME(input, base, start, end, acc, idx, dom, HAS_NULLS, HAS_IDX, WANT_MIN) \
do { \
for (int64_t i = start; i < end; i++) { \
int64_t row = (HAS_IDX) ? (idx)[i] : i; \
int64_t v = read_col_i64(base, row, (input)->type, (input)->attrs); \
if ((HAS_NULLS) && v == 0) continue; \
if ((acc)->cnt == 0) { \
if (WANT_MIN) (acc)->min_i = v; else (acc)->max_i = v; \
} else if (WANT_MIN) { \
if (sym_lex_lt(dom, v, (acc)->min_i)) (acc)->min_i = v; \
} else { \
if (sym_lex_gt(dom, v, (acc)->max_i)) (acc)->max_i = v; \
} \
(acc)->cnt++; \
} \
} while (0)
#define DISPATCH_SYM_EXTREME(input, base, start, end, acc, has_nulls, idx, dom, WANT_MIN) \
do { \
if (!(has_nulls) && !(idx)) \
REDUCE_SYM_EXTREME(input, base, start, end, acc, idx, dom, 0, 0, WANT_MIN); \
else if (!(has_nulls)) \
REDUCE_SYM_EXTREME(input, base, start, end, acc, idx, dom, 0, 1, WANT_MIN); \
else if (!(idx)) \
REDUCE_SYM_EXTREME(input, base, start, end, acc, idx, dom, 1, 0, WANT_MIN); \
else \
REDUCE_SYM_EXTREME(input, base, start, end, acc, idx, dom, 1, 1, WANT_MIN); \
} while (0)
static void reduce_range(ray_t* input, int64_t start, int64_t end,
reduce_acc_t* acc, bool has_nulls,
const int64_t* idx, uint16_t opcode) {
void* base = ray_data(input);
switch (input->type) {
case RAY_BOOL: case RAY_U8:
DISPATCH_BITMAP_I_OPCODE(uint8_t, input, base, start, end, acc,
has_nulls, idx, opcode); break;
case RAY_I16:
DISPATCH_I_OPCODE(int16_t, NULL_I16, base, start, end, acc,
has_nulls, idx, opcode); break;
case RAY_I32: case RAY_DATE: case RAY_TIME:
DISPATCH_I_OPCODE(int32_t, NULL_I32, base, start, end, acc,
has_nulls, idx, opcode); break;
case RAY_I64: case RAY_TIMESTAMP:
DISPATCH_I_OPCODE(int64_t, NULL_I64, base, start, end, acc,
has_nulls, idx, opcode); break;
case RAY_F32:
DISPATCH_F_OPCODE(float, base, start, end, acc, has_nulls, idx, opcode); break;
case RAY_F64:
DISPATCH_F_OPCODE(double, base, start, end, acc, has_nulls, idx, opcode); break;
case RAY_SYM: {
struct ray_sym_domain_s* dom = ray_sym_vec_domain(input);
if (opcode == OP_MIN)
DISPATCH_SYM_EXTREME(input, base, start, end, acc, has_nulls,
idx, dom, 1);
else if (opcode == OP_MAX)
DISPATCH_SYM_EXTREME(input, base, start, end, acc, has_nulls,
idx, dom, 0);
break;
}
default: break;
}
}
typedef struct {
ray_t* input;
reduce_acc_t* accs;
bool has_nulls;
const int64_t* idx;
uint16_t opcode;
} par_reduce_ctx_t;
static void par_reduce_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
par_reduce_ctx_t* c = (par_reduce_ctx_t*)ctx;
reduce_range(c->input, start, end, &c->accs[worker_id],
c->has_nulls, c->idx, c->opcode);
}
static void reduce_merge(reduce_acc_t* dst, const reduce_acc_t* src, int8_t in_type,
struct ray_sym_domain_s* sym_dom, uint16_t opcode) {
bool fp = in_type == RAY_F64 || in_type == RAY_F32;
switch (opcode) {
case OP_SUM:
if (fp) dst->sum_f += src->sum_f;
else dst->sum_i = (int64_t)((uint64_t)dst->sum_i + (uint64_t)src->sum_i);
break;
case OP_PROD:
if (fp) dst->prod_f *= src->prod_f;
else dst->prod_i = (int64_t)((uint64_t)dst->prod_i * (uint64_t)src->prod_i);
break;
case OP_ALL:
dst->zero_count += src->zero_count;
break;
case OP_ANY:
dst->zero_count += src->zero_count;
dst->cnt += src->cnt;
break;
case OP_MIN:
if (src->cnt > 0) {
if (dst->cnt == 0) {
if (fp) dst->min_f = src->min_f;
else dst->min_i = src->min_i;
} else if (fp) {
if (src->min_f < dst->min_f) dst->min_f = src->min_f;
} else if (in_type == RAY_SYM) {
if (sym_lex_lt(sym_dom, src->min_i, dst->min_i)) dst->min_i = src->min_i;
} else if (src->min_i < dst->min_i) {
dst->min_i = src->min_i;
}
dst->cnt += src->cnt;
}
break;
case OP_MAX:
if (src->cnt > 0) {
if (dst->cnt == 0) {
if (fp) dst->max_f = src->max_f;
else dst->max_i = src->max_i;
} else if (fp) {
if (src->max_f > dst->max_f) dst->max_f = src->max_f;
} else if (in_type == RAY_SYM) {
if (sym_lex_gt(sym_dom, src->max_i, dst->max_i)) dst->max_i = src->max_i;
} else if (src->max_i > dst->max_i) {
dst->max_i = src->max_i;
}
dst->cnt += src->cnt;
}
break;
case OP_AVG:
if (fp) dst->sum_f += src->sum_f;
else {
uint64_t lo = (uint64_t)dst->sum_i + (uint64_t)src->sum_i;
dst->sum_hi += src->sum_hi + (lo < (uint64_t)src->sum_i ? 1 : 0);
dst->sum_i = (int64_t)lo;
}
dst->cnt += src->cnt;
break;
case OP_VAR: case OP_VAR_POP: case OP_STDDEV: case OP_STDDEV_POP:
if (fp) {
dst->sum_f += src->sum_f;
dst->sum_sq_f += src->sum_sq_f;
} else {
dst->sum_d += src->sum_d;
dst->sum_sq_i = (int64_t)((uint64_t)dst->sum_sq_i +
(uint64_t)src->sum_sq_i);
}
dst->cnt += src->cnt;
break;
case OP_FIRST: case OP_LAST:
dst->cnt += src->cnt;
break;
default:
break;
}
}
#define CD_HASH_K1 0x9E3779B97F4A7C15ULL
#define CD_HASH_K2 0xBF58476D1CE4E5B9ULL
typedef struct {
int64_t* values;
int64_t* part_off;
int64_t* part_count;
} cd_part_ctx_t;
static void cd_part_dedup_fn(void* ctx, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
cd_part_ctx_t* x = (cd_part_ctx_t*)ctx;
for (int64_t p = start; p < end; p++) {
int64_t off = x->part_off[p];
int64_t cnt = x->part_off[p + 1] - off;
if (cnt == 0) { x->part_count[p] = 0; continue; }
uint64_t cap = (uint64_t)cnt * 2;
if (cap < 32) cap = 32;
uint64_t c = 1;
while (c && c < cap) c <<= 1;
if (!c) { x->part_count[p] = -1; continue; }
cap = c;
uint64_t mask = cap - 1;
ray_t* set_hdr = NULL;
ray_t* used_hdr = NULL;
int64_t* set = (int64_t*)scratch_alloc (&set_hdr,
(size_t)cap * sizeof(int64_t));
uint8_t* used = (uint8_t*)scratch_calloc(&used_hdr,
(size_t)cap * sizeof(uint8_t));
if (!set || !used) {
if (set_hdr) scratch_free(set_hdr);
if (used_hdr) scratch_free(used_hdr);
x->part_count[p] = -1;
continue;
}
int64_t* base = x->values + off;
int64_t distinct = 0;
for (int64_t i = 0; i < cnt; i++) {
int64_t v = base[i];
uint64_t h = (uint64_t)v * CD_HASH_K1;
h ^= h >> 33;
uint64_t slot = h & mask;
while (used[slot]) {
if (set[slot] == v) goto cd_next;
slot = (slot + 1) & mask;
}
set[slot] = v;
used[slot] = 1;
distinct++;
cd_next:;
}
scratch_free(set_hdr);
scratch_free(used_hdr);
x->part_count[p] = distinct;
}
}
typedef struct {
const void* base;
int64_t* counts;
uint32_t p_bits;
uint64_t p_mask;
int64_t grain;
int64_t total;
uint8_t stride_log2;
uint8_t is_f64;
int8_t type;
uint8_t attrs;
} cd_count_ctx_t;
static void cd_hist_fn(void* ctx, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
(void)end;
cd_count_ctx_t* x = (cd_count_ctx_t*)ctx;
int64_t task_idx = start;
int64_t row_start = task_idx * x->grain;
int64_t row_end = row_start + x->grain;
if (row_end > x->total) row_end = x->total;
int64_t* hist = x->counts + (size_t)task_idx * (x->p_mask + 1);
start = row_start;
end = row_end;
const void* base = x->base;
int8_t in_type = x->type;
uint8_t in_attrs = x->attrs;
uint64_t p_mask = x->p_mask;
if (x->is_f64) {
const double* d = (const double*)base;
for (int64_t i = start; i < end; i++) {
double fv = d[i];
if (fv != fv) fv = (double)NAN;
else fv = clear_neg_zero(fv);
int64_t val;
memcpy(&val, &fv, sizeof(int64_t));
uint64_t h = (uint64_t)val * CD_HASH_K1;
h ^= h >> 33;
uint64_t p = (h ^ (h >> 33)) & p_mask;
hist[p]++;
}
} else if (in_type == RAY_I64 || in_type == RAY_TIMESTAMP) {
const int64_t* d = (const int64_t*)base;
for (int64_t i = start; i < end; i++) {
int64_t val = d[i];
uint64_t h = (uint64_t)val * CD_HASH_K1;
h ^= h >> 33;
uint64_t p = (h ^ (h >> 33)) & p_mask;
hist[p]++;
}
} else if (in_type == RAY_I32 || in_type == RAY_DATE || in_type == RAY_TIME) {
const int32_t* d = (const int32_t*)base;
for (int64_t i = start; i < end; i++) {
int64_t val = d[i];
uint64_t h = (uint64_t)val * CD_HASH_K1;
h ^= h >> 33;
uint64_t p = (h ^ (h >> 33)) & p_mask;
hist[p]++;
}
} else if (in_type == RAY_I16) {
const int16_t* d = (const int16_t*)base;
for (int64_t i = start; i < end; i++) {
int64_t val = d[i];
uint64_t h = (uint64_t)val * CD_HASH_K1;
h ^= h >> 33;
uint64_t p = (h ^ (h >> 33)) & p_mask;
hist[p]++;
}
} else if (in_type == RAY_BOOL || in_type == RAY_U8) {
const uint8_t* d = (const uint8_t*)base;
for (int64_t i = start; i < end; i++) {
int64_t val = d[i];
uint64_t h = (uint64_t)val * CD_HASH_K1;
h ^= h >> 33;
uint64_t p = (h ^ (h >> 33)) & p_mask;
hist[p]++;
}
} else if (in_type == RAY_SYM) {
for (int64_t i = start; i < end; i++) {
int64_t val = read_col_i64(base, i, in_type, in_attrs);
uint64_t h = (uint64_t)val * CD_HASH_K1;
h ^= h >> 33;
uint64_t p = (h ^ (h >> 33)) & p_mask;
hist[p]++;
}
}
}
typedef struct {
const void* base;
int64_t* out_buf;
int64_t* cursor;
uint32_t p_bits;
uint64_t p_mask;
int64_t grain;
int64_t total;
uint8_t is_f64;
int8_t type;
uint8_t attrs;
} cd_scatter_ctx_t;
static void cd_scatter_fn(void* ctx, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
(void)end;
cd_scatter_ctx_t* x = (cd_scatter_ctx_t*)ctx;
int64_t task_idx = start;
int64_t row_start = task_idx * x->grain;
int64_t row_end = row_start + x->grain;
if (row_end > x->total) row_end = x->total;
int64_t* cur = x->cursor + (size_t)task_idx * (x->p_mask + 1);
start = row_start;
end = row_end;
int64_t* out = x->out_buf;
const void* base = x->base;
int8_t in_type = x->type;
uint8_t in_attrs = x->attrs;
uint64_t p_mask = x->p_mask;
#define SCATTER_BODY(LOAD) \
for (int64_t i = start; i < end; i++) { \
int64_t val = (LOAD); \
uint64_t h = (uint64_t)val * CD_HASH_K1; \
h ^= h >> 33; \
uint64_t p = (h ^ (h >> 33)) & p_mask; \
out[cur[p]++] = val; \
}
if (x->is_f64) {
const double* d = (const double*)base;
for (int64_t i = start; i < end; i++) {
double fv = d[i];
if (fv != fv) fv = (double)NAN;
else fv = clear_neg_zero(fv);
int64_t val;
memcpy(&val, &fv, sizeof(int64_t));
uint64_t h = (uint64_t)val * CD_HASH_K1;
h ^= h >> 33;
uint64_t p = (h ^ (h >> 33)) & p_mask;
out[cur[p]++] = val;
}
} else if (in_type == RAY_I64 || in_type == RAY_TIMESTAMP) {
const int64_t* d = (const int64_t*)base;
SCATTER_BODY(d[i])
} else if (in_type == RAY_I32 || in_type == RAY_DATE || in_type == RAY_TIME) {
const int32_t* d = (const int32_t*)base;
SCATTER_BODY(d[i])
} else if (in_type == RAY_I16) {
const int16_t* d = (const int16_t*)base;
SCATTER_BODY(d[i])
} else if (in_type == RAY_BOOL || in_type == RAY_U8) {
const uint8_t* d = (const uint8_t*)base;
SCATTER_BODY(d[i])
} else {
SCATTER_BODY(read_col_i64(base, i, in_type, in_attrs))
}
#undef SCATTER_BODY
}
static int64_t cd_seq_count(int8_t in_type, uint8_t in_attrs,
const void* base, int64_t len) {
uint64_t cap = (uint64_t)(len < 16 ? 32 : len) * 2;
uint64_t c = 1;
while (c && c < cap) c <<= 1;
if (!c) return -1;
cap = c;
uint64_t mask = cap - 1;
int cd_shift = 64 - __builtin_ctzll(cap);
ray_t* set_hdr = NULL;
ray_t* used_hdr = NULL;
int64_t* set = (int64_t*)scratch_alloc (&set_hdr, (size_t)cap * sizeof(int64_t));
uint8_t* used = (uint8_t*)scratch_calloc(&used_hdr, (size_t)cap * sizeof(uint8_t));
if (!set || !used) {
if (set_hdr) scratch_free(set_hdr);
if (used_hdr) scratch_free(used_hdr);
return -1;
}
int64_t count = 0;
for (int64_t i = 0; i < len; i++) {
int64_t val;
if (in_type == RAY_F64) {
double fv = ((const double*)base)[i];
if (fv != fv) fv = (double)NAN;
else fv = clear_neg_zero(fv);
memcpy(&val, &fv, sizeof(int64_t));
} else {
val = read_col_i64(base, i, in_type, in_attrs);
}
uint64_t h = (uint64_t)val * CD_HASH_K1;
uint64_t slot = h >> cd_shift;
while (used[slot]) {
if (set[slot] == val) goto cd_seq_next;
slot = (slot + 1) & mask;
}
set[slot] = val;
used[slot] = 1;
count++;
cd_seq_next:;
}
scratch_free(set_hdr);
scratch_free(used_hdr);
return count;
}
static int64_t cd_sym_dense_count(ray_t* input) {
int64_t dom_n = ray_sym_domain_count(ray_sym_vec_domain(input));
uint32_t nsyms = (dom_n > 0 && dom_n <= UINT32_MAX) ? (uint32_t)dom_n : 0;
if (nsyms == 0) return dom_n == 0 ? 0 : -2;
ray_t* seen_hdr = NULL;
uint8_t* seen = (uint8_t*)scratch_calloc(&seen_hdr, (size_t)nsyms);
if (!seen) return -1;
const void* base = ray_data(input);
int64_t distinct = 0;
int64_t len = input->len;
uint8_t esz = ray_sym_elem_size(input->type, input->attrs);
#define CD_SYM_DENSE_LOOP(T) do { \
const T* ids = (const T*)base; \
for (int64_t i = 0; i < len; i++) { \
uint64_t id = (uint64_t)ids[i]; \
if (RAY_UNLIKELY(id >= nsyms)) { \
scratch_free(seen_hdr); \
return -2; \
} \
if (!seen[id]) { seen[id] = 1; distinct++; } \
} \
} while (0)
switch (esz) {
case 1: CD_SYM_DENSE_LOOP(uint8_t); break;
case 2: CD_SYM_DENSE_LOOP(uint16_t); break;
case 4: CD_SYM_DENSE_LOOP(uint32_t); break;
default: CD_SYM_DENSE_LOOP(uint64_t); break;
}
#undef CD_SYM_DENSE_LOOP
scratch_free(seen_hdr);
return distinct;
}
static ray_t* cd_widen_f32(ray_t* src) {
ray_t* out = ray_vec_new(RAY_F64, src->len);
if (!out || RAY_IS_ERR(out)) return out;
out->len = src->len;
const float* in = ray_data(src); double* dst = ray_data(out);
for (int64_t i = 0; i < src->len; i++) dst[i] = in[i];
if (ray_vec_may_have_nulls(src)) out->attrs |= RAY_ATTR_HAS_NULLS;
return out;
}
ray_t* exec_count_distinct(ray_graph_t* g, ray_op_t* op, ray_t* input) {
(void)g; (void)op;
if (!input || RAY_IS_ERR(input)) return input;
if (input->type == RAY_F32) {
ray_t* wide = cd_widen_f32(input);
if (!wide || RAY_IS_ERR(wide)) return wide ? wide : ray_error("oom", NULL);
ray_t* result = exec_count_distinct(g, op, wide); ray_release(wide); return result;
}
int8_t in_type = input->type;
int64_t len = input->len;
if (len == 0) return ray_i64(0);
switch (in_type) {
case RAY_BOOL: case RAY_U8:
case RAY_I16: case RAY_I32: case RAY_I64:
case RAY_F64: case RAY_DATE: case RAY_TIME: case RAY_TIMESTAMP:
case RAY_SYM:
break;
case RAY_STR:
case RAY_GUID:
case RAY_LIST: {
ray_t* dist = distinct_vec_eager(input);
if (!dist || RAY_IS_ERR(dist)) return dist ? dist : ray_error("oom", NULL);
int64_t cnt = ray_len(dist);
ray_release(dist);
return ray_i64(cnt);
}
default:
return ray_error("type", "count distinct: unsupported column type, got %s", ray_type_name(in_type));
}
void* base = ray_data(input);
ray_pool_t* pool = ray_pool_get();
if (in_type == RAY_SYM) {
int64_t cnt = cd_sym_dense_count(input);
if (cnt >= 0) return ray_i64(cnt);
if (cnt == -1) return ray_error("oom", NULL);
}
if (!pool || len < (1 << 16)) {
int64_t cnt = cd_seq_count(in_type, input->attrs, base, len);
if (cnt < 0) return ray_error("oom", NULL);
return ray_i64(cnt);
}
uint32_t nw = ray_pool_total_workers(pool);
uint32_t p_bits;
if (nw <= 8) p_bits = 4;
else if (nw <= 32) p_bits = 5;
else p_bits = 6;
uint64_t P = (uint64_t)1 << p_bits;
uint64_t p_mask = P - 1;
int64_t grain = (int64_t)RAY_DISPATCH_MORSELS * RAY_MORSEL_ELEMS;
if (grain <= 0) grain = 8192;
int64_t n_tasks_64 = (len + grain - 1) / grain;
if (n_tasks_64 <= 0) n_tasks_64 = 1;
if (n_tasks_64 > (1u << 16)) {
int64_t cnt = cd_seq_count(in_type, input->attrs, base, len);
if (cnt < 0) return ray_error("oom", NULL);
return ray_i64(cnt);
}
uint32_t n_tasks = (uint32_t)n_tasks_64;
ray_t* hist_hdr = NULL;
int64_t* hist = (int64_t*)scratch_calloc(&hist_hdr,
(size_t)P * n_tasks * sizeof(int64_t));
if (!hist) {
return ray_error("oom", NULL);
}
cd_count_ctx_t hctx = {
.base = base, .counts = hist,
.p_bits = p_bits, .p_mask = p_mask,
.grain = grain, .total = len,
.stride_log2 = 0, .is_f64 = (in_type == RAY_F64),
.type = in_type, .attrs = input->attrs,
};
ray_pool_dispatch_n(pool, cd_hist_fn, &hctx, n_tasks);
ray_t* off_hdr = NULL;
int64_t* part_off = (int64_t*)scratch_alloc(&off_hdr,
(size_t)(P + 1) * sizeof(int64_t));
if (!part_off) { scratch_free(hist_hdr); return ray_error("oom", NULL); }
ray_t* cur_hdr = NULL;
int64_t* cursor = (int64_t*)scratch_alloc(&cur_hdr,
(size_t)P * n_tasks * sizeof(int64_t));
if (!cursor) {
scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("oom", NULL);
}
int64_t total = 0;
for (uint64_t p = 0; p < P; p++) {
part_off[p] = total;
for (uint32_t t = 0; t < n_tasks; t++) {
cursor[(size_t)t * P + p] = total;
total += hist[(size_t)t * P + p];
}
}
part_off[P] = total;
if (total != len) {
scratch_free(cur_hdr); scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("nyi", "count_distinct: histogram mismatch");
}
ray_t* buf_hdr = NULL;
int64_t* out_buf = (int64_t*)scratch_alloc(&buf_hdr,
(size_t)len * sizeof(int64_t));
if (!out_buf) {
scratch_free(cur_hdr); scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("oom", NULL);
}
cd_scatter_ctx_t sctx = {
.base = base, .out_buf = out_buf, .cursor = cursor,
.p_bits = p_bits, .p_mask = p_mask,
.grain = grain, .total = len,
.is_f64 = (in_type == RAY_F64),
.type = in_type, .attrs = input->attrs,
};
ray_pool_dispatch_n(pool, cd_scatter_fn, &sctx, n_tasks);
ray_t* pcnt_hdr = NULL;
int64_t* part_count = (int64_t*)scratch_alloc(&pcnt_hdr,
(size_t)P * sizeof(int64_t));
if (!part_count) {
scratch_free(buf_hdr); scratch_free(cur_hdr);
scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("oom", NULL);
}
cd_part_ctx_t dctx = {
.values = out_buf, .part_off = part_off, .part_count = part_count,
};
ray_pool_dispatch_n(pool, cd_part_dedup_fn, &dctx, (uint32_t)P);
int64_t total_distinct = 0;
for (uint64_t p = 0; p < P; p++) {
if (part_count[p] < 0) {
scratch_free(pcnt_hdr); scratch_free(buf_hdr); scratch_free(cur_hdr);
scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("oom", NULL);
}
total_distinct += part_count[p];
}
scratch_free(pcnt_hdr); scratch_free(buf_hdr); scratch_free(cur_hdr);
scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_i64(total_distinct);
}
typedef struct {
const void* base;
int64_t* out_buf;
int64_t* rid_buf;
int64_t* cursor;
uint64_t p_mask;
int64_t grain;
int64_t total;
uint8_t is_f64;
int8_t type;
} dv_scatter_ctx_t;
static void dv_scatter_fn(void* ctx, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
(void)end;
dv_scatter_ctx_t* x = (dv_scatter_ctx_t*)ctx;
int64_t task_idx = start;
int64_t row_start = task_idx * x->grain;
int64_t row_end = row_start + x->grain;
if (row_end > x->total) row_end = x->total;
int64_t* cur = x->cursor + (size_t)task_idx * (x->p_mask + 1);
int64_t* out = x->out_buf;
int64_t* rid = x->rid_buf;
const void* base = x->base;
uint64_t p_mask = x->p_mask;
#define DV_SCATTER_BODY(LOAD) \
for (int64_t i = row_start; i < row_end; i++) { \
int64_t val = (LOAD); \
uint64_t h = (uint64_t)val * CD_HASH_K1; \
h ^= h >> 33; \
uint64_t p = (h ^ (h >> 33)) & p_mask; \
int64_t c = cur[p]++; \
out[c] = val; \
rid[c] = i; \
}
if (x->is_f64) {
const double* d = (const double*)base;
for (int64_t i = row_start; i < row_end; i++) {
double fv = d[i];
if (fv != fv) fv = (double)NAN;
else fv = clear_neg_zero(fv);
int64_t val;
memcpy(&val, &fv, sizeof(int64_t));
uint64_t h = (uint64_t)val * CD_HASH_K1;
h ^= h >> 33;
uint64_t p = (h ^ (h >> 33)) & p_mask;
int64_t c = cur[p]++;
out[c] = val;
rid[c] = i;
}
} else if (x->type == RAY_I64 || x->type == RAY_TIMESTAMP) {
const int64_t* d = (const int64_t*)base;
DV_SCATTER_BODY(d[i])
} else if (x->type == RAY_I32 || x->type == RAY_DATE || x->type == RAY_TIME) {
const int32_t* d = (const int32_t*)base;
DV_SCATTER_BODY(d[i])
} else if (x->type == RAY_I16) {
const int16_t* d = (const int16_t*)base;
DV_SCATTER_BODY(d[i])
}
#undef DV_SCATTER_BODY
}
typedef struct {
const int64_t* values;
const int64_t* rids;
const int64_t* part_off;
int64_t* part_count;
uint8_t* first;
} dv_part_ctx_t;
static void dv_part_dedup_fn(void* ctx, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
dv_part_ctx_t* x = (dv_part_ctx_t*)ctx;
for (int64_t p = start; p < end; p++) {
int64_t off = x->part_off[p];
int64_t cnt = x->part_off[p + 1] - off;
if (cnt == 0) { x->part_count[p] = 0; continue; }
uint64_t cap = (uint64_t)cnt * 2;
if (cap < 32) cap = 32;
uint64_t c = 1;
while (c && c < cap) c <<= 1;
if (!c) { x->part_count[p] = -1; continue; }
cap = c;
uint64_t mask = cap - 1;
int cd_shift = 64 - __builtin_ctzll(cap);
ray_t* set_hdr = NULL;
ray_t* used_hdr = NULL;
int64_t* set = (int64_t*)scratch_alloc (&set_hdr,
(size_t)cap * sizeof(int64_t));
uint8_t* used = (uint8_t*)scratch_calloc(&used_hdr,
(size_t)cap * sizeof(uint8_t));
if (!set || !used) {
if (set_hdr) scratch_free(set_hdr);
if (used_hdr) scratch_free(used_hdr);
x->part_count[p] = -1;
continue;
}
const int64_t* vals = x->values + off;
const int64_t* rids = x->rids + off;
int64_t count = 0;
for (int64_t j = 0; j < cnt; j++) {
int64_t val = vals[j];
uint64_t h = (uint64_t)val * CD_HASH_K1;
uint64_t slot = h >> cd_shift;
while (used[slot]) {
if (set[slot] == val) goto dv_part_next;
slot = (slot + 1) & mask;
}
set[slot] = val;
used[slot] = 1;
x->first[rids[j]] = 1;
count++;
dv_part_next:;
}
scratch_free(set_hdr);
scratch_free(used_hdr);
x->part_count[p] = count;
}
}
ray_t* distinct_radix_first_ids(ray_t* input) {
if (!input || RAY_IS_ERR(input)) return NULL;
int8_t in_type = input->type;
int64_t len = input->len;
switch (in_type) {
case RAY_I16: case RAY_I32: case RAY_I64:
case RAY_DATE: case RAY_TIME: case RAY_TIMESTAMP:
break;
case RAY_F64:
if (!ray_vec_may_have_nulls(input)) return NULL;
break;
default:
return NULL;
}
ray_pool_t* pool = ray_pool_get();
if (!pool || len < (1 << 16)) return NULL;
uint32_t nw = ray_pool_total_workers(pool);
uint32_t p_bits;
if (nw <= 8) p_bits = 4;
else if (nw <= 32) p_bits = 5;
else p_bits = 6;
uint64_t P = (uint64_t)1 << p_bits;
uint64_t p_mask = P - 1;
int64_t grain = (int64_t)RAY_DISPATCH_MORSELS * RAY_MORSEL_ELEMS;
if (grain <= 0) grain = 8192;
int64_t n_tasks_64 = (len + grain - 1) / grain;
if (n_tasks_64 <= 0 || n_tasks_64 > (1u << 16)) return NULL;
uint32_t n_tasks = (uint32_t)n_tasks_64;
void* base = ray_data(input);
ray_t* hist_hdr = NULL;
int64_t* hist = (int64_t*)scratch_calloc(&hist_hdr,
(size_t)P * n_tasks * sizeof(int64_t));
if (!hist) return ray_error("oom", NULL);
cd_count_ctx_t hctx = {
.base = base, .counts = hist,
.p_bits = p_bits, .p_mask = p_mask,
.grain = grain, .total = len,
.stride_log2 = 0, .is_f64 = (in_type == RAY_F64),
.type = in_type, .attrs = input->attrs,
};
ray_pool_dispatch_n(pool, cd_hist_fn, &hctx, n_tasks);
ray_t* off_hdr = NULL;
int64_t* part_off = (int64_t*)scratch_alloc(&off_hdr,
(size_t)(P + 1) * sizeof(int64_t));
if (!part_off) { scratch_free(hist_hdr); return ray_error("oom", NULL); }
ray_t* cur_hdr = NULL;
int64_t* cursor = (int64_t*)scratch_alloc(&cur_hdr,
(size_t)P * n_tasks * sizeof(int64_t));
if (!cursor) {
scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("oom", NULL);
}
int64_t total = 0;
for (uint64_t p = 0; p < P; p++) {
part_off[p] = total;
for (uint32_t t = 0; t < n_tasks; t++) {
cursor[(size_t)t * P + p] = total;
total += hist[(size_t)t * P + p];
}
}
part_off[P] = total;
if (total != len) {
scratch_free(cur_hdr); scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("nyi", "distinct: histogram mismatch");
}
ray_t* buf_hdr = NULL;
int64_t* out_buf = (int64_t*)scratch_alloc(&buf_hdr,
(size_t)len * sizeof(int64_t));
ray_t* rid_hdr = NULL;
int64_t* rid_buf = (int64_t*)scratch_alloc(&rid_hdr,
(size_t)len * sizeof(int64_t));
if (!out_buf || !rid_buf) {
if (buf_hdr) scratch_free(buf_hdr);
if (rid_hdr) scratch_free(rid_hdr);
scratch_free(cur_hdr); scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("oom", NULL);
}
dv_scatter_ctx_t sctx = {
.base = base, .out_buf = out_buf, .rid_buf = rid_buf, .cursor = cursor,
.p_mask = p_mask, .grain = grain, .total = len,
.is_f64 = (in_type == RAY_F64), .type = in_type,
};
ray_pool_dispatch_n(pool, dv_scatter_fn, &sctx, n_tasks);
ray_t* first_hdr = NULL;
uint8_t* first = (uint8_t*)scratch_calloc(&first_hdr, (size_t)len);
ray_t* pcnt_hdr = NULL;
int64_t* part_count = first ? (int64_t*)scratch_alloc(&pcnt_hdr,
(size_t)P * sizeof(int64_t))
: NULL;
if (!first || !part_count) {
if (first_hdr) scratch_free(first_hdr);
if (pcnt_hdr) scratch_free(pcnt_hdr);
scratch_free(rid_hdr); scratch_free(buf_hdr); scratch_free(cur_hdr);
scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("oom", NULL);
}
dv_part_ctx_t dctx = {
.values = out_buf, .rids = rid_buf, .part_off = part_off,
.part_count = part_count, .first = first,
};
ray_pool_dispatch_n(pool, dv_part_dedup_fn, &dctx, (uint32_t)P);
int64_t total_distinct = 0;
for (uint64_t p = 0; p < P; p++) {
if (part_count[p] < 0) {
scratch_free(pcnt_hdr); scratch_free(first_hdr);
scratch_free(rid_hdr); scratch_free(buf_hdr); scratch_free(cur_hdr);
scratch_free(off_hdr); scratch_free(hist_hdr);
return ray_error("oom", NULL);
}
total_distinct += part_count[p];
}
scratch_free(pcnt_hdr); scratch_free(rid_hdr); scratch_free(buf_hdr);
scratch_free(cur_hdr); scratch_free(off_hdr); scratch_free(hist_hdr);
ray_t* ids = ray_vec_new(RAY_I64, total_distinct);
if (!ids || RAY_IS_ERR(ids)) {
scratch_free(first_hdr);
return ids ? ids : ray_error("oom", NULL);
}
int64_t* out_ids = (int64_t*)ray_data(ids);
int64_t k = 0;
for (int64_t i = 0; i < len; i++)
if (first[i]) out_ids[k++] = i;
ids->len = k;
scratch_free(first_hdr);
if (k != total_distinct) {
ray_release(ids);
return ray_error("nyi", "distinct: first-occurrence sweep mismatch");
}
return ids;
}
#define CDPG_HASH(GID_P1, VAL) ({ \
uint64_t _h_ = (uint64_t)(VAL) * 0x9E3779B97F4A7C15ULL; \
_h_ ^= (uint64_t)(GID_P1) * 0xBF58476D1CE4E5B9ULL; \
_h_ ^= _h_ >> 33; \
_h_ *= 0xC4CEB9FE1A85EC53ULL; \
_h_; \
})
#define CDPG_PART_HASH(GID_P1) ({ \
uint64_t _h_ = (uint64_t)(GID_P1) * 0xBF58476D1CE4E5B9ULL; \
_h_ ^= _h_ >> 33; \
_h_ *= 0xC4CEB9FE1A85EC53ULL; \
_h_; \
})
typedef struct {
int8_t in_type;
uint8_t in_attrs;
const void* base;
const int64_t* row_gid;
int64_t n_rows;
int64_t n_groups;
uint64_t p_mask;
int64_t grain;
int64_t n_tasks;
int64_t* hist;
int64_t* cursor;
int64_t* part_off;
int64_t* gids_out;
int64_t* vals_out;
int64_t* odata;
} cdpg_ctx_t;
static inline int64_t cdpg_read(const void* base, int64_t r,
int8_t in_type, uint8_t esz) {
if (in_type == RAY_F64) return canon_f64_key(((const double*)base)[r]);
switch (esz) {
case 1: return (int64_t)((const uint8_t*)base)[r];
case 2: return (int64_t)((const int16_t*)base)[r];
case 4: return (int64_t)((const int32_t*)base)[r];
default: return ((const int64_t*)base)[r];
}
}
static void cdpg_hist_fn(void* ctx_, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
cdpg_ctx_t* x = (cdpg_ctx_t*)ctx_;
uint64_t p_mask = x->p_mask;
int64_t task_id = start / x->grain;
int64_t* my_hist = &x->hist[task_id * (p_mask + 1)];
for (int64_t r = start; r < end; r++) {
int64_t gid = x->row_gid[r];
if (gid < 0 || gid >= x->n_groups) continue;
uint64_t h = CDPG_PART_HASH(gid + 1);
my_hist[h & p_mask]++;
}
}
static void cdpg_scat_fn(void* ctx_, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
cdpg_ctx_t* x = (cdpg_ctx_t*)ctx_;
uint8_t esz = ray_sym_elem_size(x->in_type, x->in_attrs);
uint64_t p_mask = x->p_mask;
int64_t task_id = start / x->grain;
int64_t* my_cur = &x->cursor[task_id * (p_mask + 1)];
for (int64_t r = start; r < end; r++) {
int64_t gid = x->row_gid[r];
if (gid < 0 || gid >= x->n_groups) continue;
int64_t val = cdpg_read(x->base, r, x->in_type, esz);
int64_t gid_p1 = gid + 1;
uint64_t h = CDPG_PART_HASH(gid_p1);
int64_t pos = my_cur[h & p_mask]++;
x->gids_out[pos] = gid_p1;
x->vals_out[pos] = val;
}
}
static void cdpg_dedup_fn(void* ctx_, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
cdpg_ctx_t* x = (cdpg_ctx_t*)ctx_;
for (int64_t p = start; p < end; p++) {
int64_t off = x->part_off[p];
int64_t cnt = x->part_off[p + 1] - off;
if (cnt == 0) continue;
uint64_t cap = (uint64_t)cnt * 2;
if (cap < 32) cap = 32;
uint64_t c = 1;
while (c && c < cap) c <<= 1;
if (!c) continue;
cap = c;
uint64_t mask = cap - 1;
ray_t* k_hdr = NULL;
ray_t* v_hdr = NULL;
int64_t* slot_gid = (int64_t*)scratch_calloc(&k_hdr,
(size_t)cap * sizeof(int64_t));
int64_t* slot_val = (int64_t*)scratch_alloc(&v_hdr,
(size_t)cap * sizeof(int64_t));
if (!slot_gid || !slot_val) {
if (k_hdr) scratch_free(k_hdr);
if (v_hdr) scratch_free(v_hdr);
continue;
}
const int64_t* gids = x->gids_out + off;
const int64_t* vals = x->vals_out + off;
for (int64_t i = 0; i < cnt; i++) {
int64_t gid_p1 = gids[i];
int64_t val = vals[i];
uint64_t h = CDPG_HASH(gid_p1, val);
uint64_t slot = h & mask;
for (;;) {
int64_t cur = slot_gid[slot];
if (cur == 0) {
slot_gid[slot] = gid_p1;
slot_val[slot] = val;
x->odata[gid_p1 - 1]++;
break;
}
if (cur == gid_p1 && slot_val[slot] == val) break;
slot = (slot + 1) & mask;
}
}
scratch_free(k_hdr);
scratch_free(v_hdr);
}
}
static ray_t* count_distinct_per_group_parallel(
ray_t* src, const int64_t* row_gid,
int64_t n_rows, int64_t n_groups, ray_t* out)
{
ray_pool_t* pool = ray_pool_get();
if (!pool) return NULL;
uint32_t nw = ray_pool_total_workers(pool);
if (nw < 2) return NULL;
uint8_t p_bits = 6;
uint64_t P = (uint64_t)1 << p_bits;
uint64_t p_mask = P - 1;
cdpg_ctx_t ctx = {
.in_type = src->type,
.in_attrs = src->attrs,
.base = ray_data(src),
.row_gid = row_gid,
.n_rows = n_rows,
.n_groups = n_groups,
.p_mask = p_mask,
.odata = (int64_t*)ray_data(out),
};
if (P > 256) return NULL;
int64_t grain = (int64_t)RAY_DISPATCH_MORSELS * RAY_MORSEL_ELEMS;
int64_t n_tasks = (n_rows + grain - 1) / grain;
if (n_tasks > 65536) {
n_tasks = 65536;
grain = (n_rows + n_tasks - 1) / n_tasks;
}
ctx.grain = grain;
ctx.n_tasks = n_tasks;
ray_t* hist_hdr = NULL;
ctx.hist = (int64_t*)scratch_calloc(&hist_hdr,
(size_t)n_tasks * (size_t)P * sizeof(int64_t));
if (!ctx.hist) { return NULL; }
ray_pool_dispatch(pool, cdpg_hist_fn, &ctx, n_rows);
ray_t* off_hdr = NULL;
ctx.part_off = (int64_t*)scratch_alloc(&off_hdr,
(size_t)(P + 1) * sizeof(int64_t));
ray_t* cur_hdr = NULL;
ctx.cursor = (int64_t*)scratch_alloc(&cur_hdr,
(size_t)n_tasks * (size_t)P * sizeof(int64_t));
if (!ctx.part_off || !ctx.cursor) {
if (off_hdr) scratch_free(off_hdr);
if (cur_hdr) scratch_free(cur_hdr);
scratch_free(hist_hdr);
return NULL;
}
int64_t total = 0;
for (uint64_t p = 0; p < P; p++) {
ctx.part_off[p] = total;
int64_t cum = total;
for (int64_t t = 0; t < n_tasks; t++) {
int64_t cnt = ctx.hist[t * P + p];
ctx.cursor[t * P + p] = cum;
cum += cnt;
}
total = cum;
}
ctx.part_off[P] = total;
ray_t* gids_hdr = NULL;
ray_t* vals_hdr = NULL;
ctx.gids_out = (int64_t*)scratch_alloc(&gids_hdr,
(size_t)total * sizeof(int64_t));
ctx.vals_out = (int64_t*)scratch_alloc(&vals_hdr,
(size_t)total * sizeof(int64_t));
if (!ctx.gids_out || !ctx.vals_out) {
if (gids_hdr) scratch_free(gids_hdr);
if (vals_hdr) scratch_free(vals_hdr);
scratch_free(cur_hdr); scratch_free(off_hdr); scratch_free(hist_hdr);
return NULL;
}
if (total > 0)
ray_pool_dispatch(pool, cdpg_scat_fn, &ctx, n_rows);
if (total > 0)
ray_pool_dispatch_n(pool, cdpg_dedup_fn, &ctx, (uint32_t)P);
scratch_free(vals_hdr); scratch_free(gids_hdr);
scratch_free(cur_hdr); scratch_free(off_hdr);
scratch_free(hist_hdr);
return out;
}
__attribute__((unused)) static ray_t* count_distinct_per_group_hll(ray_t* src, const int64_t* row_gid,
int64_t n_rows, int64_t n_groups,
ray_t* out) {
if (!src || n_rows <= 0 || n_groups <= 0) return NULL;
ray_t* cnt_hdr = NULL;
ray_t* off_hdr = NULL;
int64_t* counts = (int64_t*)scratch_calloc(&cnt_hdr,
(size_t)n_groups * sizeof(int64_t));
int64_t* offsets = (int64_t*)scratch_alloc(&off_hdr,
(size_t)n_groups * sizeof(int64_t));
if (!counts || !offsets) {
if (cnt_hdr) scratch_free(cnt_hdr);
if (off_hdr) scratch_free(off_hdr);
return NULL;
}
int64_t total = 0;
for (int64_t r = 0; r < n_rows; r++) {
int64_t g = row_gid[r];
if (g >= 0 && g < n_groups) counts[g]++;
}
for (int64_t g = 0; g < n_groups; g++) {
offsets[g] = total;
total += counts[g];
}
if (total == 0) {
scratch_free(cnt_hdr); scratch_free(off_hdr);
return out;
}
ray_t* idx_hdr = NULL;
int64_t* idx_buf = (int64_t*)scratch_alloc(&idx_hdr,
(size_t)total * sizeof(int64_t));
if (!idx_buf) {
scratch_free(cnt_hdr); scratch_free(off_hdr);
return NULL;
}
ray_t* pos_hdr = NULL;
int64_t* pos = (int64_t*)scratch_alloc(&pos_hdr,
(size_t)n_groups * sizeof(int64_t));
if (!pos) {
scratch_free(idx_hdr); scratch_free(cnt_hdr); scratch_free(off_hdr);
return NULL;
}
memcpy(pos, offsets, (size_t)n_groups * sizeof(int64_t));
for (int64_t r = 0; r < n_rows; r++) {
int64_t g = row_gid[r];
if (g >= 0 && g < n_groups) idx_buf[pos[g]++] = r;
}
scratch_free(pos_hdr);
int64_t* odata = (int64_t*)ray_data(out);
int rc = ray_count_distinct_approx_pg_buf(src, idx_buf, offsets, counts,
n_groups, RAY_HLL_DEFAULT_P, odata);
scratch_free(idx_hdr);
scratch_free(cnt_hdr);
scratch_free(off_hdr);
if (rc != 0) return NULL;
return out;
}
ray_t* ray_count_distinct_per_group(ray_t* src, const int64_t* row_gid,
int64_t n_rows, int64_t n_groups) {
if (!src || RAY_IS_ERR(src)) return ray_error("domain", "count distinct per group: invalid source column");
if (n_groups < 0) return ray_error("domain", "count distinct per group: group count must be non-negative, got %lld", (long long)n_groups);
if (src->type == RAY_F32) {
ray_t* wide = cd_widen_f32(src);
if (!wide || RAY_IS_ERR(wide)) return wide ? wide : ray_error("oom", NULL);
ray_t* result = ray_count_distinct_per_group(wide, row_gid, n_rows, n_groups);
ray_release(wide); return result;
}
int8_t in_type = src->type;
switch (in_type) {
case RAY_BOOL: case RAY_U8:
case RAY_I16: case RAY_I32: case RAY_I64:
case RAY_F64: case RAY_DATE: case RAY_TIME: case RAY_TIMESTAMP:
case RAY_SYM:
break;
default:
return NULL;
}
if (src->len < n_rows) return ray_error("domain", "count distinct per group: source has fewer rows than required, got %lld need %lld", (long long)src->len, (long long)n_rows);
ray_t* out = ray_vec_new(RAY_I64, n_groups);
if (!out || RAY_IS_ERR(out)) return out ? out : ray_error("oom", NULL);
out->len = n_groups;
int64_t* odata = (int64_t*)ray_data(out);
memset(odata, 0, (size_t)n_groups * sizeof(int64_t));
if (n_rows == 0 || n_groups == 0) return out;
if (n_rows >= 200000) {
ray_t* par = count_distinct_per_group_parallel(src, row_gid,
n_rows, n_groups, out);
if (par) return par;
}
uint64_t cap = (uint64_t)n_rows * 2;
if (cap < 32) cap = 32;
uint64_t c = 1;
while (c && c < cap) c <<= 1;
if (!c) { ray_release(out); return ray_error("oom", NULL); }
cap = c;
uint64_t mask = cap - 1;
ray_t* k_hdr = NULL;
ray_t* v_hdr = NULL;
int64_t* slot_gid = (int64_t*)scratch_calloc(&k_hdr,
(size_t)cap * sizeof(int64_t));
int64_t* slot_val = (int64_t*)scratch_alloc(&v_hdr,
(size_t)cap * sizeof(int64_t));
if (!slot_gid || !slot_val) {
if (k_hdr) scratch_free(k_hdr);
if (v_hdr) scratch_free(v_hdr);
ray_release(out);
return ray_error("oom", NULL);
}
void* base = ray_data(src);
uint8_t esz = ray_sym_elem_size(in_type, src->attrs);
#define CD_INSERT(VAL_EXPR) do { \
int64_t val = (VAL_EXPR); \
int64_t gid_p1 = gid + 1; \
uint64_t h = (uint64_t)val * 0x9E3779B97F4A7C15ULL; \
h ^= (uint64_t)gid_p1 * 0xBF58476D1CE4E5B9ULL; \
h ^= h >> 33; \
h *= 0xC4CEB9FE1A85EC53ULL; \
uint64_t slot = h & mask; \
for (;;) { \
int64_t cur = slot_gid[slot]; \
if (cur == 0) { \
slot_gid[slot] = gid_p1; \
slot_val[slot] = val; \
odata[gid]++; \
break; \
} \
if (cur == gid_p1 && slot_val[slot] == val) break; \
slot = (slot + 1) & mask; \
} \
} while (0)
{
if (in_type == RAY_F64) {
const double* d = (const double*)base;
for (int64_t r = 0; r < n_rows; r++) {
int64_t gid = row_gid[r];
if (gid < 0 || gid >= n_groups) continue;
CD_INSERT(canon_f64_key(d[r]));
}
} else if (esz == 8) {
const int64_t* d = (const int64_t*)base;
for (int64_t r = 0; r < n_rows; r++) {
int64_t gid = row_gid[r];
if (gid < 0 || gid >= n_groups) continue;
CD_INSERT(d[r]);
}
} else if (esz == 4) {
const int32_t* d = (const int32_t*)base;
for (int64_t r = 0; r < n_rows; r++) {
int64_t gid = row_gid[r];
if (gid < 0 || gid >= n_groups) continue;
CD_INSERT((int64_t)d[r]);
}
} else if (esz == 2) {
const int16_t* d = (const int16_t*)base;
for (int64_t r = 0; r < n_rows; r++) {
int64_t gid = row_gid[r];
if (gid < 0 || gid >= n_groups) continue;
CD_INSERT((int64_t)d[r]);
}
} else {
const uint8_t* d = (const uint8_t*)base;
for (int64_t r = 0; r < n_rows; r++) {
int64_t gid = row_gid[r];
if (gid < 0 || gid >= n_groups) continue;
CD_INSERT((int64_t)d[r]);
}
}
}
#undef CD_INSERT
scratch_free(k_hdr);
scratch_free(v_hdr);
return out;
}
typedef struct {
const void* base;
int8_t src_type;
bool has_nulls;
bool use_quantile;
double q;
const int64_t* idx_buf;
const int64_t* offsets;
const int64_t* grp_cnt;
double* scratch_pool;
double* out_data;
ray_t* out;
_Atomic(bool) any_null;
const uint8_t* handled;
} med_par_ctx_t;
typedef struct {
uint64_t key;
int64_t count;
int64_t first_row;
int64_t first_pos;
uint8_t used;
} mode_scalar_entry_t;
typedef struct {
uint64_t hash;
int64_t count;
int64_t first_row;
int64_t first_pos;
uint8_t used;
} mode_wide_entry_t;
typedef struct {
ray_t* src;
int64_t* winners;
const int64_t* idx_buf;
const int64_t* offsets;
const int64_t* grp_cnt;
int64_t threshold;
_Atomic(int) oom;
_Atomic(int) cancel;
} mode_par_ctx_t;
static inline double med_read_as_f64(const void* base, int8_t t, int64_t row) {
switch (t) {
case RAY_F64: { double v; memcpy(&v, (const char*)base + (size_t)row * 8, 8); return v; }
case RAY_F32: { float v; memcpy(&v, (const char*)base + (size_t)row * 4, 4); return (double)v; }
case RAY_I64: { int64_t v; memcpy(&v, (const char*)base + (size_t)row * 8, 8); return (double)v; }
case RAY_I32: { int32_t v; memcpy(&v, (const char*)base + (size_t)row * 4, 4); return (double)v; }
case RAY_DATE:
case RAY_TIME: { int32_t v; memcpy(&v, (const char*)base + (size_t)row * 4, 4); return (double)v; }
case RAY_TIMESTAMP: { int64_t v; memcpy(&v, (const char*)base + (size_t)row * 8, 8); return (double)v; }
case RAY_I16: { int16_t v; memcpy(&v, (const char*)base + (size_t)row * 2, 2); return (double)v; }
case RAY_BOOL: case RAY_U8: return (double)((const uint8_t*)base)[row];
default: return 0.0;
}
}
static inline bool med_is_null(const void* base, int8_t t, int64_t row) {
switch (t) {
case RAY_F64: { double v; memcpy(&v, (const char*)base + (size_t)row * 8, 8); return v != v; }
case RAY_F32: { float v; memcpy(&v, (const char*)base + (size_t)row * 4, 4); return v != v; }
case RAY_I64: return ((const int64_t*)base)[row] == NULL_I64;
case RAY_I32: return ((const int32_t*)base)[row] == NULL_I32;
case RAY_DATE:
case RAY_TIME: return ((const int32_t*)base)[row] == NULL_I32;
case RAY_TIMESTAMP: return ((const int64_t*)base)[row] == NULL_I64;
case RAY_I16: return ((const int16_t*)base)[row] == NULL_I16;
case RAY_BOOL: case RAY_U8: return false;
default: return false;
}
}
typedef struct {
ray_pool_fn fn;
void* context;
int64_t begin[RAY_POOL_INIT_TASKS];
int64_t end[RAY_POOL_INIT_TASKS];
} group_work_t;
static void group_work_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
group_work_t* c = raw;
for (int64_t t = start; t < end; t++) c->fn(c->context, wid, c->begin[t], c->end[t]);
}
void ray_group_dispatch(ray_pool_fn fn, void* context, const int64_t* counts, int64_t groups) {
ray_pool_t* pool = ray_pool_get();
int64_t work = 0;
for (int64_t g = 0; g < groups; g++) work += counts[g] + 1;
if (groups < 2 || !ray_pool_par_dispatch_ok(pool, work, 4096)) {
fn(context, 0, 0, groups); return;
}
uint32_t budget = ray_pool_total_workers(pool) * 8;
if (budget > RAY_POOL_INIT_TASKS / 2) budget = RAY_POOL_INIT_TASKS / 2;
int64_t grain = (work + budget - 1) / budget;
group_work_t tasks = {.fn = fn, .context = context};
uint32_t n = 0;
int64_t begin = 0, cost = 0;
for (int64_t g = 0; g < groups; g++) {
int64_t next = counts[g] + 1;
if (g > begin && cost + next > grain) {
assert(n < RAY_POOL_INIT_TASKS);
tasks.begin[n] = begin; tasks.end[n++] = g; begin = g; cost = 0;
}
cost += next;
if (cost >= grain) {
assert(n < RAY_POOL_INIT_TASKS);
tasks.begin[n] = begin; tasks.end[n++] = g + 1; begin = g + 1; cost = 0;
}
}
if (begin < groups) {
assert(n < RAY_POOL_INIT_TASKS);
tasks.begin[n] = begin; tasks.end[n++] = groups;
}
ray_pool_dispatch_n(pool, group_work_run, &tasks, n);
}
typedef struct {
ray_group_winner_fn fn;
void* context;
const int64_t* rows;
const int64_t* offsets;
const int64_t* counts;
int64_t* winners;
int64_t threshold;
struct { int64_t group, begin, count, winner; } partial[RAY_POOL_INIT_TASKS];
} group_winners_t;
static void group_winners_small(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; group_winners_t* c = raw;
for (int64_t g = start; g < end; g++) {
if (c->counts[g] >= c->threshold) continue;
c->winners[g] = c->fn(c->context, c->rows + c->offsets[g], c->counts[g]);
}
}
static void group_winners_partial(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; group_winners_t* c = raw;
for (int64_t t = start; t < end; t++)
c->partial[t].winner = c->fn(c->context, c->rows + c->partial[t].begin, c->partial[t].count);
}
void ray_group_winners(ray_group_winner_fn fn, void* context, const int64_t* rows,
const int64_t* offsets, const int64_t* counts, int64_t groups, int64_t* winners) {
ray_pool_t* pool = ray_pool_get();
int64_t total = 0;
for (int64_t g = 0; g < groups; g++) total += counts[g];
group_winners_t c = {.fn = fn, .context = context, .rows = rows,
.offsets = offsets, .counts = counts, .winners = winners, .threshold = INT64_MAX};
uint32_t tasks = 0;
if (ray_pool_par_dispatch_ok(pool, total, 262144)) {
uint32_t budget = ray_pool_total_workers(pool) * 4;
if (budget > RAY_POOL_INIT_TASKS / 2) budget = RAY_POOL_INIT_TASKS / 2;
int64_t grain = (total + budget - 1) / budget;
c.threshold = total / ray_pool_total_workers(pool);
if (c.threshold < 262144) c.threshold = 262144;
if (c.threshold < grain) c.threshold = grain;
for (int64_t g = 0; g < groups; g++) if (counts[g] >= c.threshold) {
winners[g] = -1;
for (int64_t off = 0; off < counts[g]; off += grain) {
assert(tasks < RAY_POOL_INIT_TASKS);
c.partial[tasks].group = g;
c.partial[tasks].begin = offsets[g] + off;
c.partial[tasks].count = counts[g] - off < grain ? counts[g] - off : grain;
c.partial[tasks++].winner = -1;
}
}
}
ray_group_dispatch(group_winners_small, &c, counts, groups);
if (tasks && !ray_interrupted() && !atomic_load(&pool->cancelled)) {
ray_pool_dispatch_n(pool, group_winners_partial, &c, tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) return;
for (uint32_t t = 0; t < tasks; t++) {
int64_t g = c.partial[t].group, next = c.partial[t].winner;
if (next < 0) continue;
if (winners[g] < 0) winners[g] = next;
else {
int64_t pair[] = {winners[g], next};
winners[g] = fn(context, pair, 2);
}
}
}
}
static void med_per_group_fn(void* ctx_v, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
med_par_ctx_t* c = (med_par_ctx_t*)ctx_v;
bool any_null = false;
for (int64_t g = start; g < end; g++) {
if (c->handled && c->handled[g]) continue;
int64_t cnt = c->grp_cnt[g];
int64_t off = c->offsets[g];
double* slice = c->scratch_pool ? c->scratch_pool + off : NULL;
int64_t actual = 0;
if (c->has_nulls) {
for (int64_t i = 0; i < cnt; i++) {
int64_t row = c->idx_buf[off + i];
if (med_is_null(c->base, c->src_type, row)) continue;
slice[actual++] = med_read_as_f64(c->base, c->src_type, row);
}
} else {
for (int64_t i = 0; i < cnt; i++) {
int64_t row = c->idx_buf[off + i];
slice[actual++] = med_read_as_f64(c->base, c->src_type, row);
}
}
if (actual == 0) {
c->out_data[g] = NULL_F64;
any_null = true;
} else {
c->out_data[g] = c->use_quantile
? ray_quantile_dbl_inplace(slice, actual, c->q)
: ray_median_dbl_inplace(slice, actual);
}
}
if (any_null) atomic_store_explicit(&c->any_null, true, memory_order_relaxed);
}
typedef struct {
int64_t less, equal, count, offset;
double max_less, min_greater;
} rank_task_t;
typedef struct {
const med_par_ctx_t* source;
const int64_t* rows;
double* input;
double* output;
int64_t n;
uint32_t tasks;
double pivot;
bool keep_less;
rank_task_t task[RAY_POOL_INIT_TASKS];
} rank_split_t;
static void rank_gather_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; rank_split_t* c = raw;
const med_par_ctx_t* s = c->source;
for (int64_t t = start; t < end; t++) {
int64_t begin = c->n / c->tasks * t;
int64_t limit = t + 1 == c->tasks ? c->n : c->n / c->tasks * (t + 1);
int64_t count = 0;
for (int64_t i = begin; i < limit; i++) {
int64_t row = c->rows[i];
if (s->has_nulls && med_is_null(s->base, s->src_type, row)) continue;
c->input[begin + count++] = med_read_as_f64(s->base, s->src_type, row);
}
c->task[t].count = count;
}
}
static void rank_compact_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; rank_split_t* c = raw;
for (int64_t t = start; t < end; t++)
memcpy(c->output + c->task[t].offset, c->input + c->n / c->tasks * t,
(size_t)c->task[t].count * sizeof(double));
}
static void rank_count_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; rank_split_t* c = raw;
for (int64_t t = start; t < end; t++) {
int64_t begin = c->n / c->tasks * t;
int64_t limit = t + 1 == c->tasks ? c->n : c->n / c->tasks * (t + 1);
int64_t less = 0, equal = 0;
double max_less = -INFINITY, min_greater = INFINITY;
for (int64_t i = begin; i < limit; i++) {
double v = c->input[i];
if (v < c->pivot) { less++; if (v > max_less) max_less = v; }
else if (v > c->pivot) { if (v < min_greater) min_greater = v; }
else equal++;
}
c->task[t].less = less; c->task[t].equal = equal;
c->task[t].max_less = max_less; c->task[t].min_greater = min_greater;
}
}
static void rank_scatter_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; rank_split_t* c = raw;
for (int64_t t = start; t < end; t++) {
int64_t begin = c->n / c->tasks * t;
int64_t limit = t + 1 == c->tasks ? c->n : c->n / c->tasks * (t + 1);
int64_t at = c->task[t].offset;
for (int64_t i = begin; i < limit; i++) {
double v = c->input[i];
if (c->keep_less ? v < c->pivot : v > c->pivot) c->output[at++] = v;
}
}
}
static bool rank_large_group(med_par_ctx_t* source, int64_t group,
double* temporary, ray_pool_t* pool, double* answer) {
rank_split_t c = {.source = source, .rows = source->idx_buf + source->offsets[group],
.input = source->scratch_pool + source->offsets[group], .output = temporary,
.n = source->grp_cnt[group], .tasks = ray_pool_total_workers(pool) * 4};
if (c.tasks > RAY_POOL_INIT_TASKS) c.tasks = RAY_POOL_INIT_TASKS;
ray_pool_dispatch_n(pool, rank_gather_run, &c, c.tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) return false;
int64_t valid = 0;
for (uint32_t t = 0; t < c.tasks; t++) {
c.task[t].offset = valid; valid += c.task[t].count;
}
if (!valid) { *answer = NULL_F64; atomic_store(&source->any_null, true); return true; }
if (valid != c.n) {
ray_pool_dispatch_n(pool, rank_compact_run, &c, c.tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) return false;
double* swap = c.input; c.input = c.output; c.output = swap;
}
c.n = valid;
double fraction = 0;
int64_t lo, hi;
if (source->use_quantile) {
double position = source->q * (double)(valid - 1);
lo = (int64_t)position; fraction = position - (double)lo;
hi = fraction > 0 ? lo + 1 : lo;
} else { lo = (valid - 1) / 2; hi = valid / 2; }
double lower = 0, upper = 0;
bool found = false;
for (unsigned depth = 0; depth < 32 && c.n >= RAY_PARALLEL_THRESHOLD; depth++) {
double sample[33];
int64_t step = c.n / 33;
for (int i = 0; i < 33; i++)
sample[i] = c.input[i * step + ray_hash_i64(i + depth * 33) % step];
c.pivot = ray_median_dbl_inplace(sample, 33);
ray_pool_dispatch_n(pool, rank_count_run, &c, c.tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) return false;
int64_t less = 0, equal = 0;
double max_less = -INFINITY, min_greater = INFINITY;
for (uint32_t t = 0; t < c.tasks; t++) {
less += c.task[t].less; equal += c.task[t].equal;
if (c.task[t].max_less > max_less) max_less = c.task[t].max_less;
if (c.task[t].min_greater < min_greater) min_greater = c.task[t].min_greater;
}
if (hi < less) c.keep_less = true;
else if (lo >= less + equal) c.keep_less = false;
else {
lower = lo < less ? max_less : c.pivot;
upper = hi >= less + equal ? min_greater : c.pivot;
found = true; break;
}
int64_t next = 0;
for (uint32_t t = 0; t < c.tasks; t++) {
c.task[t].offset = next;
int64_t begin = c.n / c.tasks * t;
int64_t limit = t + 1 == c.tasks ? c.n : c.n / c.tasks * (t + 1);
next += c.keep_less ? c.task[t].less : limit - begin - c.task[t].less - c.task[t].equal;
}
ray_pool_dispatch_n(pool, rank_scatter_run, &c, c.tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) return false;
if (!c.keep_less) { lo -= less + equal; hi -= less + equal; }
c.n = next;
double* swap = c.input; c.input = c.output; c.output = swap;
}
if (!found) {
lower = ray_nth_dbl_inplace(c.input, c.n, lo);
upper = hi == lo ? lower : ray_nth_dbl_inplace(c.input, c.n, hi);
}
*answer = source->use_quantile
? (fraction > 0 ? lower + fraction * (upper - lower) : lower)
: (valid % 2 ? lower : (lower + upper) / 2.0);
return true;
}
static double decode_agg_f64_param(int64_t bits) {
double v = 0.0;
memcpy(&v, &bits, sizeof(v));
return v;
}
static ray_t* rank_per_group_buf(ray_t* src,
const int64_t* idx_buf,
const int64_t* offsets,
const int64_t* grp_cnt,
int64_t n_groups,
double q,
bool use_quantile) {
if (!src || RAY_IS_ERR(src) || n_groups < 0) return NULL;
int8_t t = src->type;
if (t != RAY_F64 && t != RAY_F32 && t != RAY_I64 && t != RAY_I32 &&
t != RAY_I16 && t != RAY_U8 && t != RAY_BOOL && t != RAY_DATE &&
t != RAY_TIME && t != RAY_TIMESTAMP) return NULL;
if (use_quantile && (!__builtin_isfinite(q) || q < 0.0 || q > 1.0))
return ray_error("domain", "quantile: probability out of range");
int64_t total = 0;
for (int64_t g = 0; g < n_groups; g++) total += grp_cnt[g];
ray_t* out = ray_vec_new(RAY_F64, n_groups);
if (!out || RAY_IS_ERR(out)) return out ? out : ray_error("oom", NULL);
out->len = n_groups;
ray_t* buf_hdr = NULL;
double* scratch = NULL;
if (total > 0) {
scratch = (double*)scratch_alloc(&buf_hdr,
(size_t)total * sizeof(double));
if (!scratch) { ray_release(out); return ray_error("oom", NULL); }
}
med_par_ctx_t ctx = {
.base = ray_data(src),
.src_type = t,
.has_nulls = ray_vec_may_have_nulls(src),
.use_quantile = use_quantile,
.q = q,
.idx_buf = idx_buf,
.offsets = offsets,
.grp_cnt = grp_cnt,
.scratch_pool = scratch,
.out_data = (double*)ray_data(out),
.out = out,
};
ray_pool_t* pool = ray_pool_get();
int64_t threshold = RAY_PARALLEL_THRESHOLD, largest = 0;
if (ray_pool_par_dispatch_ok(pool, total, threshold)) {
int64_t share = total / ray_pool_total_workers(pool);
if (share > threshold) threshold = share;
for (int64_t g = 0; g < n_groups; g++)
if (grp_cnt[g] >= threshold && grp_cnt[g] > largest) largest = grp_cnt[g];
}
double* temporary = largest ? ray_alloc_raw((size_t)largest * sizeof(double)) : NULL;
uint8_t* handled = temporary ? ray_calloc_raw((size_t)n_groups) : NULL;
if (handled) {
ctx.handled = handled;
for (int64_t g = 0; g < n_groups; g++) if (grp_cnt[g] >= threshold) {
if (!rank_large_group(&ctx, g, temporary, pool, &ctx.out_data[g])) break;
handled[g] = 1;
}
}
ray_free_raw(temporary);
ray_profile_tick("rank: selected dominant groups");
ray_group_dispatch(med_per_group_fn, &ctx, grp_cnt, n_groups);
ray_profile_tick("rank: reduced remaining groups");
ray_free_raw(handled);
if (ray_interrupted() || (pool && atomic_load(&pool->cancelled))) {
if (buf_hdr) scratch_free(buf_hdr);
ray_release(out); return ray_error("cancel", NULL);
}
if (atomic_load_explicit(&ctx.any_null, memory_order_relaxed)) out->attrs |= RAY_ATTR_HAS_NULLS;
if (buf_hdr) scratch_free(buf_hdr);
return out;
}
ray_t* ray_median_per_group_buf(ray_t* src,
const int64_t* idx_buf,
const int64_t* offsets,
const int64_t* grp_cnt,
int64_t n_groups) {
return rank_per_group_buf(src, idx_buf, offsets, grp_cnt, n_groups, 0.0, false);
}
ray_t* ray_quantile_per_group_buf(ray_t* src,
const int64_t* idx_buf,
const int64_t* offsets,
const int64_t* grp_cnt,
int64_t n_groups,
double q) {
return rank_per_group_buf(src, idx_buf, offsets, grp_cnt, n_groups, q, true);
}
static uint64_t mode_hash_cap(int64_t n) {
uint64_t cap = 16;
uint64_t need = n > 0 ? (uint64_t)n * 2u : 16u;
if (need < 16) need = 16;
while (cap < need && cap <= (UINT64_MAX >> 1)) cap <<= 1;
return cap >= need ? cap : 0;
}
static inline uint64_t mode_scalar_key(ray_t* src, int64_t row) {
const void* base = ray_data(src);
switch (src->type) {
case RAY_F64: {
uint64_t k;
memcpy(&k, (const char*)base + (size_t)row * 8, 8);
if (k == 0x8000000000000000ULL) k = 0;
return k;
}
case RAY_F32: {
uint32_t k32;
memcpy(&k32, (const char*)base + (size_t)row * 4, 4);
if (k32 == 0x80000000U) k32 = 0;
return (uint64_t)k32;
}
case RAY_SYM:
return (uint64_t)ray_read_sym(base, row, src->type, src->attrs);
case RAY_I64:
case RAY_TIMESTAMP: {
int64_t v;
memcpy(&v, (const char*)base + (size_t)row * 8, 8);
return (uint64_t)v;
}
case RAY_I32:
case RAY_DATE:
case RAY_TIME: {
int32_t v;
memcpy(&v, (const char*)base + (size_t)row * 4, 4);
return (uint64_t)(int64_t)v;
}
case RAY_I16: {
int16_t v;
memcpy(&v, (const char*)base + (size_t)row * 2, 2);
return (uint64_t)(int64_t)v;
}
case RAY_BOOL:
case RAY_U8:
return (uint64_t)((const uint8_t*)base)[row];
default:
return 0;
}
}
static int64_t mode_scalar_group(ray_t* src, const int64_t* rows,
int64_t cnt, int* cancelled) {
uint64_t cap = mode_hash_cap(cnt);
if (!cap || cap > SIZE_MAX / sizeof(mode_scalar_entry_t)) return -2;
ray_t* hdr = NULL;
mode_scalar_entry_t* ht = (mode_scalar_entry_t*)scratch_calloc(
&hdr, (size_t)cap * sizeof(mode_scalar_entry_t));
if (!ht) return -2;
uint64_t mask = cap - 1;
bool has_nulls = ray_vec_may_have_nulls(src);
int64_t best_count = 0, best_row = -1, best_pos = INT64_MAX;
for (int64_t i = 0; i < cnt; i++) {
if ((i & 65535) == 0 && ray_interrupted()) {
*cancelled = 1;
scratch_free(hdr);
return -3;
}
int64_t row = rows[i];
if (has_nulls && ray_vec_is_null(src, row)) continue;
uint64_t key = mode_scalar_key(src, row);
uint64_t slot = ray_hash_i64((int64_t)key) & mask;
for (;;) {
mode_scalar_entry_t* e = &ht[slot];
if (!e->used) {
e->used = 1;
e->key = key;
e->count = 1;
e->first_row = row;
e->first_pos = i;
if (best_count == 0 || e->first_pos < best_pos) {
best_count = 1;
best_row = row;
best_pos = i;
}
break;
}
if (e->key == key) {
e->count++;
if (e->count > best_count ||
(e->count == best_count && e->first_pos < best_pos)) {
best_count = e->count;
best_row = e->first_row;
best_pos = e->first_pos;
}
break;
}
slot = (slot + 1) & mask;
}
}
scratch_free(hdr);
return best_row;
}
static int64_t mode_guid_group(ray_t* src, const int64_t* rows,
int64_t cnt, int* cancelled) {
uint64_t cap = mode_hash_cap(cnt);
if (!cap || cap > SIZE_MAX / sizeof(mode_wide_entry_t)) return -2;
ray_t* hdr = NULL;
mode_wide_entry_t* ht = (mode_wide_entry_t*)scratch_calloc(
&hdr, (size_t)cap * sizeof(mode_wide_entry_t));
if (!ht) return -2;
uint64_t mask = cap - 1;
const uint8_t* base = (const uint8_t*)ray_data(src);
bool has_nulls = ray_vec_may_have_nulls(src);
int64_t best_count = 0, best_row = -1, best_pos = INT64_MAX;
for (int64_t i = 0; i < cnt; i++) {
if ((i & 65535) == 0 && ray_interrupted()) {
*cancelled = 1;
scratch_free(hdr);
return -3;
}
int64_t row = rows[i];
if (has_nulls && ray_vec_is_null(src, row)) continue;
const uint8_t* key = base + (size_t)row * 16;
uint64_t h = ray_hash_bytes(key, 16);
uint64_t slot = h & mask;
for (;;) {
mode_wide_entry_t* e = &ht[slot];
if (!e->used) {
e->used = 1;
e->hash = h;
e->count = 1;
e->first_row = row;
e->first_pos = i;
if (best_count == 0 || e->first_pos < best_pos) {
best_count = 1;
best_row = row;
best_pos = i;
}
break;
}
if (e->hash == h &&
memcmp(key, base + (size_t)e->first_row * 16, 16) == 0) {
e->count++;
if (e->count > best_count ||
(e->count == best_count && e->first_pos < best_pos)) {
best_count = e->count;
best_row = e->first_row;
best_pos = e->first_pos;
}
break;
}
slot = (slot + 1) & mask;
}
}
scratch_free(hdr);
return best_row;
}
static void mode_fixed_per_group_fn(void* ctx_v, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
mode_par_ctx_t* c = (mode_par_ctx_t*)ctx_v;
for (int64_t g = start; g < end; g++) {
if (atomic_load_explicit(&c->oom, memory_order_relaxed) ||
atomic_load_explicit(&c->cancel, memory_order_relaxed))
return;
int64_t off = c->offsets[g];
int64_t cnt = c->grp_cnt[g];
if (cnt >= c->threshold) continue;
int cancelled = 0;
int64_t best = (c->src->type == RAY_GUID)
? mode_guid_group(c->src, &c->idx_buf[off], cnt, &cancelled)
: mode_scalar_group(c->src, &c->idx_buf[off], cnt, &cancelled);
if (cancelled || best == -3) {
atomic_store_explicit(&c->cancel, 1, memory_order_relaxed);
return;
}
if (best == -2) {
atomic_store_explicit(&c->oom, 1, memory_order_relaxed);
return;
}
c->winners[g] = best;
}
}
typedef struct {
ray_t* src;
const int64_t* rows;
const int64_t* offsets;
const int64_t* counts;
int64_t* winners;
int64_t threshold;
_Atomic(bool) oom, cancel;
} mode_str_work_t;
static void mode_str_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; mode_str_work_t* c = raw;
ray_t* src = c->src;
const int64_t* idx_buf = c->rows;
const int64_t* offsets = c->offsets;
const int64_t* grp_cnt = c->counts;
for (int64_t g = start; g < end; g++) {
int64_t cnt = grp_cnt[g];
if (cnt >= c->threshold) continue;
int64_t off = offsets[g];
uint64_t cap = mode_hash_cap(cnt);
if (!cap || cap > SIZE_MAX / sizeof(mode_wide_entry_t)) {
atomic_store(&c->oom, true); return;
}
ray_t* hdr = NULL;
mode_wide_entry_t* ht = (mode_wide_entry_t*)scratch_calloc(
&hdr, (size_t)cap * sizeof(mode_wide_entry_t));
if (!ht) {
atomic_store(&c->oom, true); return;
}
uint64_t mask = cap - 1;
int64_t best_count = 0, best_row = -1, best_pos = INT64_MAX;
for (int64_t i = 0; i < cnt; i++) {
if ((i & 65535) == 0 && ray_interrupted()) {
scratch_free(hdr);
atomic_store(&c->cancel, true); return;
}
int64_t row = idx_buf[off + i];
size_t len = 0;
const char* key = ray_str_vec_get(src, row, &len);
if (!key) { key = ""; len = 0; }
uint64_t h = ray_hash_bytes(key, len);
uint64_t slot = h & mask;
for (;;) {
mode_wide_entry_t* e = &ht[slot];
if (!e->used) {
e->used = 1;
e->hash = h;
e->count = 1;
e->first_row = row;
e->first_pos = i;
if (best_count == 0 || e->first_pos < best_pos) {
best_count = 1;
best_row = row;
best_pos = i;
}
break;
}
size_t olen = 0;
const char* old = ray_str_vec_get(src, e->first_row, &olen);
if (!old) { old = ""; olen = 0; }
if (e->hash == h && olen == len &&
(len == 0 || memcmp(old, key, len) == 0)) {
e->count++;
if (e->count > best_count ||
(e->count == best_count && e->first_pos < best_pos)) {
best_count = e->count;
best_row = e->first_row;
best_pos = e->first_pos;
}
break;
}
slot = (slot + 1) & mask;
}
}
scratch_free(hdr);
c->winners[g] = best_row;
}
}
typedef struct {
uint64_t hash;
int64_t row, position, count;
} mode_record_t;
typedef struct {
ray_t* src;
const int64_t* rows;
int64_t count;
uint32_t tasks, partitions;
bool has_nulls;
mode_record_t *records, *scattered;
int64_t *histogram, *ends;
int64_t emitted[RAY_POOL_INIT_TASKS];
mode_record_t best[RAY_POOL_INIT_TASKS];
_Atomic(bool) oom;
} mode_split_t;
static uint64_t mode_row_hash(ray_t* src, int64_t row) {
if (src->type == RAY_GUID)
return ray_hash_bytes((const char*)ray_data(src) + (size_t)row * 16, 16);
if (src->type == RAY_STR) {
size_t len = 0;
const char* value = ray_str_vec_get(src, row, &len);
return ray_hash_bytes(value ? value : "", value ? len : 0);
}
return ray_hash_i64((int64_t)mode_scalar_key(src, row));
}
static bool mode_rows_equal(ray_t* src, int64_t a, int64_t b) {
if (src->type == RAY_GUID) {
const char* base = ray_data(src);
return memcmp(base + (size_t)a * 16, base + (size_t)b * 16, 16) == 0;
}
if (src->type == RAY_STR) {
size_t na = 0, nb = 0;
const char* va = ray_str_vec_get(src, a, &na);
const char* vb = ray_str_vec_get(src, b, &nb);
if (!va) na = 0;
if (!vb) nb = 0;
return na == nb && (!na || memcmp(va, vb, na) == 0);
}
return mode_scalar_key(src, a) == mode_scalar_key(src, b);
}
static bool mode_better(mode_record_t a, mode_record_t b) {
return a.count > b.count || (a.count == b.count && a.position < b.position);
}
static void mode_emit_record(mode_split_t* c, int64_t task, int64_t begin,
int64_t* emitted, mode_record_t record) {
c->records[begin + (*emitted)++] = record;
c->histogram[task * c->partitions + (record.hash & (c->partitions - 1))]++;
}
static void mode_combine_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; mode_split_t* c = raw;
for (int64_t t = start; t < end; t++) {
int64_t begin = c->count / c->tasks * t;
int64_t limit = t + 1 == c->tasks ? c->count : c->count / c->tasks * (t + 1);
mode_record_t local[1024] = {0};
int64_t emitted = 0;
unsigned used = 0;
for (int64_t i = begin; i < limit; i++) {
if ((i & 65535) == 0 && ray_interrupted()) return;
int64_t row = c->rows[i];
if (c->has_nulls && ray_vec_is_null(c->src, row)) continue;
uint64_t hash = mode_row_hash(c->src, row);
unsigned at = hash & 1023;
while (local[at].count && (local[at].hash != hash ||
!mode_rows_equal(c->src, local[at].row, row))) at = (at + 1) & 1023;
if (local[at].count) local[at].count++;
else {
local[at] = (mode_record_t){hash, row, i, 1};
if (++used == 768) {
for (unsigned j = 0; j < 1024; j++)
if (local[j].count) mode_emit_record(c, t, begin, &emitted, local[j]);
memset(local, 0, sizeof(local)); used = 0;
}
}
}
for (unsigned i = 0; i < 1024; i++)
if (local[i].count) mode_emit_record(c, t, begin, &emitted, local[i]);
c->emitted[t] = emitted;
}
}
static void mode_scatter_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; mode_split_t* c = raw;
for (int64_t t = start; t < end; t++) {
int64_t begin = c->count / c->tasks * t;
int64_t* cursor = c->histogram + t * c->partitions;
for (int64_t i = 0; i < c->emitted[t]; i++) {
if ((i & 65535) == 0 && ray_interrupted()) return;
mode_record_t record = c->records[begin + i];
c->scattered[cursor[record.hash & (c->partitions - 1)]++] = record;
}
}
}
static void mode_count_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; mode_split_t* c = raw;
for (int64_t p = start; p < end; p++) {
int64_t begin = p ? c->ends[p - 1] : 0, limit = c->ends[p];
if (begin == limit) continue;
uint64_t cap = mode_hash_cap(limit - begin);
if (!cap || cap > SIZE_MAX / sizeof(mode_record_t)) { atomic_store(&c->oom, true); return; }
ray_t* hdr = NULL;
mode_record_t* table = scratch_calloc(&hdr, (size_t)cap * sizeof(mode_record_t));
if (!table) { atomic_store(&c->oom, true); return; }
mode_record_t best = {.position = INT64_MAX};
for (int64_t i = begin; i < limit; i++) {
if ((i & 65535) == 0 && ray_interrupted()) { scratch_free(hdr); return; }
mode_record_t record = c->scattered[i];
uint64_t at = ray_hash_i64((int64_t)record.hash) & (cap - 1);
while (table[at].count && (table[at].hash != record.hash ||
!mode_rows_equal(c->src, table[at].row, record.row)))
at = (at + 1) & (cap - 1);
if (!table[at].count) table[at] = record;
else {
table[at].count += record.count;
if (record.position < table[at].position) {
table[at].position = record.position;
table[at].row = record.row;
}
}
if (mode_better(table[at], best)) best = table[at];
}
c->best[p] = best;
scratch_free(hdr);
}
}
static int64_t mode_large_group(ray_t* src, const int64_t* rows, int64_t count) {
ray_pool_t* pool = ray_pool_get();
mode_split_t c = {.src = src, .rows = rows, .count = count,
.tasks = ray_pool_total_workers(pool) * 4,
.has_nulls = src->type != RAY_STR && ray_vec_may_have_nulls(src)};
if (c.tasks > RAY_POOL_INIT_TASKS / 2) c.tasks = RAY_POOL_INIT_TASKS / 2;
c.partitions = 1;
while (c.partitions < c.tasks) c.partitions *= 2;
if ((uint64_t)count > SIZE_MAX / 160 ||
(uint64_t)count * 160 > (uint64_t)ray_heap_anon_watermark() / 4) return -4;
ray_t *records_hdr = NULL, *scatter_hdr = NULL, *hist_hdr = NULL;
c.records = scratch_alloc(&records_hdr, (size_t)count * sizeof(mode_record_t));
c.histogram = scratch_calloc(&hist_hdr,
((size_t)c.tasks * c.partitions + c.partitions) * sizeof(int64_t));
int64_t answer = -2;
if (!c.records || !c.histogram) goto done;
c.ends = c.histogram + (size_t)c.tasks * c.partitions;
ray_pool_dispatch_n(pool, mode_combine_run, &c, c.tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) { answer = -3; goto done; }
int64_t total = 0;
for (uint32_t p = 0; p < c.partitions; p++) {
for (uint32_t t = 0; t < c.tasks; t++) {
int64_t* slot = &c.histogram[(size_t)t * c.partitions + p];
int64_t n = *slot; *slot = total; total += n;
}
c.ends[p] = total;
}
if (!total) { answer = -1; goto done; }
c.scattered = scratch_alloc(&scatter_hdr, (size_t)total * sizeof(mode_record_t));
if (!c.scattered) goto done;
ray_pool_dispatch_n(pool, mode_scatter_run, &c, c.tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) { answer = -3; goto done; }
scratch_free(records_hdr); records_hdr = NULL;
ray_pool_dispatch_n(pool, mode_count_run, &c, c.partitions);
if (ray_interrupted() || atomic_load(&pool->cancelled)) { answer = -3; goto done; }
if (atomic_load(&c.oom)) goto done;
mode_record_t best = {.position = INT64_MAX};
for (uint32_t p = 0; p < c.partitions; p++)
if (c.best[p].count && mode_better(c.best[p], best)) best = c.best[p];
answer = best.count ? best.row : -1;
done:
scratch_free(records_hdr); scratch_free(scatter_hdr); scratch_free(hist_hdr);
return answer;
}
ray_t* ray_mode_per_group_buf(ray_t* src, const int64_t* rows,
const int64_t* offsets, const int64_t* counts, int64_t ng) {
if (!src || RAY_IS_ERR(src) || ng < 0) return NULL;
switch (src->type) {
case RAY_BOOL: case RAY_U8: case RAY_I16: case RAY_I32: case RAY_I64:
case RAY_F32: case RAY_F64: case RAY_DATE: case RAY_TIME: case RAY_TIMESTAMP:
case RAY_GUID: case RAY_SYM: case RAY_STR: break;
default: return NULL;
}
int64_t* winners = ray_alloc_raw((size_t)(ng ? ng : 1) * sizeof(int64_t));
if (!winners) return ray_error("oom", NULL);
ray_pool_t* pool = ray_pool_get();
int64_t total = 0, threshold = INT64_MAX;
for (int64_t g = 0; g < ng; g++) total += counts[g];
if (ray_pool_par_dispatch_ok(pool, total, 262144)) {
threshold = total / ray_pool_total_workers(pool);
if (threshold < 262144) threshold = 262144;
}
mode_str_work_t strings = {.src = src, .rows = rows, .offsets = offsets,
.counts = counts, .winners = winners, .threshold = threshold};
mode_par_ctx_t fixed = {.src = src, .idx_buf = rows, .offsets = offsets,
.grp_cnt = counts, .winners = winners, .threshold = threshold};
if (src->type == RAY_STR) ray_group_dispatch(mode_str_run, &strings, counts, ng);
else ray_group_dispatch(mode_fixed_per_group_fn, &fixed, counts, ng);
for (int64_t g = 0; g < ng; g++) {
if (ray_interrupted() || (pool && atomic_load(&pool->cancelled)) ||
atomic_load(&strings.cancel) || atomic_load(&fixed.cancel) ||
atomic_load(&strings.oom) || atomic_load(&fixed.oom)) break;
if (counts[g] < threshold) continue;
int64_t winner = mode_large_group(src, rows + offsets[g], counts[g]);
if (winner == -4) {
strings.threshold = fixed.threshold = INT64_MAX;
if (src->type == RAY_STR) mode_str_run(&strings, 0, g, g + 1);
else mode_fixed_per_group_fn(&fixed, 0, g, g + 1);
} else if (winner == -2) atomic_store(&fixed.oom, 1);
else if (winner == -3) atomic_store(&fixed.cancel, 1);
else winners[g] = winner;
}
bool cancel = ray_interrupted() || (pool && atomic_load(&pool->cancelled)) ||
atomic_load(&strings.cancel) || atomic_load(&fixed.cancel);
ray_t* out = cancel ? ray_error("cancel", NULL) :
atomic_load(&strings.oom) || atomic_load(&fixed.oom) ? ray_error("oom", NULL) :
ray_group_gather(src, winners, ng);
ray_free_raw(winners);
return out;
}
typedef struct {
const void* base;
int8_t src_type;
bool has_nulls;
int64_t k;
uint8_t desc;
const int64_t* idx_buf;
const int64_t* offsets;
const int64_t* grp_cnt;
ray_t* out_list;
_Atomic(int)* oom;
} topk_par_ctx_t;
static inline double topk_read_f64(const void* base, int64_t row) {
double v; memcpy(&v, (const char*)base + (size_t)row * 8, 8); return v;
}
static inline int64_t topk_read_i64(const void* base, int8_t t, int64_t row) {
switch (t) {
case RAY_I64: case RAY_TIMESTAMP:
{ int64_t v; memcpy(&v, (const char*)base + (size_t)row * 8, 8); return v; }
case RAY_I32: case RAY_DATE: case RAY_TIME:
{ int32_t v; memcpy(&v, (const char*)base + (size_t)row * 4, 4); return (int64_t)v; }
case RAY_I16:
{ int16_t v; memcpy(&v, (const char*)base + (size_t)row * 2, 2); return (int64_t)v; }
case RAY_BOOL: case RAY_U8:
return (int64_t)((const uint8_t*)base)[row];
default: return 0;
}
}
static inline void topk_write_i64(void* dst, int64_t idx, int64_t v, uint8_t esz) {
switch (esz) {
case 1: ((uint8_t*)dst)[idx] = (uint8_t)v; break;
case 2: ((int16_t*)dst)[idx] = (int16_t)v; break;
case 4: ((int32_t*)dst)[idx] = (int32_t)v; break;
default: ((int64_t*)dst)[idx] = v; break;
}
}
static inline void topk_sift_down_dbl(double* h, int64_t n, int64_t i, int max_heap) {
for (;;) {
int64_t l = 2*i+1, r = 2*i+2, w = i;
if (max_heap) {
if (l < n && h[l] > h[w]) w = l;
if (r < n && h[r] > h[w]) w = r;
} else {
if (l < n && h[l] < h[w]) w = l;
if (r < n && h[r] < h[w]) w = r;
}
if (w == i) break;
double t = h[i]; h[i] = h[w]; h[w] = t;
i = w;
}
}
static inline void topk_sift_down_i64(int64_t* h, int64_t n, int64_t i, int max_heap) {
for (;;) {
int64_t l = 2*i+1, r = 2*i+2, w = i;
if (max_heap) {
if (l < n && h[l] > h[w]) w = l;
if (r < n && h[r] > h[w]) w = r;
} else {
if (l < n && h[l] < h[w]) w = l;
if (r < n && h[r] < h[w]) w = r;
}
if (w == i) break;
int64_t t = h[i]; h[i] = h[w]; h[w] = t;
i = w;
}
}
static void topk_per_group_fn(void* ctx_v, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
topk_par_ctx_t* c = (topk_par_ctx_t*)ctx_v;
int8_t t = c->src_type;
int64_t K = c->k;
uint8_t desc = c->desc;
for (int64_t gi = start; gi < end; gi++) {
if (atomic_load_explicit(c->oom, memory_order_relaxed)) break;
ray_t* cell = ray_list_get(c->out_list, gi);
if (!cell) continue;
int64_t cnt = c->grp_cnt[gi];
int64_t off = c->offsets[gi];
const int64_t* idxs = &c->idx_buf[off];
int max_heap = desc ? 0 : 1;
if (t == RAY_F64) {
double* dst = (double*)ray_data(cell);
int64_t kept = 0;
int64_t init_end = 0;
for (int64_t i = 0; i < cnt && kept < K; i++) {
int64_t row = idxs[i];
init_end = i + 1;
if (c->has_nulls && med_is_null(c->base, c->src_type, row)) continue;
dst[kept++] = topk_read_f64(c->base, row);
}
if (kept == K) {
for (int64_t j = K/2 - 1; j >= 0; j--)
topk_sift_down_dbl(dst, K, j, max_heap);
for (int64_t i = init_end; i < cnt; i++) {
int64_t row = idxs[i];
if (c->has_nulls && med_is_null(c->base, c->src_type, row)) continue;
double v = topk_read_f64(c->base, row);
if (desc ? (v > dst[0]) : (v < dst[0])) {
dst[0] = v;
topk_sift_down_dbl(dst, K, 0, max_heap);
}
}
}
if (kept < K)
for (int64_t h = kept / 2; h > 0; h--) topk_sift_down_dbl(dst, kept, h - 1, max_heap);
int64_t n = kept;
while (n > 1) {
double tmp = dst[0]; dst[0] = dst[n-1]; dst[n-1] = tmp;
n--;
topk_sift_down_dbl(dst, n, 0, max_heap);
}
cell->len = kept;
} else {
void* dst = ray_data(cell);
uint8_t esz = ray_sym_elem_size(t, cell->attrs);
int64_t heap_nmax = cnt < K ? cnt : K;
ray_t* heap_hdr = NULL;
int64_t* heap = NULL;
if (heap_nmax > 0) {
if ((uint64_t)heap_nmax > SIZE_MAX / sizeof(int64_t)) {
atomic_store_explicit(c->oom, 1, memory_order_relaxed);
break;
}
heap = (int64_t*)scratch_alloc(&heap_hdr,
(size_t)heap_nmax * sizeof(int64_t));
if (!heap) {
atomic_store_explicit(c->oom, 1, memory_order_relaxed);
break;
}
}
int64_t kept = 0;
int64_t init_end = 0;
for (int64_t i = 0; i < cnt && kept < K; i++) {
int64_t row = idxs[i];
init_end = i + 1;
if (c->has_nulls && med_is_null(c->base, c->src_type, row)) continue;
heap[kept++] = topk_read_i64(c->base, t, row);
}
if (kept == K) {
for (int64_t j = K/2 - 1; j >= 0; j--)
topk_sift_down_i64(heap, K, j, max_heap);
for (int64_t i = init_end; i < cnt; i++) {
int64_t row = idxs[i];
if (c->has_nulls && med_is_null(c->base, c->src_type, row)) continue;
int64_t v = topk_read_i64(c->base, t, row);
if (desc ? (v > heap[0]) : (v < heap[0])) {
heap[0] = v;
topk_sift_down_i64(heap, K, 0, max_heap);
}
}
}
if (kept < K)
for (int64_t h = kept / 2; h > 0; h--) topk_sift_down_i64(heap, kept, h - 1, max_heap);
int64_t n = kept;
while (n > 1) {
int64_t tmp = heap[0]; heap[0] = heap[n-1]; heap[n-1] = tmp;
n--;
topk_sift_down_i64(heap, n, 0, max_heap);
}
for (int64_t i = 0; i < kept; i++)
topk_write_i64(dst, i, heap[i], esz);
cell->len = kept;
scratch_free(heap_hdr);
}
}
}
static int topk_row_cmp(ray_t* src, int64_t a, int64_t b, const ray_group_sym_view_t* view) {
const void* data = ray_data(src);
if (src->type == RAY_F32) {
float x = ((const float*)data)[a], y = ((const float*)data)[b];
return (x > y) - (x < y);
}
if (src->type == RAY_F64) {
double x = ((const double*)data)[a], y = ((const double*)data)[b];
return (x > y) - (x < y);
}
if (src->type != RAY_GUID && src->type != RAY_SYM && src->type != RAY_STR) {
int64_t x = topk_read_i64(data, src->type, a), y = topk_read_i64(data, src->type, b);
return (x > y) - (x < y);
}
if (src->type == RAY_GUID)
return memcmp((const char*)data + (size_t)a * 16, (const char*)data + (size_t)b * 16, 16);
if (src->type == RAY_SYM) {
ray_t* x = ray_group_sym_read(view, ray_sym_vec_domain(src), ray_read_sym(data, a, src->type, src->attrs));
ray_t* y = ray_group_sym_read(view, ray_sym_vec_domain(src), ray_read_sym(data, b, src->type, src->attrs));
return ray_str_cmp(x, y);
}
size_t na = 0, nb = 0;
const char* x = ray_str_vec_get(src, a, &na);
const char* y = ray_str_vec_get(src, b, &nb);
size_t common = na < nb ? na : nb;
int cmp = common ? memcmp(x, y, common) : 0;
return cmp ? cmp : (na > nb) - (na < nb);
}
static void topk_rows_sift(ray_t* src, int64_t* heap, int64_t n, int64_t root, bool desc, const ray_group_sym_view_t* view) {
#define TOPK_SIFT_NATIVE(TYPE) do { \
const TYPE* data = ray_data(src); \
for (;;) { \
int64_t child = root * 2 + 1; \
if (child >= n) return; \
if (child + 1 < n && (desc ? data[heap[child + 1]] < data[heap[child]] : \
data[heap[child + 1]] > data[heap[child]])) child++; \
if (!(desc ? data[heap[child]] < data[heap[root]] : data[heap[child]] > data[heap[root]])) return; \
int64_t swap = heap[root]; heap[root] = heap[child]; heap[child] = swap; \
root = child; \
} \
} while (0)
switch (src->type) {
case RAY_BOOL: case RAY_U8: TOPK_SIFT_NATIVE(uint8_t);
case RAY_I16: TOPK_SIFT_NATIVE(int16_t);
case RAY_I32: case RAY_DATE: case RAY_TIME: TOPK_SIFT_NATIVE(int32_t);
case RAY_I64: case RAY_TIMESTAMP: TOPK_SIFT_NATIVE(int64_t);
case RAY_F32: TOPK_SIFT_NATIVE(float);
case RAY_F64: TOPK_SIFT_NATIVE(double);
default: break;
}
#undef TOPK_SIFT_NATIVE
for (;;) {
int64_t child = root * 2 + 1;
if (child >= n) return;
if (child + 1 < n) {
int cmp = topk_row_cmp(src, heap[child + 1], heap[child], view);
if (desc ? cmp < 0 : cmp > 0) child++;
}
int cmp = topk_row_cmp(src, heap[child], heap[root], view);
if (!(desc ? cmp < 0 : cmp > 0)) return;
int64_t tmp = heap[root]; heap[root] = heap[child]; heap[child] = tmp;
root = child;
}
}
typedef struct {
ray_t* src;
int64_t k;
bool desc;
const int64_t* rows;
const int64_t* offsets;
const int64_t* counts;
ray_t** cells;
_Atomic(bool) oom;
_Atomic(ray_t*) error;
ray_group_sym_view_t symbols;
} topk_wide_work_t;
static void topk_wide_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; topk_wide_work_t* c = raw;
ray_t* src = c->src;
const int64_t* rows = c->rows;
const int64_t* offsets = c->offsets;
const int64_t* counts = c->counts;
bool desc = c->desc;
int64_t max_count = 0;
for (int64_t g = start; g < end; g++) if (counts[g] > max_count) max_count = counts[g];
int64_t cap = c->k < max_count ? c->k : max_count;
int64_t* heap = ray_alloc_raw((size_t)(cap > 0 ? cap : 1) * sizeof(int64_t));
if (!heap) { atomic_store(&c->oom, true); return; }
for (int64_t g = start; g < end; g++) {
int64_t kept = 0;
for (int64_t j = 0; j < counts[g]; j++) {
int64_t row = rows[offsets[g] + j];
if (ray_vec_is_null(src, row)) continue;
if (kept < cap) {
heap[kept++] = row;
if (kept == cap)
for (int64_t h = kept / 2; h > 0; h--) topk_rows_sift(src, heap, kept, h - 1, desc, &c->symbols);
} else {
int cmp = topk_row_cmp(src, row, heap[0], &c->symbols);
if (desc ? cmp > 0 : cmp < 0) { heap[0] = row; topk_rows_sift(src, heap, kept, 0, desc, &c->symbols); }
}
}
if (kept < cap)
for (int64_t h = kept / 2; h > 0; h--) topk_rows_sift(src, heap, kept, h - 1, desc, &c->symbols);
for (int64_t n = kept; n > 1; n--) {
int64_t tmp = heap[0]; heap[0] = heap[n - 1]; heap[n - 1] = tmp;
topk_rows_sift(src, heap, n - 1, 0, desc, &c->symbols);
}
ray_t* cell = gather_by_idx(src, heap, kept);
if (!cell || RAY_IS_ERR(cell)) {
if (cell) {
ray_t* expected = NULL;
if (!atomic_compare_exchange_strong(&c->error, &expected, cell)) ray_release(cell);
}
atomic_store(&c->oom, true); break;
}
c->cells[g] = cell;
}
ray_free_raw(heap);
}
static ray_t* topk_wide_per_group_buf(ray_t* src, int64_t k, bool desc,
const int64_t* rows, const int64_t* offsets, const int64_t* counts, int64_t ng) {
ray_t* out = ray_list_new(ng);
if (!out || RAY_IS_ERR(out)) return out ? out : ray_error("oom", NULL);
out->len = ng;
memset(ray_data(out), 0, (size_t)ng * sizeof(ray_t*));
topk_wide_work_t work = {.src = src, .k = k, .desc = desc, .rows = rows,
.offsets = offsets, .counts = counts, .cells = ray_data(out), .oom = false, .error = NULL};
if (src->type == RAY_SYM) ray_sym_strings_borrow(&work.symbols.strings, &work.symbols.count);
ray_group_dispatch(topk_wide_run, &work, counts, ng);
ray_pool_t* pool = ray_pool_get();
bool cancel = pool && atomic_load(&pool->cancelled);
ray_t* error = atomic_load(&work.error);
if (cancel || atomic_load(&work.oom)) {
ray_release(out);
if (error) return error;
return ray_error(cancel ? "cancel" : "oom", NULL);
}
return out;
}
static ray_t* topk_regular_per_group_buf(ray_t* src,
int64_t k,
uint8_t desc,
const int64_t* idx_buf,
const int64_t* offsets,
const int64_t* grp_cnt,
int64_t n_groups) {
if (!src || RAY_IS_ERR(src) || n_groups < 0) return NULL;
if (k < 1) return NULL;
int8_t t = src->type;
if (t == RAY_F32 || t == RAY_GUID || t == RAY_SYM || t == RAY_STR)
return topk_wide_per_group_buf(src, k, desc, idx_buf, offsets, grp_cnt, n_groups);
if (t != RAY_F64 && t != RAY_I64 && t != RAY_I32 && t != RAY_I16 &&
t != RAY_U8 && t != RAY_BOOL && t != RAY_DATE && t != RAY_TIME &&
t != RAY_TIMESTAMP)
return NULL;
ray_t* out = ray_list_new(n_groups);
if (!out || RAY_IS_ERR(out)) return out ? out : ray_error("oom", NULL);
for (int64_t gi = 0; gi < n_groups; gi++) {
int64_t kk = grp_cnt[gi] < k ? grp_cnt[gi] : k;
ray_t* cell = col_vec_new(src, kk);
if (!cell || RAY_IS_ERR(cell)) {
ray_release(out);
return cell ? cell : ray_error("oom", NULL);
}
cell->len = 0;
ray_t* new_out = ray_list_append(out, cell);
ray_release(cell);
if (!new_out || RAY_IS_ERR(new_out)) {
ray_release(out);
return new_out ? new_out : ray_error("oom", NULL);
}
out = new_out;
}
_Atomic(int) oom = 0;
topk_par_ctx_t ctx = {
.base = ray_data(src),
.src_type = t,
.has_nulls = ray_vec_may_have_nulls(src),
.k = k,
.desc = desc,
.idx_buf = idx_buf,
.offsets = offsets,
.grp_cnt = grp_cnt,
.out_list = out,
.oom = &oom,
};
ray_group_dispatch(topk_per_group_fn, &ctx, grp_cnt, n_groups);
if (atomic_load_explicit(&oom, memory_order_relaxed)) {
ray_release(out);
return ray_error("oom", NULL);
}
return out;
}
typedef struct {
ray_t* source;
const int64_t* rows;
int64_t count, k, heap_stride;
uint32_t tasks, source_tasks;
bool desc, sorted;
int64_t* heaps;
int64_t kept[RAY_POOL_INIT_TASKS];
ray_group_sym_view_t symbols;
} topk_chunks_t;
static void topk_keep_row(topk_chunks_t* c, int64_t* heap, int64_t* kept, int64_t row) {
if (*kept < c->k) {
heap[(*kept)++] = row;
if (*kept == c->k)
for (int64_t h = *kept / 2; h > 0; h--)
topk_rows_sift(c->source, heap, *kept, h - 1, c->desc, &c->symbols);
} else {
int cmp = topk_row_cmp(c->source, row, heap[0], &c->symbols);
if (c->desc ? cmp > 0 : cmp < 0) {
heap[0] = row;
topk_rows_sift(c->source, heap, *kept, 0, c->desc, &c->symbols);
}
}
}
static void topk_sort_chunk(topk_chunks_t* c, int64_t t, int64_t kept) {
int64_t* heap = c->heaps + t * (c->heap_stride ? c->heap_stride : c->k);
if (kept < c->k)
for (int64_t h = kept / 2; h > 0; h--)
topk_rows_sift(c->source, heap, kept, h - 1, c->desc, &c->symbols);
for (int64_t n = kept; n > 1; n--) {
int64_t swap = heap[0]; heap[0] = heap[n - 1]; heap[n - 1] = swap;
topk_rows_sift(c->source, heap, n - 1, 0, c->desc, &c->symbols);
}
}
static void topk_sort_chunks_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; topk_chunks_t* c = raw;
for (int64_t t = start; t < end; t++) topk_sort_chunk(c, t, c->kept[t]);
}
static void topk_chunks_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
topk_chunks_t* c = raw;
uint32_t sources = c->source_tasks ? c->source_tasks : c->tasks;
for (int64_t t = start; t < end; t++) {
int64_t remainder = c->count % sources;
int64_t begin = c->count / sources * t + (t < remainder ? t : remainder);
int64_t limit = begin + c->count / sources + (t < remainder);
int64_t heap_id = c->source_tasks ? wid : t;
int64_t kept = c->source_tasks ? c->kept[heap_id] : 0;
int64_t* heap = c->heaps + heap_id * (c->heap_stride ? c->heap_stride : c->k);
bool has_nulls = ray_vec_may_have_nulls(c->source);
#define TOPK_SCAN_NATIVE(TYPE, VALID) do { \
const TYPE* data = ray_data(c->source); \
for (int64_t i = begin; i < limit; i++) { \
if ((i & 65535) == 0 && ray_interrupted()) return; \
int64_t row = c->rows[i]; \
TYPE value = data[row]; \
if (has_nulls && !(VALID)) continue; \
if (kept < c->k) { \
topk_keep_row(c, heap, &kept, row); \
} else if (c->desc ? value > data[heap[0]] : value < data[heap[0]]) { \
heap[0] = row; \
int64_t root = 0; \
for (;;) { \
int64_t child = root * 2 + 1; \
if (child >= kept) break; \
if (child + 1 < kept && (c->desc ? data[heap[child + 1]] < data[heap[child]] : \
data[heap[child + 1]] > data[heap[child]])) child++; \
if (!(c->desc ? data[heap[child]] < data[heap[root]] : \
data[heap[child]] > data[heap[root]])) break; \
int64_t swap = heap[root]; heap[root] = heap[child]; heap[child] = swap; \
root = child; \
} \
} \
} \
} while (0)
switch (c->source->type) {
case RAY_BOOL: case RAY_U8: TOPK_SCAN_NATIVE(uint8_t, true); break;
case RAY_I16: TOPK_SCAN_NATIVE(int16_t, value != NULL_I16); break;
case RAY_I32: case RAY_DATE: case RAY_TIME:
TOPK_SCAN_NATIVE(int32_t, value != NULL_I32); break;
case RAY_I64: case RAY_TIMESTAMP:
TOPK_SCAN_NATIVE(int64_t, value != NULL_I64); break;
case RAY_F32: TOPK_SCAN_NATIVE(float, value == value); break;
case RAY_F64: TOPK_SCAN_NATIVE(double, value == value); break;
default:
for (int64_t i = begin; i < limit; i++) {
if ((i & 65535) == 0 && ray_interrupted()) return;
if (!ray_vec_is_null(c->source, c->rows[i]))
topk_keep_row(c, heap, &kept, c->rows[i]);
}
break;
}
#undef TOPK_SCAN_NATIVE
if (c->sorted && !c->source_tasks) topk_sort_chunk(c, heap_id, kept);
c->kept[heap_id] = kept;
}
}
typedef struct {
const topk_chunks_t* values;
const int64_t* input;
int64_t* output;
struct {
int64_t left, nl, right, nr, output, begin, end;
} job[RAY_POOL_INIT_TASKS];
} topk_merge_t;
static bool topk_before(const topk_chunks_t* c, int64_t a, int64_t b, bool equal) {
int cmp = topk_row_cmp(c->source, a, b, &c->symbols);
return c->desc ? (equal ? cmp >= 0 : cmp > 0) : (equal ? cmp <= 0 : cmp < 0);
}
static void topk_merge_run(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; topk_merge_t* c = raw;
for (int64_t t = start; t < end; t++) {
const int64_t* left = c->input + c->job[t].left;
const int64_t* right = c->input + c->job[t].right;
int64_t nl = c->job[t].nl, nr = c->job[t].nr, diagonal = c->job[t].begin;
int64_t low = diagonal > nr ? diagonal - nr : 0;
int64_t high = diagonal < nl ? diagonal : nl, i = 0, j = 0;
while (low <= high) {
i = low + (high - low) / 2; j = diagonal - i;
if (i > 0 && j < nr && !topk_before(c->values, left[i - 1], right[j], true)) high = i - 1;
else if (j > 0 && i < nl && topk_before(c->values, left[i], right[j - 1], true)) low = i + 1;
else break;
}
for (int64_t pos = diagonal; pos < c->job[t].end; pos++) {
if ((pos & 65535) == 0 && ray_interrupted()) return;
bool take_left = i < nl && (j == nr || topk_before(c->values, left[i], right[j], true));
c->output[c->job[t].output + pos] = take_left ? left[i++] : right[j++];
}
}
}
typedef struct {
const topk_chunks_t* values;
int64_t *input, *output;
int64_t base, n, pivot;
struct { int64_t begin, end, less, equal, lo, eq, hi; } task[RAY_POOL_INIT_TASKS];
} topk_select_t;
static void topk_select_ranges(topk_select_t* c) {
int64_t tasks = c->values->tasks, remainder = c->n % tasks;
for (int64_t t = 0; t < tasks; t++) {
c->task[t].begin = c->base + c->n / tasks * t + (t < remainder ? t : remainder);
c->task[t].end = c->task[t].begin + c->n / tasks + (t < remainder);
}
}
static void topk_select_gather(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; topk_select_t* c = raw;
for (int64_t t = start; t < end; t++) {
int64_t at = c->task[t].begin;
for (int64_t i = at; i < c->task[t].end; i++) {
if ((i & 65535) == 0 && ray_interrupted()) return;
int64_t row = c->values->rows[i];
if (!ray_vec_is_null(c->values->source, row)) c->input[at++] = row;
}
c->task[t].less = at - c->task[t].begin;
}
}
static void topk_select_compact(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; topk_select_t* c = raw;
for (int64_t t = start; t < end; t++)
memcpy(c->output + c->task[t].lo, c->input + c->task[t].begin,
(size_t)c->task[t].less * sizeof(int64_t));
}
static void topk_select_count(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; topk_select_t* c = raw;
ray_t* src = c->values->source;
#define TOPK_SELECT_COUNT(TYPE) do { \
const TYPE* data = ray_data(src); \
TYPE pivot = data[c->pivot]; \
for (int64_t t = start; t < end; t++) { \
int64_t less = 0, equal = 0; \
for (int64_t i = c->task[t].begin; i < c->task[t].end; i++) { \
if ((i & 65535) == 0 && ray_interrupted()) return; \
TYPE value = data[c->input[i]]; \
int cmp = (value > pivot) - (value < pivot); \
if (!cmp) equal++; \
else if (c->values->desc ? cmp > 0 : cmp < 0) less++; \
} \
c->task[t].less = less; c->task[t].equal = equal; \
} \
return; \
} while (0)
switch (src->type) {
case RAY_BOOL: case RAY_U8: TOPK_SELECT_COUNT(uint8_t);
case RAY_I16: TOPK_SELECT_COUNT(int16_t);
case RAY_I32: case RAY_DATE: case RAY_TIME: TOPK_SELECT_COUNT(int32_t);
case RAY_I64: case RAY_TIMESTAMP: TOPK_SELECT_COUNT(int64_t);
case RAY_F32: TOPK_SELECT_COUNT(float);
case RAY_F64: TOPK_SELECT_COUNT(double);
default: break;
}
#undef TOPK_SELECT_COUNT
for (int64_t t = start; t < end; t++) {
int64_t less = 0, equal = 0;
for (int64_t i = c->task[t].begin; i < c->task[t].end; i++) {
if ((i & 65535) == 0 && ray_interrupted()) return;
int cmp = topk_row_cmp(c->values->source, c->input[i], c->pivot, &c->values->symbols);
if (!cmp) equal++;
else if (c->values->desc ? cmp > 0 : cmp < 0) less++;
}
c->task[t].less = less; c->task[t].equal = equal;
}
}
static void topk_select_scatter(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; topk_select_t* c = raw;
ray_t* src = c->values->source;
#define TOPK_SELECT_SCATTER(TYPE) do { \
const TYPE* data = ray_data(src); \
TYPE pivot = data[c->pivot]; \
for (int64_t t = start; t < end; t++) { \
int64_t lo = c->task[t].lo, eq = c->task[t].eq, hi = c->task[t].hi; \
\
int64_t begin = c->base / c->values->tasks * t; \
int64_t limit = t + 1 == c->values->tasks ? c->base : c->base / c->values->tasks * (t + 1); \
memcpy(c->output + begin, c->input + begin, (size_t)(limit - begin) * sizeof(int64_t)); \
for (int64_t i = c->task[t].begin; i < c->task[t].end; i++) { \
if ((i & 65535) == 0 && ray_interrupted()) return; \
int64_t row = c->input[i]; \
TYPE value = data[row]; \
int cmp = (value > pivot) - (value < pivot); \
if (!cmp) c->output[eq++] = row; \
else if (c->values->desc ? cmp > 0 : cmp < 0) c->output[lo++] = row; \
else c->output[hi++] = row; \
} \
} \
return; \
} while (0)
switch (src->type) {
case RAY_BOOL: case RAY_U8: TOPK_SELECT_SCATTER(uint8_t);
case RAY_I16: TOPK_SELECT_SCATTER(int16_t);
case RAY_I32: case RAY_DATE: case RAY_TIME: TOPK_SELECT_SCATTER(int32_t);
case RAY_I64: case RAY_TIMESTAMP: TOPK_SELECT_SCATTER(int64_t);
case RAY_F32: TOPK_SELECT_SCATTER(float);
case RAY_F64: TOPK_SELECT_SCATTER(double);
default: break;
}
#undef TOPK_SELECT_SCATTER
for (int64_t t = start; t < end; t++) {
int64_t lo = c->task[t].lo, eq = c->task[t].eq, hi = c->task[t].hi;
int64_t begin = c->base / c->values->tasks * t;
int64_t limit = t + 1 == c->values->tasks ? c->base : c->base / c->values->tasks * (t + 1);
memcpy(c->output + begin, c->input + begin, (size_t)(limit - begin) * sizeof(int64_t));
for (int64_t i = c->task[t].begin; i < c->task[t].end; i++) {
if ((i & 65535) == 0 && ray_interrupted()) return;
int64_t row = c->input[i];
int cmp = topk_row_cmp(c->values->source, row, c->pivot, &c->values->symbols);
if (!cmp) c->output[eq++] = row;
else if (c->values->desc ? cmp > 0 : cmp < 0) c->output[lo++] = row;
else c->output[hi++] = row;
}
}
}
static int64_t topk_select_rows(const topk_chunks_t* values, int64_t** input,
int64_t** scratch, int64_t k) {
topk_select_t c = {.values = values, .input = *input, .output = *scratch, .n = values->count};
ray_pool_t* pool = ray_pool_get();
topk_select_ranges(&c);
ray_pool_dispatch_n(pool, topk_select_gather, &c, values->tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) return -1;
int64_t valid = 0;
for (uint32_t t = 0; t < values->tasks; t++) { c.task[t].lo = valid; valid += c.task[t].less; }
if (valid != c.n) {
ray_pool_dispatch_n(pool, topk_select_compact, &c, values->tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) return -1;
int64_t* swap = c.input; c.input = c.output; c.output = swap;
}
c.n = valid;
if (k > valid) k = valid;
unsigned depth = 0;
while (k > 0 && c.base + c.n > k) {
if (ray_interrupted() || atomic_load(&pool->cancelled)) return -1;
int64_t sample[33], step = c.n / 33;
unsigned ns = step ? 33 : (unsigned)c.n;
for (unsigned i = 0; i < ns; i++) {
sample[i] = c.input[c.base + (step ? i * step + ray_hash_i64(i + depth * 33) % step : i)];
for (unsigned j = i; j && topk_before(values, sample[j], sample[j - 1], false); j--) {
int64_t swap = sample[j]; sample[j] = sample[j - 1]; sample[j - 1] = swap;
}
}
c.pivot = sample[ns / 2];
int64_t less = 0, equal = 0;
if (c.n < 65536 || depth >= 32) {
if (depth >= 32) {
int64_t* heap = c.input + c.base;
for (int64_t h = c.n / 2; h > 0; h--) topk_rows_sift(values->source, heap, c.n, h - 1, values->desc, &values->symbols);
for (int64_t n = c.n; n > 1; n--) {
int64_t swap = heap[0]; heap[0] = heap[n - 1]; heap[n - 1] = swap;
topk_rows_sift(values->source, heap, n - 1, 0, values->desc, &values->symbols);
}
break;
}
int64_t lo = c.base, at = lo, hi = c.base + c.n;
while (at < hi) {
int cmp = topk_row_cmp(values->source, c.input[at], c.pivot, &values->symbols);
if (!cmp) at++;
else if (values->desc ? cmp > 0 : cmp < 0) {
int64_t swap = c.input[lo]; c.input[lo++] = c.input[at]; c.input[at++] = swap;
} else {
int64_t swap = c.input[--hi]; c.input[hi] = c.input[at]; c.input[at] = swap;
}
}
less = lo - c.base; equal = hi - lo;
} else {
topk_select_ranges(&c);
ray_pool_dispatch_n(pool, topk_select_count, &c, values->tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) return -1;
for (uint32_t t = 0; t < values->tasks; t++) { less += c.task[t].less; equal += c.task[t].equal; }
int64_t lo = c.base, eq = lo + less, hi = eq + equal;
for (uint32_t t = 0; t < values->tasks; t++) {
c.task[t].lo = lo; lo += c.task[t].less;
c.task[t].eq = eq; eq += c.task[t].equal;
c.task[t].hi = hi; hi += c.task[t].end - c.task[t].begin - c.task[t].less - c.task[t].equal;
}
ray_pool_dispatch_n(pool, topk_select_scatter, &c, values->tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) return -1;
int64_t* swap = c.input; c.input = c.output; c.output = swap;
}
depth++;
if (k <= c.base + less) c.n = less;
else if (k <= c.base + less + equal) break;
else { c.base += less + equal; c.n -= less + equal; }
}
*input = c.input; *scratch = c.output;
return k;
}
static ray_t* topk_sorted_large_group(ray_t* src, const int64_t* rows, int64_t count,
int64_t k, bool desc, uint32_t tasks) {
unsigned levels = 64 - (unsigned)__builtin_clzll((uint64_t)k);
bool select = k > count / tasks / levels / 2;
if (select) {
uint32_t limit = tasks * 4;
if (limit > RAY_POOL_INIT_TASKS / 2) limit = RAY_POOL_INIT_TASKS / 2;
while (tasks < limit && k / tasks > 262144) tasks *= 2;
if (tasks > limit) tasks = limit;
}
int64_t cap = (count + tasks - 1) / tasks;
if (cap > k) cap = k;
if ((uint64_t)cap > SIZE_MAX / tasks / sizeof(int64_t) ||
(uint64_t)count > SIZE_MAX / sizeof(int64_t) - tasks) return ray_error("oom", NULL);
size_t bytes = (select ? (size_t)count + tasks : (size_t)tasks * (size_t)cap) * sizeof(int64_t);
topk_chunks_t c = {.source = src, .rows = rows, .count = count, .k = cap,
.tasks = tasks, .desc = desc, .sorted = true};
int64_t* first = ray_alloc_raw(bytes);
int64_t* second = ray_alloc_raw(bytes);
if (!first || !second) { ray_free_raw(first); ray_free_raw(second); return ray_error("oom", NULL); }
c.heaps = first;
if (src->type == RAY_SYM) ray_sym_strings_borrow(&c.symbols.strings, &c.symbols.count);
ray_pool_t* pool = ray_pool_get();
ray_t* out = NULL;
int64_t *input = first, *output = second;
if (select) {
int64_t kept = topk_select_rows(&c, &input, &output, k);
if (kept < 0) goto cancelled;
ray_profile_tick("topK: selected rows");
c.rows = input; c.count = kept;
cap = (kept + tasks - 1) / tasks;
if (cap < 1) cap = 1;
c.k = cap; c.heaps = output;
ray_pool_dispatch_n(pool, topk_chunks_run, &c, tasks);
int64_t* swap = input; input = output; output = swap;
} else {
if (tasks == ray_pool_total_workers(pool)) {
c.source_tasks = tasks * 4;
if (c.source_tasks > RAY_POOL_INIT_TASKS / 2) c.source_tasks = RAY_POOL_INIT_TASKS / 2;
}
ray_pool_dispatch_n(pool, topk_chunks_run, &c, c.source_tasks ? c.source_tasks : tasks);
if (ray_interrupted() || atomic_load(&pool->cancelled)) goto cancelled;
ray_profile_tick("topK: scanned chunks");
if (c.source_tasks) ray_pool_dispatch_n(pool, topk_sort_chunks_run, &c, tasks);
}
ray_profile_tick("topK: sorted chunks");
int64_t offsets[RAY_POOL_INIT_TASKS], next_offsets[RAY_POOL_INIT_TASKS], next_counts[RAY_POOL_INIT_TASKS];
for (uint32_t t = 0; t < tasks; t++) offsets[t] = t * cap;
topk_merge_t merge = {.values = &c, .input = input, .output = output};
uint32_t runs = tasks;
while (runs > 1) {
if (ray_interrupted() || atomic_load(&pool->cancelled)) goto cancelled;
int64_t total = 0;
uint32_t pairs = (runs + 1) / 2;
for (uint32_t p = 0; p < pairs; p++) {
int64_t length = c.kept[p * 2] + (p * 2 + 1 < runs ? c.kept[p * 2 + 1] : 0);
next_offsets[p] = total; next_counts[p] = length < k ? length : k;
total += next_counts[p];
}
int64_t grain = (total + tasks - 1) / tasks;
if (grain < 1) grain = 1;
uint32_t jobs = 0;
for (uint32_t p = 0; p < pairs; p++) {
uint32_t l = p * 2, r = l + 1;
for (int64_t pos = 0; pos < next_counts[p]; pos += grain) {
merge.job[jobs].left = offsets[l]; merge.job[jobs].nl = c.kept[l];
merge.job[jobs].right = r < runs ? offsets[r] : 0;
merge.job[jobs].nr = r < runs ? c.kept[r] : 0;
merge.job[jobs].output = next_offsets[p]; merge.job[jobs].begin = pos;
merge.job[jobs++].end = pos + grain < next_counts[p] ? pos + grain : next_counts[p];
}
}
if (jobs) ray_pool_dispatch_n(pool, topk_merge_run, &merge, jobs);
ray_profile_tick("topK: merged level");
const int64_t* swap = merge.input; merge.input = merge.output; merge.output = (int64_t*)swap;
memcpy(offsets, next_offsets, pairs * sizeof(int64_t));
memcpy(c.kept, next_counts, pairs * sizeof(int64_t));
runs = pairs;
}
if (ray_interrupted() || atomic_load(&pool->cancelled)) goto cancelled;
out = ray_group_gather(src, merge.input, c.kept[0]);
ray_profile_tick("topK: gathered output");
goto done;
cancelled:
out = ray_error("cancel", NULL);
done:
ray_free_raw(first); ray_free_raw(second);
return out;
}
static ray_t* topk_large_group(ray_t* src, const int64_t* rows, int64_t count,
int64_t k, bool desc, uint32_t tasks) {
if (k >= 64 || k > count / tasks / 2)
return topk_sorted_large_group(src, rows, count, k, desc, tasks);
topk_chunks_t c = {.source = src, .rows = rows, .count = count, .k = k,
.tasks = tasks, .desc = desc, .heap_stride = ((k + 7) & ~INT64_C(7)) + 8};
c.heaps = ray_alloc_raw((size_t)(tasks + 1) * (size_t)c.heap_stride * sizeof(int64_t));
if (!c.heaps) return ray_error("oom", NULL);
if (src->type == RAY_SYM) ray_sym_strings_borrow(&c.symbols.strings, &c.symbols.count);
ray_pool_t* pool = ray_pool_get();
if (tasks == ray_pool_total_workers(pool)) {
c.source_tasks = tasks * 4;
if (c.source_tasks > RAY_POOL_INIT_TASKS / 2) c.source_tasks = RAY_POOL_INIT_TASKS / 2;
}
ray_pool_dispatch_n(pool, topk_chunks_run, &c, c.source_tasks ? c.source_tasks : tasks);
ray_profile_tick("topK: selected chunk candidates");
if (ray_interrupted() || atomic_load(&pool->cancelled)) {
ray_free_raw(c.heaps); return ray_error("cancel", NULL);
}
int64_t* heap = c.heaps + tasks * c.heap_stride;
int64_t kept = 0;
for (uint32_t t = 0; t < tasks; t++)
for (int64_t i = 0; i < c.kept[t]; i++) topk_keep_row(&c, heap, &kept, c.heaps[t * c.heap_stride + i]);
if (kept < k)
for (int64_t h = kept / 2; h > 0; h--) topk_rows_sift(src, heap, kept, h - 1, desc, &c.symbols);
for (int64_t n = kept; n > 1; n--) {
int64_t swap = heap[0]; heap[0] = heap[n - 1]; heap[n - 1] = swap;
topk_rows_sift(src, heap, n - 1, 0, desc, &c.symbols);
}
ray_t* out = ray_group_gather(src, heap, kept);
ray_free_raw(c.heaps);
return out;
}
ray_t* ray_topk_per_group_buf(ray_t* src, int64_t k, uint8_t desc,
const int64_t* rows, const int64_t* offsets, const int64_t* counts, int64_t groups) {
if (!src || RAY_IS_ERR(src) || groups < 0 || k < 1) return NULL;
int64_t total = 0;
for (int64_t g = 0; g < groups; g++) total += counts[g];
ray_pool_t* pool = ray_pool_get();
if (ray_interrupted() || (pool && atomic_load(&pool->cancelled))) return ray_error("cancel", NULL);
if (!ray_pool_par_dispatch_ok(pool, total, 262144))
return topk_regular_per_group_buf(src, k, desc, rows, offsets, counts, groups);
uint32_t tasks = ray_pool_total_workers(pool);
if (tasks > RAY_POOL_INIT_TASKS / 2) tasks = RAY_POOL_INIT_TASKS / 2;
int64_t threshold = total / ray_pool_total_workers(pool);
if (threshold < 262144) threshold = 262144;
int64_t* regular = NULL;
for (int64_t g = 0; g < groups; g++) if (counts[g] >= threshold) {
if (!regular) {
regular = ray_alloc_raw((size_t)groups * sizeof(int64_t));
if (!regular) return ray_error("oom", NULL);
memcpy(regular, counts, (size_t)groups * sizeof(int64_t));
}
regular[g] = 0;
}
if (!regular) return topk_regular_per_group_buf(src, k, desc, rows, offsets, counts, groups);
ray_t* out = topk_regular_per_group_buf(src, k, desc, rows, offsets, regular, groups);
if (out && !RAY_IS_ERR(out)) {
ray_t** cells = ray_data(out);
for (int64_t g = 0; g < groups; g++) if (counts[g] && !regular[g]) {
ray_t* cell = topk_large_group(src, rows + offsets[g], counts[g], k, desc, tasks);
if (!cell || RAY_IS_ERR(cell)) { ray_release(out); out = cell; break; }
ray_release(cells[g]); cells[g] = cell;
}
}
ray_free_raw(regular);
return out;
}
typedef struct {
ray_t* src;
uint16_t op;
bool has_nulls;
} wide_winners_t;
static int64_t wide_winner_rows(void* raw, const int64_t* rows, int64_t count) {
wide_winners_t* c = raw;
return wide_winner_row(c->src, c->op, rows, count, c->has_nulls);
}
ray_t* ray_wide_minmax_per_group_buf(ray_t* src, uint16_t op,
const int64_t* idx_buf,
const int64_t* offsets,
const int64_t* grp_cnt,
int64_t n_groups) {
if (!src || RAY_IS_ERR(src) || n_groups < 0) return NULL;
if (!agg_needs_row_gather(src->type, op)) return NULL;
bool has_nulls = ray_vec_may_have_nulls(src);
int64_t* winners = ray_alloc_raw((size_t)(n_groups ? n_groups : 1) * sizeof(int64_t));
if (!winners) return ray_error("oom", NULL);
wide_winners_t work = {src, op, has_nulls};
ray_group_winners(wide_winner_rows, &work, idx_buf, offsets, grp_cnt, n_groups, winners);
ray_pool_t* pool = ray_pool_get();
if (ray_interrupted() || (pool && atomic_load_explicit(&pool->cancelled, memory_order_relaxed))) {
ray_free_raw(winners); return ray_error("cancel", NULL);
}
ray_t* out = ray_group_gather(src, winners, n_groups);
ray_free_raw(winners);
return out;
}
static ray_t* reduction_i64_result(int64_t val, int8_t out_type, ray_t* src) {
switch (out_type) {
case RAY_BOOL: return ray_bool((bool)val);
case RAY_DATE: return ray_date((int32_t)val);
case RAY_TIME: return ray_time(val);
case RAY_TIMESTAMP: return ray_timestamp(val);
case RAY_I32: return ray_i32((int32_t)val);
case RAY_I16: return ray_i16((int16_t)val);
case RAY_U8: return ray_u8((uint8_t)val);
case RAY_SYM: return ray_sym(src ? sym_id_runtime(src, val) : val);
default: return ray_i64(val);
}
}
static ray_t* reduction_extreme_result(ray_op_t* op, int8_t in_type, bool found,
double fval, int64_t ival, ray_t* src) {
int8_t out_type = op->out_type ? op->out_type : in_type;
if (!found) return ray_typed_null(-out_type);
if (out_type == RAY_F64) return ray_f64(ray_f64_fin(fval));
if (out_type == RAY_F32) return ray_f32((float)fval);
return reduction_i64_result(ival, out_type, out_type == RAY_SYM ? src : NULL);
}
ray_t* exec_reduction(ray_graph_t* g, ray_op_t* op, ray_t* input) {
if (!input || RAY_IS_ERR(input)) return input;
if (input->type == RAY_TABLE) {
if (op->opcode == OP_COUNT)
return ray_i64(ray_table_nrows(input));
return ray_error("type", "reduction: %s requires a column, got a table", ray_opcode_name(op->opcode));
}
if (ray_is_atom(input)) {
switch (op->opcode) {
case OP_COUNT:
if ((-input->type) == RAY_STR)
return ray_i64((int64_t)ray_str_len(input));
return ray_i64(1);
case OP_SUM: return ray_sum_fn(input);
case OP_ALL: return ray_all_fn(input);
case OP_ANY: return ray_any_fn(input);
case OP_AVG: return ray_avg_fn(input);
case OP_MIN: return ray_min_fn(input);
case OP_MAX: return ray_max_fn(input);
case OP_FIRST: return ray_first_fn(input);
case OP_LAST: return ray_last_fn(input);
case OP_MEDIAN: return ray_med_fn(input);
case OP_VAR: return ray_var_fn(input);
case OP_VAR_POP: return ray_var_pop_fn(input);
case OP_STDDEV: return ray_stddev_fn(input);
case OP_STDDEV_POP: return ray_stddev_pop_fn(input);
case OP_PROD: ray_retain(input); return input;
default: return ray_error("type", "reduction: unsupported op %s on atom of type %s", ray_opcode_name(op->opcode), ray_type_name(input->type));
}
}
int8_t in_type = input->type;
int64_t len = input->len;
bool has_nulls = ray_vec_may_have_nulls(input);
ray_t* sel_idx_block = NULL;
const int64_t* sel_idx = NULL;
int64_t scan_n = len;
if (g && g->selection) {
ray_rowsel_t* sm = ray_rowsel_meta(g->selection);
if (sm->nrows == len) {
scan_n = sm->total_pass;
if (op->opcode != OP_COUNT) {
sel_idx_block = ray_rowsel_to_indices(g->selection);
if (!sel_idx_block) return ray_error("oom", NULL);
sel_idx = (const int64_t*)ray_data(sel_idx_block);
}
}
}
if (op->opcode == OP_COUNT)
return ray_i64(scan_n);
if (agg_needs_row_gather(in_type, op->opcode)) {
ray_t* r = agg_wide_reduce(input, op->opcode, sel_idx, scan_n, has_nulls);
if (sel_idx_block) ray_release(sel_idx_block);
return r;
}
if ((op->opcode == OP_FIRST || op->opcode == OP_LAST) &&
(in_type == RAY_I64 || in_type == RAY_F64 || in_type == RAY_F32 || in_type == RAY_I32 ||
in_type == RAY_I16 || in_type == RAY_BOOL || in_type == RAY_U8 ||
in_type == RAY_TIMESTAMP || in_type == RAY_DATE || in_type == RAY_TIME ||
in_type == RAY_SYM)) {
int64_t row = -1;
if (op->opcode == OP_FIRST) {
for (int64_t i = 0; i < scan_n; i++) {
int64_t r = sel_idx ? sel_idx[i] : i;
if (!has_nulls || !ray_vec_is_null(input, r)) { row = r; break; }
}
} else {
for (int64_t i = scan_n - 1; i >= 0; i--) {
int64_t r = sel_idx ? sel_idx[i] : i;
if (!has_nulls || !ray_vec_is_null(input, r)) { row = r; break; }
}
}
if (sel_idx_block) ray_release(sel_idx_block);
if (row < 0 || row >= len)
return ray_typed_null(-in_type);
void* base = ray_data(input);
if (in_type == RAY_F64) return ray_f64(((const double*)base)[row]);
if (in_type == RAY_F32) return ray_f64((double)((const float*)base)[row]);
return reduction_i64_result(read_col_i64(base, row, in_type, input->attrs), in_type,
in_type == RAY_SYM ? input : NULL);
}
ray_pool_t* pool = ray_pool_get();
if (pool && scan_n >= RAY_PARALLEL_THRESHOLD) {
uint32_t nw = ray_pool_total_workers(pool);
ray_t* accs_hdr;
reduce_acc_t* accs = (reduce_acc_t*)scratch_calloc(&accs_hdr, nw * sizeof(reduce_acc_t));
if (!accs) { if (sel_idx_block) ray_release(sel_idx_block); return ray_error("oom", NULL); }
for (uint32_t i = 0; i < nw; i++) reduce_acc_init(&accs[i]);
par_reduce_ctx_t ctx = { .input = input, .accs = accs,
.has_nulls = has_nulls, .idx = sel_idx,
.opcode = op->opcode };
ray_pool_dispatch(pool, par_reduce_fn, &ctx, scan_n);
reduce_acc_t merged;
reduce_acc_init(&merged);
merged = accs[0];
for (uint32_t i = 1; i < nw; i++) {
reduce_merge(&merged, &accs[i], in_type, ray_sym_vec_domain(input),
op->opcode);
}
for (uint32_t i = 0; i < nw; i++) {
if (accs[i].has_first) {
if (in_type == RAY_F64 || in_type == RAY_F32) merged.first_f = accs[i].first_f;
else merged.first_i = accs[i].first_i;
break;
}
}
for (int32_t i = (int32_t)nw - 1; i >= 0; i--) {
if (accs[i].has_first) {
if (in_type == RAY_F64 || in_type == RAY_F32) merged.last_f = accs[i].last_f;
else merged.last_i = accs[i].last_i;
break;
}
}
ray_t* result;
switch (op->opcode) {
case OP_SUM: result = (in_type == RAY_F64 || in_type == RAY_F32) ? ray_f64(ray_f64_fin(merged.sum_f)) : (in_type == RAY_TIME ? ray_time(merged.sum_i) : ray_i64(merged.sum_i)); break;
case OP_PROD: result = (in_type == RAY_F64 || in_type == RAY_F32) ? ray_f64(ray_f64_fin(merged.prod_f)) : ray_i64(merged.prod_i); break;
case OP_ALL: result = ray_bool(merged.zero_count == 0); break;
case OP_ANY: result = ray_bool(merged.cnt > merged.zero_count); break;
case OP_MIN: result = reduction_extreme_result(op, in_type, merged.cnt > 0, merged.min_f, merged.min_i, input); break;
case OP_MAX: result = reduction_extreme_result(op, in_type, merged.cnt > 0, merged.max_f, merged.max_i, input); break;
case OP_COUNT: result = ray_i64(scan_n); break;
case OP_AVG: result = merged.cnt > 0 ? ray_f64(ray_f64_fin((in_type == RAY_F64 || in_type == RAY_F32) ? merged.sum_f / merged.cnt : ray_i128_to_f64(merged.sum_hi, (uint64_t)merged.sum_i) / merged.cnt)) : ray_typed_null(-RAY_F64); break;
case OP_FIRST: result = merged.has_first ? (group_fp_type(in_type) ? ray_f64(merged.first_f) : reduction_i64_result(merged.first_i, in_type, in_type == RAY_SYM ? input : NULL)) : ray_typed_null(-(op->out_type ? op->out_type : in_type)); break;
case OP_LAST: result = merged.has_first ? (group_fp_type(in_type) ? ray_f64(merged.last_f) : reduction_i64_result(merged.last_i, in_type, in_type == RAY_SYM ? input : NULL)) : ray_typed_null(-(op->out_type ? op->out_type : in_type)); break;
case OP_VAR: case OP_VAR_POP:
case OP_STDDEV: case OP_STDDEV_POP: {
bool insufficient = (op->opcode == OP_VAR || op->opcode == OP_STDDEV) ? merged.cnt <= 1 : merged.cnt <= 0;
if (insufficient) { result = ray_typed_null(-RAY_F64); break; }
double mean, var_pop;
if (in_type == RAY_F64 || in_type == RAY_F32) { mean = merged.sum_f / merged.cnt; var_pop = merged.sum_sq_f / merged.cnt - mean * mean; }
else { mean = merged.sum_d / merged.cnt; var_pop = (double)merged.sum_sq_i / merged.cnt - mean * mean; }
if (var_pop < 0) var_pop = 0;
double val;
if (op->opcode == OP_VAR_POP) val = var_pop;
else if (op->opcode == OP_VAR) val = var_pop * merged.cnt / (merged.cnt - 1);
else if (op->opcode == OP_STDDEV_POP) val = sqrt(var_pop);
else val = sqrt(var_pop * merged.cnt / (merged.cnt - 1));
result = ray_f64(ray_f64_fin(val));
break;
}
default: result = ray_error("nyi", NULL); break;
}
scratch_free(accs_hdr);
if (sel_idx_block) ray_release(sel_idx_block);
return result;
}
reduce_acc_t acc;
reduce_acc_init(&acc);
reduce_range(input, 0, scan_n, &acc, has_nulls, sel_idx, op->opcode);
if (sel_idx_block) ray_release(sel_idx_block);
switch (op->opcode) {
case OP_SUM: return (in_type == RAY_F64 || in_type == RAY_F32) ? ray_f64(ray_f64_fin(acc.sum_f)) : (in_type == RAY_TIME ? ray_time(acc.sum_i) : ray_i64(acc.sum_i));
case OP_PROD: return (in_type == RAY_F64 || in_type == RAY_F32) ? ray_f64(ray_f64_fin(acc.prod_f)) : ray_i64(acc.prod_i);
case OP_ALL: return ray_bool(acc.zero_count == 0);
case OP_ANY: return ray_bool(acc.cnt > acc.zero_count);
case OP_MIN: return reduction_extreme_result(op, in_type, acc.cnt > 0, acc.min_f, acc.min_i, input);
case OP_MAX: return reduction_extreme_result(op, in_type, acc.cnt > 0, acc.max_f, acc.max_i, input);
case OP_COUNT: return ray_i64(scan_n);
case OP_AVG: return acc.cnt > 0 ? ray_f64(ray_f64_fin((in_type == RAY_F64 || in_type == RAY_F32) ? acc.sum_f / acc.cnt : ray_i128_to_f64(acc.sum_hi, (uint64_t)acc.sum_i) / acc.cnt)) : ray_typed_null(-RAY_F64);
case OP_FIRST: return acc.has_first ? (group_fp_type(in_type) ? ray_f64(acc.first_f) : reduction_i64_result(acc.first_i, in_type, in_type == RAY_SYM ? input : NULL)) : ray_typed_null(-(op->out_type ? op->out_type : in_type));
case OP_LAST: return acc.has_first ? (group_fp_type(in_type) ? ray_f64(acc.last_f) : reduction_i64_result(acc.last_i, in_type, in_type == RAY_SYM ? input : NULL)) : ray_typed_null(-(op->out_type ? op->out_type : in_type));
case OP_VAR: case OP_VAR_POP:
case OP_STDDEV: case OP_STDDEV_POP: {
bool insufficient = (op->opcode == OP_VAR || op->opcode == OP_STDDEV) ? acc.cnt <= 1 : acc.cnt <= 0;
if (insufficient) return ray_typed_null(-RAY_F64);
double mean, var_pop;
if (in_type == RAY_F64 || in_type == RAY_F32) { mean = acc.sum_f / acc.cnt; var_pop = acc.sum_sq_f / acc.cnt - mean * mean; }
else { mean = acc.sum_d / acc.cnt; var_pop = (double)acc.sum_sq_i / acc.cnt - mean * mean; }
if (var_pop < 0) var_pop = 0;
double val;
if (op->opcode == OP_VAR_POP) val = var_pop;
else if (op->opcode == OP_VAR) val = var_pop * acc.cnt / (acc.cnt - 1);
else if (op->opcode == OP_STDDEV_POP) val = sqrt(var_pop);
else val = sqrt(var_pop * acc.cnt / (acc.cnt - 1));
return ray_f64(ray_f64_fin(val));
}
default: return ray_error("nyi", NULL);
}
}
static inline void ght_layout_point_inline(ght_layout_t* ly) {
ly->agg_val_slot = ly->agg_val_slot_in;
ly->agg_flags = ly->agg_flags_in;
ly->agg_flags2 = ly->agg_flags2_in;
ly->agg_null_sentinel = ly->agg_null_sentinel_in;
ly->agg_dom = ly->agg_dom_in;
ly->key_off = ly->key_off_in;
ly->key_flags = ly->key_flags_in;
ly->wide_key_esz = ly->wide_key_esz_in;
ly->wide_key_type = ly->wide_key_type_in;
}
static bool ght_layout_alloc_spill(ght_layout_t* ly, uint32_t n_keys, uint32_t n_aggs) {
size_t off = 0;
size_t dom_off = off; off += (size_t)n_aggs * sizeof(void*);
size_t sent_off = off; off += (size_t)n_aggs * sizeof(int64_t);
size_t koff_off = off; off += (size_t)(n_keys + 1) * sizeof(uint16_t);
off = (off + 1u) & ~(size_t)1u;
size_t vslot_off = off; off += (size_t)n_aggs;
size_t aflags_off = off; off += (size_t)n_aggs;
size_t aflags2_off = off; off += (size_t)n_aggs;
size_t kflags_off = off; off += (size_t)n_keys;
size_t wesz_off = off; off += (size_t)n_keys;
size_t wtype_off = off; off += (size_t)n_keys;
char* base = (char*)scratch_calloc(&ly->spill_hdr, off ? off : 1);
if (!base) { ly->spill_hdr = NULL; return false; }
ly->agg_dom = (struct ray_sym_domain_s**)(void*)(base + dom_off);
ly->agg_null_sentinel = (int64_t*)(void*)(base + sent_off);
ly->key_off = (uint16_t*)(void*)(base + koff_off);
ly->agg_val_slot = (int8_t*)(base + vslot_off);
ly->agg_flags = (uint8_t*)(base + aflags_off);
ly->agg_flags2 = (uint8_t*)(base + aflags2_off);
ly->key_flags = (uint8_t*)(base + kflags_off);
ly->wide_key_esz = (uint8_t*)(base + wesz_off);
ly->wide_key_type = (int8_t*)(base + wtype_off);
return true;
}
void ght_layout_free(ght_layout_t* ly) {
if (ly && ly->spill_hdr) { scratch_free(ly->spill_hdr); ly->spill_hdr = NULL; }
}
void ght_layout_copy(ght_layout_t* dst, const ght_layout_t* src) {
*dst = *src;
if (src->agg_val_slot == src->agg_val_slot_in) {
ght_layout_point_inline(dst);
dst->spill_hdr = NULL;
} else {
dst->spill_hdr = NULL;
}
}
static inline bool ray_key_may_be_null(const ray_t* kv);
bool ght_compute_layout(ght_layout_t* out, uint32_t n_keys, uint32_t n_aggs,
ray_t** agg_vecs, ray_t** agg_vecs2,
uint8_t need_flags,
const uint16_t* agg_ops,
const int8_t* key_types,
ray_t* const* key_vecs) {
memset(out, 0, sizeof(*out));
out->n_keys = (uint16_t)n_keys;
out->n_aggs = (uint16_t)n_aggs;
out->need_flags = need_flags;
if (n_keys > GHT_INLINE || n_aggs > GHT_INLINE) {
if (!ght_layout_alloc_spill(out, n_keys, n_aggs)) return false;
} else {
ght_layout_point_inline(out);
}
if (key_types) {
for (uint32_t k = 0; k < n_keys; k++) {
if (key_types[k] == RAY_GUID) {
out->key_flags[k] = GHT_KEYF_WIDE;
out->wide_key_esz[k] = 16;
out->wide_key_type[k] = RAY_GUID;
out->any_wide_key = 1;
} else if (key_types[k] == RAY_STR) {
out->key_flags[k] = (uint8_t)(GHT_KEYF_WIDE | GHT_KEYF_INLINE_STR);
out->wide_key_esz[k] = (uint8_t)sizeof(ray_str_t);
out->wide_key_type[k] = RAY_STR;
out->any_wide_key = 1;
out->any_inline_str = 1;
}
}
if (n_keys >= 1 && !out->any_wide_key) {
bool packed = true;
for (uint32_t k = 0; k < n_keys && packed; k++) {
switch (key_types[k]) {
case RAY_BOOL: case RAY_U8: case RAY_I16: case RAY_I32:
case RAY_I64: case RAY_DATE: case RAY_TIME:
case RAY_TIMESTAMP: case RAY_SYM:
break;
default:
packed = false;
break;
}
}
out->packed_key = packed ? 1 : 0;
}
}
uint16_t nv = 0;
uint8_t agg_any = 0;
for (uint32_t a = 0; a < n_aggs; a++) {
bool wide_mm = agg_ops && agg_vecs[a] &&
agg_needs_row_gather(agg_vecs[a]->type, agg_ops[a]);
bool holistic = agg_ops && (agg_ops[a] == OP_MEDIAN ||
agg_ops[a] == OP_QUANTILE ||
agg_ops[a] == OP_MODE ||
agg_ops[a] == OP_TOP_N ||
agg_ops[a] == OP_BOT_N || wide_mm);
bool valueless = agg_ops && agg_ops[a] == OP_COUNT;
uint8_t af = 0;
if (holistic) {
af |= GHT_AF_HOLISTIC;
if (wide_mm) af |= GHT_AF_WIDE;
out->agg_val_slot[a] = -1;
} else if (valueless) {
out->agg_val_slot[a] = -1;
} else if (agg_vecs[a]) {
out->agg_val_slot[a] = (int8_t)nv;
bool pair = agg_ops && agg_is_binary_agg(agg_ops[a]);
bool pair_fp = pair && agg_vecs2 && agg_vecs2[a] &&
group_fp_type(agg_vecs2[a]->type);
if (group_fp_type(agg_vecs[a]->type) || pair_fp)
af |= GHT_AF_F64;
if (agg_vecs[a]->type == RAY_SYM) {
af |= GHT_AF_SYM;
out->agg_dom[a] = ray_sym_vec_domain(agg_vecs[a]);
}
nv++;
if (agg_ops && agg_is_binary_agg(agg_ops[a])) {
af |= GHT_AF_BINARY;
nv++;
}
} else {
out->agg_val_slot[a] = -1;
}
if (agg_ops && !wide_mm) {
if (agg_ops[a] == OP_FIRST) af |= GHT_AF_FIRST;
if (agg_ops[a] == OP_LAST) af |= GHT_AF_LAST;
if (agg_ops[a] == OP_PROD) af |= GHT_AF_PROD;
}
out->agg_flags[a] = af;
agg_any |= af;
uint8_t af2 = 0;
int64_t sent = 0;
int8_t vslot = out->agg_val_slot[a];
if (agg_ops && (agg_ops[a] == OP_ALL || agg_ops[a] == OP_ANY))
af2 |= GHT_AF2_TRUTHY;
if (vslot >= 0 && agg_vecs[a]) {
ray_t* src = (agg_vecs[a]->attrs & RAY_ATTR_SLICE)
? agg_vecs[a]->slice_parent : agg_vecs[a];
if (src && ray_vec_may_have_nulls(src)) {
af2 |= GHT_AF2_NULLABLE;
sent = agg_int_null_sentinel_for(agg_vecs[a]->type);
out->any_agg_null = 1;
}
}
if (vslot >= 0 && agg_ops && agg_is_binary_agg(agg_ops[a]) &&
agg_vecs2 && agg_vecs2[a]) {
ray_t* ysrc = (agg_vecs2[a]->attrs & RAY_ATTR_SLICE)
? agg_vecs2[a]->slice_parent : agg_vecs2[a];
if (ysrc && ray_vec_may_have_nulls(ysrc)) {
af2 |= GHT_AF2_Y_NULLABLE;
out->any_agg_null = 1;
}
}
out->agg_flags2[a] = af2;
out->agg_null_sentinel[a] = sent;
}
out->n_agg_vals = nv;
out->agg_flags_any = agg_any;
uint32_t null_words = n_keys ? ((n_keys + 63u) >> 6) : 1u;
if (key_vecs && n_keys >= 1 && n_keys <= 64) {
bool any_nullable = false;
for (uint32_t k = 0; k < n_keys && !any_nullable; k++)
any_nullable = ray_key_may_be_null(key_vecs[k]);
if (!any_nullable) null_words = 0;
}
out->null_words = (uint16_t)null_words;
{
uint32_t koff = 0;
for (uint32_t k = 0; k < n_keys; k++) {
if (koff > UINT16_MAX) { ght_layout_free(out); return false; }
out->key_off[k] = (uint16_t)koff;
koff += (out->key_flags[k] & GHT_KEYF_INLINE_STR) ? 16 : 8;
}
if (koff > UINT16_MAX) { ght_layout_free(out); return false; }
out->key_off[n_keys] = (uint16_t)koff;
koff += null_words * 8;
if (koff > UINT16_MAX) { ght_layout_free(out); return false; }
out->key_region = (uint16_t)koff;
}
uint16_t key_region = out->key_region;
bool has_first_last = (agg_any & (GHT_AF_FIRST | GHT_AF_LAST)) != 0;
uint32_t entry_stride = 8u + key_region + (uint32_t)nv * 8 + 8u;
uint32_t off = 8u + key_region;
uint32_t block = (uint32_t)nv * 8;
if (need_flags & GHT_NEED_SUM) { out->off_sum = (uint16_t)off; off += block; }
if (need_flags & GHT_NEED_MIN) { out->off_min = (uint16_t)off; off += block; }
if (need_flags & GHT_NEED_MAX) { out->off_max = (uint16_t)off; off += block; }
if (need_flags & GHT_NEED_SUMSQ) { out->off_sumsq = (uint16_t)off; off += block; }
if (need_flags & GHT_NEED_SUM128) { out->off_sum_hi = (uint16_t)off; off += block; }
if (has_first_last) {
out->off_first_row = (uint16_t)off; off += block;
out->off_last_row = (uint16_t)off; off += block;
}
if (need_flags & GHT_NEED_PEARSON) {
out->off_sum_y = (uint16_t)off; off += block;
out->off_sumsq_y = (uint16_t)off; off += block;
out->off_sumxy = (uint16_t)off; off += block;
}
out->off_group_first = (uint16_t)off; off += 8u;
if (out->any_agg_null) { out->off_nn = (uint16_t)off; off += block; }
if (nv > INT8_MAX || entry_stride > UINT16_MAX || off > UINT16_MAX) {
ght_layout_free(out);
return false;
}
out->entry_stride = (uint16_t)entry_stride;
out->row_stride = (uint16_t)off;
return true;
}
#define HT_EMPTY UINT32_MAX
#define HT_PACK(salt, gid) (((uint32_t)(uint8_t)(salt) << 24) | ((gid) & 0xFFFFFF))
#define HT_GID(s) ((s) & 0xFFFFFF)
#define HT_SALT_V(s) ((uint8_t)((s) >> 24))
static bool group_ht_wire_wide(group_ht_t* ht, uint32_t nk) {
uint32_t n = nk > 8 ? nk : 8;
if (nk <= 8) {
ht->key_wide_hdr = NULL;
ht->key_data = ht->key_data_in;
ht->key_pool = ht->key_pool_in;
} else {
ht->key_wide_hdr = NULL;
void* blk = scratch_alloc(&ht->key_wide_hdr, (size_t)nk * 2 * sizeof(void*));
if (!blk) return false;
ht->key_data = (void**)blk;
ht->key_pool = (const void**)((void**)blk + nk);
}
memset(ht->key_data, 0, (size_t)n * sizeof(void*));
memset((void*)ht->key_pool, 0, (size_t)n * sizeof(void*));
return true;
}
static bool group_ht_init_sized(group_ht_t* ht, uint32_t cap,
const ght_layout_t* ly, uint32_t init_grp_cap) {
ht->ht_cap = cap;
ht->oom = 0;
ht->grow_cap = 0;
ght_layout_copy(&ht->layout, ly);
if (!group_ht_wire_wide(ht, ly->n_keys)) return false;
ht->slots = (uint32_t*)scratch_alloc(&ht->_h_slots, (size_t)cap * sizeof(uint32_t));
if (!ht->slots) return false;
memset(ht->slots, 0xFF, (size_t)cap * sizeof(uint32_t));
ht->grp_cap = init_grp_cap;
ht->grp_count = 0;
ht->rows = (char*)scratch_alloc(&ht->_h_rows,
(size_t)init_grp_cap * ly->row_stride);
if (!ht->rows) return false;
return true;
}
bool group_ht_init(group_ht_t* ht, uint32_t cap, const ght_layout_t* ly) {
return group_ht_init_sized(ht, cap, ly, 1);
}
static inline void group_ht_set_key_data(group_ht_t* ht, void** kd) {
if (!ht->layout.any_wide_key || !kd) return;
const uint8_t* const kflags = ht->layout.key_flags;
for (uint16_t k = 0; k < ht->layout.n_keys; k++) {
if (kflags[k] & GHT_KEYF_WIDE) ht->key_data[k] = kd[k];
}
}
static inline void derive_key_pool(const ght_layout_t* ly, ray_t* const* key_vecs,
const void** out) {
for (uint16_t k = 0; k < ly->n_keys; k++) {
out[k] = NULL;
if ((ly->key_flags[k] & GHT_KEYF_WIDE) && ly->wide_key_type[k] == RAY_STR
&& key_vecs && key_vecs[k]) {
ray_t* kv = key_vecs[k];
ray_t* owner = (kv->attrs & RAY_ATTR_SLICE) ? kv->slice_parent : kv;
ray_t* pool = owner ? owner->str_pool : NULL;
out[k] = (pool && !RAY_IS_ERR(pool)) ? ray_data(pool) : NULL;
}
}
}
static inline void group_ht_set_key_pool(group_ht_t* ht, ray_t* const* key_vecs) {
if (ht->layout.any_wide_key && key_vecs)
derive_key_pool(&ht->layout, key_vecs, ht->key_pool);
}
static inline void group_ht_copy_key_pool(group_ht_t* ht, const void* const* kp) {
if (!ht->layout.any_wide_key || !kp) return;
const uint8_t* const kflags = ht->layout.key_flags;
for (uint16_t k = 0; k < ht->layout.n_keys; k++)
if (kflags[k] & GHT_KEYF_WIDE) ht->key_pool[k] = kp[k];
}
void group_ht_free(group_ht_t* ht) {
scratch_free(ht->_h_slots);
scratch_free(ht->_h_rows);
scratch_free(ht->key_wide_hdr);
}
static bool group_ht_grow(group_ht_t* ht) {
uint32_t old_cap = ht->grp_cap;
uint32_t new_cap = old_cap * 2;
uint16_t rs = ht->layout.row_stride;
char* new_rows = (char*)scratch_realloc(
&ht->_h_rows, (size_t)old_cap * rs, (size_t)new_cap * rs);
if (!new_rows) return false;
ht->rows = new_rows;
ht->grp_cap = new_cap;
return true;
}
static inline uint64_t wide_key_hash_at(const ght_layout_t* ly, uint8_t k,
void* const* key_data,
const void* const* key_pool, int64_t row) {
if (ly->wide_key_type[k] == RAY_STR) {
const ray_str_t* d = &((const ray_str_t*)key_data[k])[row];
return ray_str_t_hash(d, key_pool ? (const char*)key_pool[k] : NULL);
}
uint8_t esz = ly->wide_key_esz[k];
return ray_hash_bytes((const char*)key_data[k] + (size_t)row * esz, esz);
}
static inline bool wide_key_eq_at(const ght_layout_t* ly, uint8_t k,
void* const* key_data, const void* const* key_pool,
int64_t ra, int64_t rb) {
if (ra == rb) return true;
if (ly->wide_key_type[k] == RAY_STR) {
const ray_str_t* a = &((const ray_str_t*)key_data[k])[ra];
const ray_str_t* b = &((const ray_str_t*)key_data[k])[rb];
const char* pool = key_pool ? (const char*)key_pool[k] : NULL;
return ray_str_t_eq(a, pool, b, pool);
}
uint8_t esz = ly->wide_key_esz[k];
const char* base = (const char*)key_data[k];
return memcmp(base + (size_t)ra * esz, base + (size_t)rb * esz, esz) == 0;
}
static inline uint64_t ght_hash_null_words(uint64_t h, const int64_t* nullw,
uint32_t null_words) {
for (uint32_t w = 0; w < null_words; w++)
if (nullw[w]) h = ray_hash_combine(h, ray_hash_i64(nullw[w]));
return h;
}
static inline uint64_t ght_hash_lanes(const int64_t* lane, uint32_t lanes) {
uint64_t a = 0x2d358dccaa6c78a5ULL;
uint64_t b = 0x8bb84b93962eacc9ULL ^ (uint64_t)lanes;
uint32_t k = 0;
for (; k + 2 <= lanes; k += 2) {
a ^= (uint64_t)lane[k];
b ^= (uint64_t)lane[k + 1];
ray__wymum(&a, &b);
}
if (k < lanes) {
uint64_t v = (uint64_t)lane[k];
a ^= v;
b ^= v;
ray__wymum(&a, &b);
}
return ray__wymix(a ^ 0x2d358dccaa6c78a5ULL, b ^ 0x8bb84b93962eacc9ULL);
}
static inline void group_keys_gather(int64_t* dst, const void* data,
const int64_t* rows, uint32_t m,
int8_t type, uint8_t attrs) {
switch (type) {
case RAY_I64: case RAY_TIMESTAMP:
for (uint32_t j = 0; j < m; j++) dst[j] = ((const int64_t*)data)[rows[j]];
return;
case RAY_SYM:
switch (attrs & RAY_SYM_W_MASK) {
case RAY_SYM_W8:
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)((const uint8_t*)data)[rows[j]];
return;
case RAY_SYM_W16:
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)((const uint16_t*)data)[rows[j]];
return;
case RAY_SYM_W32:
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)((const uint32_t*)data)[rows[j]];
return;
default:
for (uint32_t j = 0; j < m; j++) dst[j] = ((const int64_t*)data)[rows[j]];
return;
}
case RAY_I32: case RAY_DATE: case RAY_TIME:
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)((const int32_t*)data)[rows[j]];
return;
case RAY_I16:
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)((const int16_t*)data)[rows[j]];
return;
default:
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)((const uint8_t*)data)[rows[j]];
return;
}
}
static inline void group_keys_gather_seq(int64_t* dst, const void* data,
int64_t row0, uint32_t m,
int8_t type, uint8_t attrs) {
switch (type) {
case RAY_I64: case RAY_TIMESTAMP: {
const int64_t* p = (const int64_t*)data + row0;
for (uint32_t j = 0; j < m; j++) dst[j] = p[j];
return;
}
case RAY_SYM:
switch (attrs & RAY_SYM_W_MASK) {
case RAY_SYM_W8: {
const uint8_t* p = (const uint8_t*)data + row0;
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)p[j];
return;
}
case RAY_SYM_W16: {
const uint16_t* p = (const uint16_t*)data + row0;
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)p[j];
return;
}
case RAY_SYM_W32: {
const uint32_t* p = (const uint32_t*)data + row0;
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)p[j];
return;
}
default: {
const int64_t* p = (const int64_t*)data + row0;
for (uint32_t j = 0; j < m; j++) dst[j] = p[j];
return;
}
}
case RAY_I32: case RAY_DATE: case RAY_TIME: {
const int32_t* p = (const int32_t*)data + row0;
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)p[j];
return;
}
case RAY_I16: {
const int16_t* p = (const int16_t*)data + row0;
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)p[j];
return;
}
default: {
const uint8_t* p = (const uint8_t*)data + row0;
for (uint32_t j = 0; j < m; j++) dst[j] = (int64_t)p[j];
return;
}
}
}
static const int64_t ght_null_words_none[1] = { 0 };
static inline const int64_t* ght_null_words_at(const ght_layout_t* ly,
const void* keybase) {
if (!ly->null_words) return ght_null_words_none;
return (const int64_t*)(const void*)((const char*)keybase +
ly->key_off[ly->n_keys]);
}
static inline bool ray_key_may_be_null(const ray_t* kv) {
if (!kv || kv->type == RAY_SYM || kv->type == RAY_STR) return false;
const ray_t* src = (kv->attrs & RAY_ATTR_SLICE) ? kv->slice_parent : kv;
return src && ray_vec_may_have_nulls(src);
}
static inline uint64_t inline_str_hash(const ght_layout_t* ly, uint8_t k,
const void* keybase,
const void* const* key_pool) {
const ray_str_t* d = (const ray_str_t*)((const char*)keybase + ly->key_off[k]);
return ray_str_t_hash(d, key_pool ? (const char*)key_pool[k] : NULL);
}
static inline bool inline_str_eq(const ght_layout_t* ly, uint8_t k,
const void* a_base, const void* b_base,
const void* const* key_pool) {
const ray_str_t* a = (const ray_str_t*)((const char*)a_base + ly->key_off[k]);
const ray_str_t* b = (const ray_str_t*)((const char*)b_base + ly->key_off[k]);
const char* pool = key_pool ? (const char*)key_pool[k] : NULL;
return ray_str_t_eq(a, pool, b, pool);
}
static inline uint64_t inline_layout_key_hash(const ght_layout_t* ly, uint8_t k,
const int8_t* key_types,
const void* keybase, void* const* key_data,
const void* const* key_pool) {
if (ly->key_flags[k] & GHT_KEYF_INLINE_STR)
return inline_str_hash(ly, k, keybase, key_pool);
int64_t v = *(const int64_t*)((const char*)keybase + ly->key_off[k]);
if (ly->key_flags[k] & GHT_KEYF_WIDE)
return wide_key_hash_at(ly, k, key_data, key_pool, v);
if (key_types[k] == RAY_F64) {
double dv; memcpy(&dv, &v, 8); return ray_hash_f64(dv);
}
return ray_hash_i64(v);
}
static inline uint64_t inline_build_keys(const ght_layout_t* ly, const int8_t* key_types,
void* const* key_data, const uint8_t* key_attrs, const void* const* key_pool,
ray_t* const* key_vecs, uint8_t any_nullable, int64_t row,
char* keybase) {
uint64_t h = 0;
uint16_t nk = ly->n_keys;
const uint8_t* const kflags = ly->key_flags;
const uint16_t* const koff = ly->key_off;
int64_t* nullw = (int64_t*)(keybase + koff[nk]);
uint32_t null_words = ly->null_words;
for (uint32_t w = 0; w < null_words; w++) nullw[w] = 0;
for (uint32_t k = 0; k < nk; k++) {
char* slot = keybase + koff[k];
uint64_t kh;
bool is_null = any_nullable && key_vecs && ray_key_may_be_null(key_vecs[k])
&& ray_vec_is_null(key_vecs[k], row);
if (is_null) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
if (kflags[k] & GHT_KEYF_INLINE_STR) memset(slot, 0, sizeof(ray_str_t));
else *(int64_t*)slot = 0;
kh = ray_hash_i64(0);
} else if (kflags[k] & GHT_KEYF_INLINE_STR) {
*(ray_str_t*)slot = ((const ray_str_t*)key_data[k])[row];
kh = ray_str_t_hash((const ray_str_t*)slot,
key_pool ? (const char*)key_pool[k] : NULL);
} else if (kflags[k] & GHT_KEYF_WIDE) {
*(int64_t*)slot = row;
kh = wide_key_hash_at(ly, k, key_data, key_pool, row);
} else if (key_types[k] == RAY_F64) {
int64_t kv = group_key_f64_bits(key_data[k], row);
double dv; memcpy(&dv, &kv, 8);
*(int64_t*)slot = kv;
kh = ray_hash_f64(dv);
} else {
int64_t kv = read_col_i64(key_data[k], row, key_types[k], key_attrs[k]);
*(int64_t*)slot = kv;
kh = ray_hash_i64(kv);
}
h = (k == 0) ? kh : ray_hash_combine(h, kh);
}
return ght_hash_null_words(h, nullw, null_words);
}
static inline uint64_t hash_keys_inline(const int64_t* keys, const int8_t* key_types,
uint16_t n_keys, void* const* key_data,
const ght_layout_t* ly,
const void* const* key_pool) {
if (ly->any_inline_str) {
uint64_t h = 0;
const int64_t* nullw = (const int64_t*)((const char*)keys + ly->key_off[n_keys]);
const bool nullable = ly->null_words != 0;
for (uint32_t k = 0; k < n_keys; k++) {
uint64_t kh = (nullable && ((uint64_t)nullw[k >> 6] >> (k & 63)) & 1u)
? ray_hash_i64(0)
: inline_layout_key_hash(ly, k, key_types, keys, key_data, key_pool);
h = (k == 0) ? kh : ray_hash_combine(h, kh);
}
return ght_hash_null_words(h, nullw, ly->null_words);
}
if (ly->packed_key)
return ght_hash_lanes(keys, (uint32_t)n_keys + ly->null_words);
const uint8_t* const kflags = ly->key_flags;
const bool nullable = ly->null_words != 0;
uint64_t h = 0;
for (uint32_t k = 0; k < n_keys; k++) {
uint64_t kh;
if (nullable && ((uint64_t)keys[n_keys + (k >> 6)] >> (k & 63)) & 1u) {
kh = ray_hash_i64(0);
} else if (kflags[k] & GHT_KEYF_WIDE) {
kh = wide_key_hash_at(ly, k, key_data, key_pool, keys[k]);
} else if (key_types[k] == RAY_F64) {
double dv;
memcpy(&dv, &keys[k], 8);
kh = ray_hash_f64(dv);
} else {
kh = ray_hash_i64(keys[k]);
}
h = (k == 0) ? kh : ray_hash_combine(h, kh);
}
return ght_hash_null_words(h, keys + n_keys, ly->null_words);
}
static void group_ht_rehash(group_ht_t* ht, const int8_t* key_types) {
uint32_t new_cap = ht->ht_cap * 2;
uint32_t tgt = ht->grow_cap;
uint32_t lim = ht->ht_cap << 2;
if (tgt > lim) tgt = lim;
if (tgt > new_cap) new_cap = tgt;
ray_t* new_h = NULL;
uint32_t* new_slots = (uint32_t*)scratch_alloc(&new_h, (size_t)new_cap * sizeof(uint32_t));
if (!new_slots) return;
scratch_free(ht->_h_slots);
ht->_h_slots = new_h;
ht->slots = new_slots;
memset(ht->slots, 0xFF, (size_t)new_cap * sizeof(uint32_t));
ht->ht_cap = new_cap;
uint32_t mask = new_cap - 1;
uint16_t rs = ht->layout.row_stride;
uint16_t nk = ht->layout.n_keys;
for (uint32_t gi = 0; gi < ht->grp_count; gi++) {
const int64_t* row_keys = (const int64_t*)(ht->rows + (size_t)gi * rs + 8);
uint64_t h = hash_keys_inline(row_keys, key_types, nk, ht->key_data,
&ht->layout, ht->key_pool);
uint32_t slot = (uint32_t)(h & mask);
while (ht->slots[slot] != HT_EMPTY)
slot = (slot + 1) & mask;
ht->slots[slot] = HT_PACK(HT_SALT(h), gi);
}
}
static void group_ht_rebuild_slots(group_ht_t* ht, const int8_t* key_types) {
if (!ht->slots || ht->ht_cap == 0) return;
memset(ht->slots, 0xFF, (size_t)ht->ht_cap * sizeof(uint32_t));
uint32_t mask = ht->ht_cap - 1;
uint16_t rs = ht->layout.row_stride;
uint16_t nk = ht->layout.n_keys;
for (uint32_t gi = 0; gi < ht->grp_count; gi++) {
const int64_t* row_keys = (const int64_t*)(ht->rows + (size_t)gi * rs + 8);
uint64_t h = hash_keys_inline(row_keys, key_types, nk, ht->key_data,
&ht->layout, ht->key_pool);
uint32_t slot = (uint32_t)(h & mask);
while (ht->slots[slot] != HT_EMPTY)
slot = (slot + 1) & mask;
ht->slots[slot] = HT_PACK(HT_SALT(h), gi);
}
}
static void init_accum_from_entry_nullable(char* row, const char* entry,
const ght_layout_t* ly);
static void accum_from_entry_nullable(char* row, const char* entry,
const ght_layout_t* ly);
static inline void init_accum_from_entry(char* row, const char* entry,
const ght_layout_t* ly) {
if (ly->any_agg_null) { init_accum_from_entry_nullable(row, entry, ly); return; }
uint16_t accum_start = (uint16_t)(8 + ly->key_region);
if (ly->row_stride > accum_start)
memset(row + accum_start, 0, ly->row_stride - accum_start);
const char* agg_data = entry + 8 + (size_t)ly->key_region;
uint16_t na = ly->n_aggs;
uint8_t nf = ly->need_flags;
bool has_fl = (ly->agg_flags_any & (GHT_AF_FIRST | GHT_AF_LAST)) != 0;
int64_t entry_row = 0;
if (has_fl)
memcpy(&entry_row, entry + ly->entry_stride - 8, 8);
const int8_t* const vslot = ly->agg_val_slot;
const uint8_t* const aflags = ly->agg_flags;
const uint8_t* const aflags2 = ly->agg_flags2;
for (uint32_t a = 0; a < na; a++) {
int8_t s = vslot[a];
if (s < 0) continue;
uint8_t af = aflags[a];
if (nf & GHT_NEED_SUM) {
if (aflags2[a] & GHT_AF2_TRUTHY) {
int64_t tv;
if (af & GHT_AF_F64) { double v; memcpy(&v, agg_data + s * 8, 8); tv = (v != 0.0); }
else { int64_t v; memcpy(&v, agg_data + s * 8, 8); tv = (v != 0); }
memcpy(row + ly->off_sum + s * 8, &tv, 8);
} else {
memcpy(row + ly->off_sum + s * 8, agg_data + s * 8, 8);
}
if ((nf & GHT_NEED_SUM128) && !(af & GHT_AF_F64) &&
!(aflags2[a] & GHT_AF2_TRUTHY)) {
int64_t v; memcpy(&v, agg_data + s * 8, 8);
if (v < 0) { int64_t hi = -1; memcpy(row + ly->off_sum_hi + s * 8, &hi, 8); }
}
}
if (nf & GHT_NEED_MIN) memcpy(row + ly->off_min + s * 8, agg_data + s * 8, 8);
if (nf & GHT_NEED_MAX) memcpy(row + ly->off_max + s * 8, agg_data + s * 8, 8);
if (nf & GHT_NEED_SUMSQ) {
if (af & GHT_AF_F64) {
double v; memcpy(&v, agg_data + s * 8, 8);
double sq = v * v;
memcpy(row + ly->off_sumsq + s * 8, &sq, 8);
} else {
int64_t v; memcpy(&v, agg_data + s * 8, 8);
double sq = (double)v * (double)v;
memcpy(row + ly->off_sumsq + s * 8, &sq, 8);
}
}
if ((nf & GHT_NEED_PEARSON) && (af & GHT_AF_BINARY)) {
double x, y;
if (af & GHT_AF_F64) {
memcpy(&x, agg_data + s * 8, 8);
memcpy(&y, agg_data + (s + 1) * 8, 8);
} else {
int64_t xi, yi;
memcpy(&xi, agg_data + s * 8, 8);
memcpy(&yi, agg_data + (s + 1) * 8, 8);
x = (double)xi; y = (double)yi;
}
memcpy(row + ly->off_sum_y + s * 8, &y, 8);
double yy = y * y;
memcpy(row + ly->off_sumsq_y + s * 8, &yy, 8);
double xy = x * y;
memcpy(row + ly->off_sumxy + s * 8, &xy, 8);
}
if (has_fl) {
memcpy(row + ly->off_first_row + s * 8, &entry_row, 8);
memcpy(row + ly->off_last_row + s * 8, &entry_row, 8);
}
}
}
static inline void accum_from_entry(char* row, const char* entry,
const ght_layout_t* ly) {
if (ly->any_agg_null) { accum_from_entry_nullable(row, entry, ly); return; }
const char* agg_data = entry + 8 + (size_t)ly->key_region;
uint16_t na = ly->n_aggs;
uint8_t nf = ly->need_flags;
bool has_fl = (ly->agg_flags_any & (GHT_AF_FIRST | GHT_AF_LAST)) != 0;
int64_t entry_row = 0;
if (has_fl)
memcpy(&entry_row, entry + ly->entry_stride - 8, 8);
const int8_t* const vslot = ly->agg_val_slot;
const uint8_t* const aflags = ly->agg_flags;
const uint8_t* const aflags2 = ly->agg_flags2;
struct ray_sym_domain_s* const* const adom = ly->agg_dom;
for (uint32_t a = 0; a < na; a++) {
int8_t s = vslot[a];
if (s < 0) continue;
const char* val = agg_data + s * 8;
uint8_t af = aflags[a];
bool take_first = false, take_last = false;
if (has_fl && (af & GHT_AF_FIRST)) {
int64_t fr; memcpy(&fr, row + ly->off_first_row + s * 8, 8);
take_first = (entry_row < fr);
}
if (has_fl && (af & GHT_AF_LAST)) {
int64_t lr; memcpy(&lr, row + ly->off_last_row + s * 8, 8);
take_last = (entry_row > lr);
}
if (af & GHT_AF_F64) {
double v;
memcpy(&v, val, 8);
if (nf & GHT_NEED_SUM) {
if (af & GHT_AF_FIRST) { if (take_first) memcpy(row + ly->off_sum + s * 8, val, 8); }
else if (af & GHT_AF_LAST) { if (take_last) memcpy(row + ly->off_sum + s * 8, val, 8); }
else if (aflags2[a] & GHT_AF2_TRUTHY) { ROW_WR_I64(row, ly->off_sum, s) += (v != 0.0); }
else if (af & GHT_AF_PROD) { ROW_WR_F64(row, ly->off_sum, s) *= v; }
else { ROW_WR_F64(row, ly->off_sum, s) += v; }
}
if (nf & GHT_NEED_MIN) { double* p = &ROW_WR_F64(row, ly->off_min, s); if (v < *p) *p = v; }
if (nf & GHT_NEED_MAX) { double* p = &ROW_WR_F64(row, ly->off_max, s); if (v > *p) *p = v; }
if (nf & GHT_NEED_SUMSQ) { ROW_WR_F64(row, ly->off_sumsq, s) += v * v; }
if ((nf & GHT_NEED_PEARSON) && (af & GHT_AF_BINARY)) {
double y;
memcpy(&y, agg_data + (s + 1) * 8, 8);
ROW_WR_F64(row, ly->off_sum_y, s) += y;
ROW_WR_F64(row, ly->off_sumsq_y, s) += y * y;
ROW_WR_F64(row, ly->off_sumxy, s) += v * y;
}
} else {
int64_t v;
memcpy(&v, val, 8);
if (nf & GHT_NEED_SUM) {
if (af & GHT_AF_FIRST) { if (take_first) memcpy(row + ly->off_sum + s * 8, val, 8); }
else if (af & GHT_AF_LAST) { if (take_last) memcpy(row + ly->off_sum + s * 8, val, 8); }
else if (aflags2[a] & GHT_AF2_TRUTHY) { ROW_WR_I64(row, ly->off_sum, s) += (v != 0); }
else if (af & GHT_AF_PROD) { ROW_WR_I64(row, ly->off_sum, s) = (int64_t)((uint64_t)ROW_RD_I64(row, ly->off_sum, s) * (uint64_t)v); }
else if (nf & GHT_NEED_SUM128) { ray_i128_add(&ROW_WR_I64(row, ly->off_sum_hi, s), (uint64_t*)&ROW_WR_I64(row, ly->off_sum, s), v); }
else { ROW_WR_I64(row, ly->off_sum, s) = wrap_add_i64(ROW_RD_I64(row, ly->off_sum, s), v); }
}
if (nf & GHT_NEED_MIN) {
int64_t* p = &ROW_WR_I64(row, ly->off_min, s);
if (af & GHT_AF_SYM) {
if (*p == INT64_MAX || sym_lex_lt(adom[a], v, *p)) *p = v;
} else if (v < *p) *p = v;
}
if (nf & GHT_NEED_MAX) {
int64_t* p = &ROW_WR_I64(row, ly->off_max, s);
if (af & GHT_AF_SYM) {
if (*p == INT64_MIN || sym_lex_gt(adom[a], v, *p)) *p = v;
} else if (v > *p) *p = v;
}
if (nf & GHT_NEED_SUMSQ) { ROW_WR_F64(row, ly->off_sumsq, s) += (double)v * (double)v; }
if ((nf & GHT_NEED_PEARSON) && (af & GHT_AF_BINARY)) {
int64_t yi; memcpy(&yi, agg_data + (s + 1) * 8, 8);
double y = (double)yi;
double xd = (double)v;
ROW_WR_F64(row, ly->off_sum_y, s) += y;
ROW_WR_F64(row, ly->off_sumsq_y, s) += y * y;
ROW_WR_F64(row, ly->off_sumxy, s) += xd * y;
}
}
if (take_first) memcpy(row + ly->off_first_row + s * 8, &entry_row, 8);
if (take_last) memcpy(row + ly->off_last_row + s * 8, &entry_row, 8);
}
}
static inline bool agg_entry_is_null(uint8_t af, uint8_t af2, int64_t sentinel,
const char* val) {
if (!(af2 & GHT_AF2_NULLABLE)) return false;
if (af & GHT_AF_F64) { double v; memcpy(&v, val, 8); return v != v; }
int64_t v; memcpy(&v, val, 8); return v == sentinel;
}
static void accum_from_entry_nullable(char* row, const char* entry,
const ght_layout_t* ly) {
const char* agg_data = entry + 8 + (size_t)ly->key_region;
uint16_t na = ly->n_aggs;
uint8_t nf = ly->need_flags;
bool has_fl = (ly->agg_flags_any & (GHT_AF_FIRST | GHT_AF_LAST)) != 0;
int64_t entry_row = 0;
if (has_fl)
memcpy(&entry_row, entry + ly->entry_stride - 8, 8);
const int8_t* const vslot = ly->agg_val_slot;
const uint8_t* const aflags = ly->agg_flags;
const uint8_t* const aflags2 = ly->agg_flags2;
const int64_t* const asent = ly->agg_null_sentinel;
struct ray_sym_domain_s* const* const adom = ly->agg_dom;
int64_t* const nn = (int64_t*)(void*)(row + ly->off_nn);
for (uint32_t a = 0; a < na; a++) {
int8_t s = vslot[a];
if (s < 0) continue;
const char* val = agg_data + s * 8;
uint8_t af = aflags[a];
if (agg_entry_is_null(af, aflags2[a], asent[a], val))
continue;
if ((af & GHT_AF_BINARY) && (aflags2[a] & GHT_AF2_Y_NULLABLE)) {
if (af & GHT_AF_F64) {
double y; memcpy(&y, agg_data + (s + 1) * 8, 8);
if (y != y) continue;
} else {
int64_t yi; memcpy(&yi, agg_data + (s + 1) * 8, 8);
if (yi == NULL_I64) continue;
}
}
nn[s]++;
if (af & GHT_AF_F64) {
double v; memcpy(&v, val, 8);
if (nf & GHT_NEED_SUM) {
if (af & GHT_AF_FIRST) {
if (entry_row < ROW_RD_I64(row, ly->off_first_row, s)) {
memcpy(row + ly->off_sum + s * 8, val, 8);
ROW_WR_I64(row, ly->off_first_row, s) = entry_row;
}
} else if (af & GHT_AF_LAST) {
if (entry_row > ROW_RD_I64(row, ly->off_last_row, s)) {
memcpy(row + ly->off_sum + s * 8, val, 8);
ROW_WR_I64(row, ly->off_last_row, s) = entry_row;
}
} else if (aflags2[a] & GHT_AF2_TRUTHY) { ROW_WR_I64(row, ly->off_sum, s) += (v != 0.0); }
else if (af & GHT_AF_PROD) { ROW_WR_F64(row, ly->off_sum, s) *= v; }
else { ROW_WR_F64(row, ly->off_sum, s) += v; }
}
if (nf & GHT_NEED_MIN) { double* p = &ROW_WR_F64(row, ly->off_min, s); if (v < *p) *p = v; }
if (nf & GHT_NEED_MAX) { double* p = &ROW_WR_F64(row, ly->off_max, s); if (v > *p) *p = v; }
if (nf & GHT_NEED_SUMSQ) { ROW_WR_F64(row, ly->off_sumsq, s) += v * v; }
if ((nf & GHT_NEED_PEARSON) && (af & GHT_AF_BINARY)) {
double y; memcpy(&y, agg_data + (s + 1) * 8, 8);
ROW_WR_F64(row, ly->off_sum_y, s) += y;
ROW_WR_F64(row, ly->off_sumsq_y, s) += y * y;
ROW_WR_F64(row, ly->off_sumxy, s) += v * y;
}
} else {
int64_t v; memcpy(&v, val, 8);
if (nf & GHT_NEED_SUM) {
if (af & GHT_AF_FIRST) {
if (entry_row < ROW_RD_I64(row, ly->off_first_row, s)) {
memcpy(row + ly->off_sum + s * 8, val, 8);
ROW_WR_I64(row, ly->off_first_row, s) = entry_row;
}
} else if (af & GHT_AF_LAST) {
if (entry_row > ROW_RD_I64(row, ly->off_last_row, s)) {
memcpy(row + ly->off_sum + s * 8, val, 8);
ROW_WR_I64(row, ly->off_last_row, s) = entry_row;
}
} else if (aflags2[a] & GHT_AF2_TRUTHY) { ROW_WR_I64(row, ly->off_sum, s) += (v != 0); }
else if (af & GHT_AF_PROD) { ROW_WR_I64(row, ly->off_sum, s) = (int64_t)((uint64_t)ROW_RD_I64(row, ly->off_sum, s) * (uint64_t)v); }
else if (nf & GHT_NEED_SUM128) { ray_i128_add(&ROW_WR_I64(row, ly->off_sum_hi, s), (uint64_t*)&ROW_WR_I64(row, ly->off_sum, s), v); }
else { ROW_WR_I64(row, ly->off_sum, s) = wrap_add_i64(ROW_RD_I64(row, ly->off_sum, s), v); }
}
if (nf & GHT_NEED_MIN) {
int64_t* p = &ROW_WR_I64(row, ly->off_min, s);
if (af & GHT_AF_SYM) { if (*p == INT64_MAX || sym_lex_lt(adom[a], v, *p)) *p = v; }
else if (v < *p) *p = v;
}
if (nf & GHT_NEED_MAX) {
int64_t* p = &ROW_WR_I64(row, ly->off_max, s);
if (af & GHT_AF_SYM) { if (*p == INT64_MIN || sym_lex_gt(adom[a], v, *p)) *p = v; }
else if (v > *p) *p = v;
}
if (nf & GHT_NEED_SUMSQ) { ROW_WR_F64(row, ly->off_sumsq, s) += (double)v * (double)v; }
if ((nf & GHT_NEED_PEARSON) && (af & GHT_AF_BINARY)) {
int64_t yi; memcpy(&yi, agg_data + (s + 1) * 8, 8);
double y = (double)yi; double xd = (double)v;
ROW_WR_F64(row, ly->off_sum_y, s) += y;
ROW_WR_F64(row, ly->off_sumsq_y, s) += y * y;
ROW_WR_F64(row, ly->off_sumxy, s) += xd * y;
}
}
}
}
static void init_accum_from_entry_nullable(char* row, const char* entry,
const ght_layout_t* ly) {
uint16_t accum_start = (uint16_t)(8 + ly->key_region);
if (ly->row_stride > accum_start)
memset(row + accum_start, 0, ly->row_stride - accum_start);
uint16_t na = ly->n_aggs;
uint8_t nf = ly->need_flags;
bool has_fl = (ly->agg_flags_any & (GHT_AF_FIRST | GHT_AF_LAST)) != 0;
const int8_t* const vslot = ly->agg_val_slot;
const uint8_t* const aflags = ly->agg_flags;
for (uint32_t a = 0; a < na; a++) {
int8_t s = vslot[a];
if (s < 0) continue;
uint8_t af = aflags[a];
if (nf & GHT_NEED_MIN) {
if (af & GHT_AF_F64) ROW_WR_F64(row, ly->off_min, s) = DBL_MAX;
else ROW_WR_I64(row, ly->off_min, s) = INT64_MAX;
}
if (nf & GHT_NEED_MAX) {
if (af & GHT_AF_F64) ROW_WR_F64(row, ly->off_max, s) = -DBL_MAX;
else ROW_WR_I64(row, ly->off_max, s) = INT64_MIN;
}
if ((nf & GHT_NEED_SUM) && (af & GHT_AF_PROD)) {
if (af & GHT_AF_F64) ROW_WR_F64(row, ly->off_sum, s) = 1.0;
else ROW_WR_I64(row, ly->off_sum, s) = 1;
}
if (has_fl) {
ROW_WR_I64(row, ly->off_first_row, s) = INT64_MAX;
ROW_WR_I64(row, ly->off_last_row, s) = INT64_MIN;
}
}
accum_from_entry_nullable(row, entry, ly);
}
#define GHT_LANES_INLINE 8
#if defined(DEBUG) || defined(RAY_HARDENED)
#define GHT_LANES_WHOLE(bytes) assert((((unsigned)(bytes)) & 7u) == 0 && \
"ght_lanes_* length must be a whole number of 8-byte lanes")
#else
#define GHT_LANES_WHOLE(bytes) ((void)0)
#endif
static inline void ght_lanes_copy(void* dst, const void* src, uint16_t bytes) {
uint64_t* d = (uint64_t*)dst;
const uint64_t* s = (const uint64_t*)src;
GHT_LANES_WHOLE(bytes);
switch ((unsigned)bytes >> 3) {
default: memcpy(dst, src, bytes); return;
case 8: d[7] = s[7]; __attribute__((fallthrough));
case 7: d[6] = s[6]; __attribute__((fallthrough));
case 6: d[5] = s[5]; __attribute__((fallthrough));
case 5: d[4] = s[4]; __attribute__((fallthrough));
case 4: d[3] = s[3]; __attribute__((fallthrough));
case 3: d[2] = s[2]; __attribute__((fallthrough));
case 2: d[1] = s[1]; __attribute__((fallthrough));
case 1: d[0] = s[0]; __attribute__((fallthrough));
case 0: return;
}
}
static inline bool ght_lanes_equal(const void* a, const void* b, uint16_t bytes) {
const uint64_t* x = (const uint64_t*)a;
const uint64_t* y = (const uint64_t*)b;
GHT_LANES_WHOLE(bytes);
switch ((unsigned)bytes >> 3) {
default: return memcmp(a, b, bytes) == 0;
case 8: if (x[7] != y[7]) return false; __attribute__((fallthrough));
case 7: if (x[6] != y[6]) return false; __attribute__((fallthrough));
case 6: if (x[5] != y[5]) return false; __attribute__((fallthrough));
case 5: if (x[4] != y[4]) return false; __attribute__((fallthrough));
case 4: if (x[3] != y[3]) return false; __attribute__((fallthrough));
case 3: if (x[2] != y[2]) return false; __attribute__((fallthrough));
case 2: if (x[1] != y[1]) return false; __attribute__((fallthrough));
case 1: if (x[0] != y[0]) return false; __attribute__((fallthrough));
case 0: return true;
}
}
static inline bool group_keys_equal(const int64_t* a_keys, const int64_t* b_keys,
const ght_layout_t* ly, void* const* key_data,
const void* const* key_pool) {
uint16_t nk = ly->n_keys;
if (!ly->any_wide_key) {
return ght_lanes_equal(a_keys, b_keys, ly->key_region);
}
const uint8_t* const kflags = ly->key_flags;
const uint16_t* const koff = ly->key_off;
if (ly->any_inline_str) {
const char* a = (const char*)a_keys;
const char* b = (const char*)b_keys;
for (uint32_t k = 0; k < nk; k++) {
if (kflags[k] & GHT_KEYF_INLINE_STR) {
if (!inline_str_eq(ly, k, a_keys, b_keys, key_pool)) return false;
} else if (kflags[k] & GHT_KEYF_WIDE) {
int64_t ra = *(const int64_t*)(a + koff[k]);
int64_t rb = *(const int64_t*)(b + koff[k]);
if (!wide_key_eq_at(ly, k, key_data, key_pool, ra, rb)) return false;
} else {
if (*(const int64_t*)(a + koff[k]) !=
*(const int64_t*)(b + koff[k])) return false;
}
}
return memcmp(a + koff[nk], b + koff[nk],
(size_t)ly->null_words * 8) == 0;
}
for (uint32_t k = 0; k < nk; k++) {
if (kflags[k] & GHT_KEYF_WIDE) {
if (!wide_key_eq_at(ly, k, key_data, key_pool, a_keys[k], b_keys[k]))
return false;
} else {
if (a_keys[k] != b_keys[k]) return false;
}
}
if (memcmp(a_keys + nk, b_keys + nk, (size_t)ly->null_words * 8) != 0)
return false;
return true;
}
static inline uint32_t group_probe_entry(group_ht_t* ht,
const char* entry, const int8_t* key_types, uint32_t mask) {
const ght_layout_t* ly = &ht->layout;
uint64_t hash = *(const uint64_t*)entry;
const char* ekeys = entry + 8;
uint8_t salt = HT_SALT(hash);
int64_t entry_row;
memcpy(&entry_row, entry + ly->entry_stride - 8, 8);
uint32_t slot = (uint32_t)(hash & mask);
uint16_t key_bytes = ly->key_region;
uint8_t accum_skip = (ly->need_flags == 0
&& (ly->agg_flags_any & (GHT_AF_FIRST | GHT_AF_LAST | GHT_AF_BINARY)) == 0);
for (;;) {
uint32_t sv = ht->slots[slot];
if (sv == HT_EMPTY) {
if (ht->grp_count >= ht->grp_cap) {
if (!group_ht_grow(ht)) { ht->oom = 1; return mask; }
}
uint32_t gid = ht->grp_count++;
char* row = ht->rows + (size_t)gid * ly->row_stride;
*(int64_t*)row = 1;
ght_lanes_copy(row + 8, ekeys, key_bytes);
if (!accum_skip)
init_accum_from_entry(row, entry, ly);
else if (ly->row_stride > 8 + key_bytes)
memset(row + 8 + key_bytes, 0,
ly->row_stride - (8 + key_bytes));
memcpy(row + ly->off_group_first, &entry_row, 8);
ht->slots[slot] = HT_PACK(salt, gid);
if (ht->grp_count * 2 > ht->ht_cap) {
group_ht_rehash(ht, key_types);
mask = ht->ht_cap - 1;
}
return mask;
}
if (HT_SALT_V(sv) == salt) {
uint32_t gid = HT_GID(sv);
char* row = ht->rows + (size_t)gid * ly->row_stride;
if (group_keys_equal((const int64_t*)(row + 8),
(const int64_t*)ekeys, ly, ht->key_data, ht->key_pool)) {
(*(int64_t*)row)++;
int64_t first;
memcpy(&first, row + ly->off_group_first, 8);
if (entry_row < first)
memcpy(row + ly->off_group_first, &entry_row, 8);
if (!accum_skip)
accum_from_entry(row, entry, ly);
return mask;
}
}
slot = (slot + 1) & mask;
}
}
#define GROUP_PREFETCH_BATCH 16
static inline int64_t group_strlen_at(const ray_t* col, int64_t row);
static inline bool group_rowsel_pass(ray_t* sel, int64_t row) {
if (!sel) return true;
ray_rowsel_t* m = ray_rowsel_meta(sel);
if (row < 0 || row >= m->nrows) return false;
uint32_t seg = (uint32_t)(row / RAY_MORSEL_ELEMS);
uint8_t f = ray_rowsel_flags(sel)[seg];
if (f == RAY_SEL_ALL) return true;
if (f == RAY_SEL_NONE) return false;
uint16_t local = (uint16_t)(row - (int64_t)seg * RAY_MORSEL_ELEMS);
uint32_t lo = ray_rowsel_offsets(sel)[seg];
uint32_t hi = ray_rowsel_offsets(sel)[seg + 1];
const uint16_t* idx = ray_rowsel_idx(sel);
while (lo < hi) {
uint32_t mid = lo + ((hi - lo) >> 1);
uint16_t v = idx[mid];
if (v == local) return true;
if (v < local) lo = mid + 1;
else hi = mid;
}
return false;
}
void group_rows_range(group_ht_t* ht, void** key_data, int8_t* key_types,
uint8_t* key_attrs, ray_t** key_vecs, ray_t** agg_vecs,
ray_t** agg_vecs2,
uint8_t* agg_strlen,
ray_t* rowsel,
int64_t start, int64_t end,
const int64_t* match_idx) {
const ght_layout_t* ly = &ht->layout;
uint16_t nk = ly->n_keys;
uint16_t na = ly->n_aggs;
const uint8_t* const kflags = ly->key_flags;
const uint8_t* const aflags = ly->agg_flags;
uint8_t wide_any = ly->any_wide_key;
uint32_t mask = ht->ht_cap - 1;
char ebuf_stk[8 + 9 * 8 + 8 * 8 + 8];
char* ebuf = ebuf_stk;
ray_t* ebuf_hdr = NULL;
if (ly->entry_stride > sizeof(ebuf_stk)) {
ebuf = (char*)scratch_alloc(&ebuf_hdr, ly->entry_stride);
if (!ebuf) { ht->oom = 1; return; }
}
uint8_t any_nullable = 0;
for (uint32_t k = 0; k < nk; k++)
if (key_vecs && ray_key_may_be_null(key_vecs[k])) { any_nullable = 1; break; }
if (wide_any) { group_ht_set_key_data(ht, key_data); group_ht_set_key_pool(ht, key_vecs); }
for (int64_t i = start; i < end; i++) {
if (((i - start) & 65535) == 0 && ray_interrupted()) break;
int64_t row = match_idx ? match_idx[i] : i;
if (!match_idx && rowsel && !group_rowsel_pass(rowsel, row)) continue;
uint64_t h = 0;
int64_t* ek = (int64_t*)(ebuf + 8);
if (ly->any_inline_str) {
h = inline_build_keys(ly, key_types, key_data, key_attrs, ht->key_pool,
key_vecs, any_nullable, row, ebuf + 8);
} else {
int64_t* nullw = ek + nk;
uint32_t null_words = ly->null_words;
for (uint32_t w = 0; w < null_words; w++) nullw[w] = 0;
if (ly->packed_key) {
for (uint32_t k = 0; k < nk; k++) {
if (__builtin_expect(any_nullable && ray_key_may_be_null(key_vecs[k])
&& ray_vec_is_null(key_vecs[k], row), 0)) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
ek[k] = 0;
} else {
ek[k] = read_col_i64(key_data[k], row, key_types[k], key_attrs[k]);
}
}
h = ght_hash_lanes(ek, (uint32_t)nk + null_words);
} else {
for (uint32_t k = 0; k < nk; k++) {
int8_t t = key_types[k];
uint64_t kh;
bool is_null = any_nullable && ray_key_may_be_null(key_vecs[k])
&& ray_vec_is_null(key_vecs[k], row);
if (is_null) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
ek[k] = 0;
kh = ray_hash_i64(0);
} else if (wide_any && (kflags[k] & GHT_KEYF_WIDE)) {
ek[k] = row;
kh = wide_key_hash_at(ly, k, key_data, ht->key_pool, row);
} else if (t == RAY_F64) {
int64_t kv = group_key_f64_bits(key_data[k], row);
double dv; memcpy(&dv, &kv, 8);
ek[k] = kv;
kh = ray_hash_f64(dv);
} else {
int64_t kv = read_col_i64(key_data[k], row, t, key_attrs[k]);
ek[k] = kv;
kh = ray_hash_i64(kv);
}
h = (k == 0) ? kh : ray_hash_combine(h, kh);
}
h = ght_hash_null_words(h, nullw, null_words);
}
}
*(uint64_t*)ebuf = h;
int64_t* ev = (int64_t*)(ebuf + 8 + (size_t)ly->key_region);
for (uint32_t a = 0; a < na; a++) {
uint8_t af = aflags[a];
int8_t vs = ly->agg_val_slot[a];
if (vs < 0) continue;
uint8_t vi = (uint8_t)vs;
ray_t* ac = agg_vecs[a];
if (!ac) continue;
if (agg_strlen && agg_strlen[a])
ev[vi] = group_strlen_at(ac, row);
else if (af & GHT_AF_F64) {
double v = group_fp_type(ac->type)
? group_fp_at(ray_data(ac), ac->type, row)
: group_pack_i64_as_f64(ac, row);
memcpy(&ev[vi], &v, sizeof(v));
}
else
ev[vi] = read_col_i64(ray_data(ac), row, ac->type, ac->attrs);
vi++;
if ((af & GHT_AF_BINARY) && agg_vecs2 && agg_vecs2[a]) {
ray_t* ay = agg_vecs2[a];
if (af & GHT_AF_F64) {
double v = group_fp_type(ay->type)
? group_fp_at(ray_data(ay), ay->type, row)
: group_pack_i64_as_f64(ay, row);
memcpy(&ev[vi], &v, sizeof(v));
}
else
ev[vi] = group_pack_y_i64(ay, row);
vi++;
}
}
memcpy(ebuf + ly->entry_stride - 8, &row, 8);
mask = group_probe_entry(ht, ebuf, key_types, mask);
}
scratch_free(ebuf_hdr);
}
static uint32_t group_radix_part_count(uint32_t n_workers, int64_t n_rows) {
uint32_t n_parts = 1;
uint64_t rows = n_rows > 0 ? (uint64_t)n_rows : 1;
while ((n_parts < n_workers ||
(uint64_t)n_parts < rows / n_parts + (rows % n_parts != 0)) &&
n_parts <= UINT32_MAX / 2)
n_parts <<= 1;
return n_parts;
}
static inline uint32_t group_radix_part(uint64_t hash, uint32_t n_parts) {
return ((uint32_t)(hash >> 16)) & (n_parts - 1);
}
typedef struct {
int64_t value;
uint32_t part;
uint32_t gid;
} group_topn_item_t;
#define TOPN_BETTER(desc_dir, a, b) ((desc_dir) ? ((a) > (b)) : ((a) < (b)))
static inline int64_t topn_heap_push(group_topn_item_t* h, int64_t hn,
int64_t k, uint8_t desc_dir, int64_t v,
uint32_t part, uint32_t gid) {
if (hn < k) {
int64_t j = hn++;
h[j] = (group_topn_item_t){v, part, gid};
while (j > 0) {
int64_t pr = (j - 1) >> 1;
if (!TOPN_BETTER(desc_dir, h[pr].value, h[j].value)) break;
group_topn_item_t t = h[pr]; h[pr] = h[j]; h[j] = t;
j = pr;
}
return hn;
}
if (!TOPN_BETTER(desc_dir, v, h[0].value)) return hn;
h[0] = (group_topn_item_t){v, part, gid};
for (int64_t j = 0;;) {
int64_t l = j * 2 + 1, r = l + 1, m = j;
if (l < hn && TOPN_BETTER(desc_dir, h[m].value, h[l].value)) m = l;
if (r < hn && TOPN_BETTER(desc_dir, h[m].value, h[r].value)) m = r;
if (m == j) break;
group_topn_item_t t = h[m]; h[m] = h[j]; h[j] = t;
j = m;
}
return hn;
}
static int topn_gid_cmp(const void* a, const void* b) {
uint32_t ga = ((const group_topn_item_t*)a)->gid;
uint32_t gb = ((const group_topn_item_t*)b)->gid;
return ga < gb ? -1 : (ga > gb ? 1 : 0);
}
static int topn_part_gid_cmp(const void* a, const void* b) {
const group_topn_item_t* x = (const group_topn_item_t*)a;
const group_topn_item_t* y = (const group_topn_item_t*)b;
if (x->part != y->part) return x->part < y->part ? -1 : 1;
return x->gid < y->gid ? -1 : (x->gid > y->gid ? 1 : 0);
}
typedef struct {
group_ht_t* part_hts;
group_topn_item_t* items;
const uint32_t* item_off;
uint32_t* item_cnt;
uint16_t order_off;
uint8_t desc_dir;
} topn_scan_ctx_t;
static void topn_scan_fn(void* vctx, uint32_t worker_id, int64_t start,
int64_t end) {
(void)worker_id;
topn_scan_ctx_t* c = (topn_scan_ctx_t*)vctx;
for (int64_t pi = start; pi < end; pi++) {
uint32_t p = (uint32_t)pi;
group_ht_t* ph = &c->part_hts[p];
uint32_t gc = ph->grp_count;
uint16_t rs = ph->layout.row_stride;
group_topn_item_t* h = c->items + c->item_off[p];
int64_t cap = (int64_t)(c->item_off[p + 1] - c->item_off[p]);
int64_t hn = 0;
for (uint32_t gi = 0; gi < gc; gi++) {
int64_t v = *(const int64_t*)(const void*)
(ph->rows + (size_t)gi * rs + c->order_off);
hn = topn_heap_push(h, hn, cap, c->desc_dir, v, p, gi);
}
if (hn > 1) qsort(h, (size_t)hn, sizeof(*h), topn_gid_cmp);
c->item_cnt[p] = (uint32_t)hn;
}
}
typedef struct {
group_topn_item_t* items;
uint32_t* item_cnt;
uint32_t cap;
uint16_t order_off;
uint8_t desc_dir;
} topn_fuse_t;
static inline void topn_fuse_partition(const topn_fuse_t* f, group_ht_t* ph,
uint32_t p) {
uint32_t gc = ph->grp_count;
uint16_t rs = ph->layout.row_stride;
group_topn_item_t* h = f->items + (size_t)p * f->cap;
int64_t cap = (int64_t)f->cap;
int64_t hn = 0;
for (uint32_t gi = 0; gi < gc; gi++) {
int64_t v = *(const int64_t*)(const void*)
(ph->rows + (size_t)gi * rs + f->order_off);
hn = topn_heap_push(h, hn, cap, f->desc_dir, v, p, gi);
}
if (hn > 1) qsort(h, (size_t)hn, sizeof(*h), topn_gid_cmp);
f->item_cnt[p] = (uint32_t)hn;
}
typedef struct {
group_ht_t* part_hts;
const int8_t* key_types;
const uint8_t* dirty;
} topn_rebuild_ctx_t;
static void topn_rebuild_fn(void* vctx, uint32_t worker_id, int64_t start,
int64_t end) {
(void)worker_id;
topn_rebuild_ctx_t* c = (topn_rebuild_ctx_t*)vctx;
for (int64_t p = start; p < end; p++)
if (c->dirty[p]) group_ht_rebuild_slots(&c->part_hts[p], c->key_types);
}
typedef struct {
int64_t first;
uint32_t flat;
} group_order_pair_t;
static int group_order_pair_cmp(const void* a, const void* b) {
int64_t fa = ((const group_order_pair_t*)a)->first;
int64_t fb = ((const group_order_pair_t*)b)->first;
return (fa > fb) - (fa < fb);
}
#define SEL_MATCH_GATE_SHIFT 1
#define P1_MORSEL 256u
#define P1_MORSEL_KEYS 8u
typedef struct {
char* data;
uint32_t count;
uint32_t cap;
bool oom;
ray_t* _hdr;
} radix_buf_t;
static inline void radix_buf_push(radix_buf_t* buf, uint16_t entry_stride,
uint64_t hash, const int64_t* key_region_buf,
const int64_t* agg_vals, uint16_t n_agg_vals,
int64_t row, uint16_t key_region,
uint32_t cap0) {
if (__builtin_expect(buf->count >= buf->cap, 0)) {
uint32_t old_cap = buf->cap;
if (old_cap > UINT32_MAX / 2) { buf->oom = true; return; }
uint32_t new_cap = old_cap ? old_cap * 2 : (cap0 ? cap0 : 1);
char* new_data = (char*)scratch_realloc(
&buf->_hdr, (size_t)old_cap * entry_stride,
(size_t)new_cap * entry_stride);
if (!new_data) { buf->oom = true; return; }
buf->data = new_data;
buf->cap = new_cap;
}
char* dst = buf->data + (size_t)buf->count * entry_stride;
*(uint64_t*)dst = hash;
ght_lanes_copy(dst + 8, key_region_buf, key_region);
if (n_agg_vals)
ght_lanes_copy(dst + 8 + key_region, agg_vals, (uint16_t)(n_agg_vals * 8));
memcpy(dst + entry_stride - 8, &row, 8);
buf->count++;
}
typedef struct {
void** key_data;
const void** key_pool;
int8_t* key_types;
uint8_t* key_attrs;
ray_t** key_vecs;
uint8_t nullable_mask;
ray_t** agg_vecs;
ray_t** agg_vecs2;
uint8_t* agg_strlen;
uint32_t n_workers;
uint32_t n_parts;
radix_buf_t* bufs;
ght_layout_t layout;
ray_t* rowsel;
const int64_t* match_idx;
uint32_t buf_prime;
} radix_phase1_ctx_t;
static inline void radix_phase1_pack_aggs(const radix_phase1_ctx_t* c,
const ght_layout_t* ly,
int64_t* agg_vals, int64_t row) {
const uint8_t* const aflags = ly->agg_flags;
uint16_t na = ly->n_aggs;
for (uint32_t a = 0; a < na; a++) {
uint8_t af = aflags[a];
int8_t vs = ly->agg_val_slot[a];
if (vs < 0) continue;
uint8_t vi = (uint8_t)vs;
ray_t* ac = c->agg_vecs[a];
if (!ac) continue;
if (c->agg_strlen && c->agg_strlen[a])
agg_vals[vi] = group_strlen_at(ac, row);
else if (af & GHT_AF_F64) {
double v = group_fp_type(ac->type)
? group_fp_at(ray_data(ac), ac->type, row)
: group_pack_i64_as_f64(ac, row);
memcpy(&agg_vals[vi], &v, sizeof(v));
}
else
agg_vals[vi] = read_col_i64(ray_data(ac), row, ac->type, ac->attrs);
vi++;
if ((af & GHT_AF_BINARY) && c->agg_vecs2 && c->agg_vecs2[a]) {
ray_t* ay = c->agg_vecs2[a];
if (af & GHT_AF_F64) {
double v = group_fp_type(ay->type)
? group_fp_at(ray_data(ay), ay->type, row)
: group_pack_i64_as_f64(ay, row);
memcpy(&agg_vals[vi], &v, sizeof(v));
}
else
agg_vals[vi] = group_pack_y_i64(ay, row);
vi++;
}
}
}
static void radix_phase1_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
radix_phase1_ctx_t* c = (radix_phase1_ctx_t*)ctx;
const ght_layout_t* ly = &c->layout;
radix_buf_t* my_bufs = &c->bufs[(size_t)worker_id * c->n_parts];
uint16_t nk = ly->n_keys;
uint16_t nv = ly->n_agg_vals;
const uint8_t* const kflags = ly->key_flags;
uint8_t wide_any = ly->any_wide_key;
uint8_t inline_str = ly->any_inline_str;
uint16_t estride = ly->entry_stride;
const int64_t* match_idx = c->match_idx;
int64_t keys_stk[9];
int64_t agg_vals_stk[8];
char keybuf_stk[136];
int64_t* keys = keys_stk;
int64_t* agg_vals = agg_vals_stk;
char* keybuf = keybuf_stk;
ray_t* stage_hdr = NULL;
if (ly->key_region > sizeof(keys_stk) || nv > 8) {
size_t kb = ly->key_region;
char* blk = (char*)scratch_alloc(&stage_hdr, kb + kb + (size_t)nv * 8);
if (!blk) return;
keybuf = blk;
keys = (int64_t*)(blk + kb);
agg_vals = (int64_t*)(blk + kb + kb);
}
uint8_t nullable = c->nullable_mask;
uint8_t packed = ly->packed_key;
if (packed && !nullable && !inline_str && ly->null_words == 0 &&
nk >= 1 && nk <= P1_MORSEL_KEYS) {
int64_t rowbuf[P1_MORSEL];
static_assert(sizeof(int64_t) * P1_MORSEL * P1_MORSEL_KEYS <= 64u * 1024u,
"phase1 morsel stage must stay a modest worker-stack block");
int64_t stage[P1_MORSEL_KEYS][P1_MORSEL];
bool contiguous = !match_idx && !c->rowsel;
for (int64_t i = start; i < end; ) {
if (ray_interrupted()) break;
uint32_t m = 0;
for (; i < end && m < P1_MORSEL; i++) {
int64_t row = match_idx ? match_idx[i] : i;
if (!match_idx && c->rowsel && !group_rowsel_pass(c->rowsel, row))
continue;
rowbuf[m++] = row;
}
if (m == 0) continue;
for (uint32_t k = 0; k < nk; k++) {
if (contiguous)
group_keys_gather_seq(stage[k], c->key_data[k], rowbuf[0], m,
c->key_types[k], c->key_attrs[k]);
else
group_keys_gather(stage[k], c->key_data[k], rowbuf, m,
c->key_types[k], c->key_attrs[k]);
}
for (uint32_t j = 0; j < m; j++) {
int64_t row = rowbuf[j];
for (uint32_t k = 0; k < nk; k++) keys[k] = stage[k][j];
uint64_t h = ght_hash_lanes(keys, nk);
radix_phase1_pack_aggs(c, ly, agg_vals, row);
uint32_t part = group_radix_part(h, c->n_parts);
radix_buf_push(&my_bufs[part], estride, h, keys,
agg_vals, nv, row, ly->key_region, c->buf_prime);
}
}
scratch_free(stage_hdr);
return;
}
for (int64_t i = start; i < end; i++) {
if (((i - start) & 65535) == 0 && ray_interrupted()) break;
int64_t row = match_idx ? match_idx[i] : i;
if (!match_idx && c->rowsel && !group_rowsel_pass(c->rowsel, row)) continue;
uint64_t h = 0;
if (inline_str) {
h = inline_build_keys(ly, c->key_types, c->key_data, c->key_attrs,
c->key_pool, c->key_vecs, nullable, row,
keybuf);
} else {
int64_t* nullw = keys + nk;
uint32_t null_words = ly->null_words;
for (uint32_t w = 0; w < null_words; w++) nullw[w] = 0;
if (packed) {
for (uint32_t k = 0; k < nk; k++) {
if (__builtin_expect(nullable && ray_key_may_be_null(c->key_vecs[k])
&& ray_vec_is_null(c->key_vecs[k], row), 0)) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
keys[k] = 0;
} else {
keys[k] = read_col_i64(c->key_data[k], row, c->key_types[k],
c->key_attrs[k]);
}
}
h = ght_hash_lanes(keys, (uint32_t)nk + null_words);
} else {
for (uint32_t k = 0; k < nk; k++) {
int8_t t = c->key_types[k];
uint64_t kh;
bool is_null = nullable && ray_key_may_be_null(c->key_vecs[k])
&& ray_vec_is_null(c->key_vecs[k], row);
if (is_null) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
keys[k] = 0;
kh = ray_hash_i64(0);
} else if (wide_any && (kflags[k] & GHT_KEYF_WIDE)) {
keys[k] = row;
kh = wide_key_hash_at(ly, k, c->key_data, c->key_pool, row);
} else if (t == RAY_F64) {
int64_t kv = group_key_f64_bits(c->key_data[k], row);
double dv; memcpy(&dv, &kv, 8);
keys[k] = kv;
kh = ray_hash_f64(dv);
} else {
int64_t kv = read_col_i64(c->key_data[k], row, t, c->key_attrs[k]);
keys[k] = kv;
kh = ray_hash_i64(kv);
}
h = (k == 0) ? kh : ray_hash_combine(h, kh);
}
h = ght_hash_null_words(h, nullw, null_words);
}
}
radix_phase1_pack_aggs(c, ly, agg_vals, row);
uint32_t part = group_radix_part(h, c->n_parts);
radix_buf_push(&my_bufs[part], estride, h,
inline_str ? (const int64_t*)keybuf : keys,
agg_vals, nv, row, ly->key_region, c->buf_prime);
}
scratch_free(stage_hdr);
}
static void group_rows_indirect(group_ht_t* ht, const int8_t* key_types,
const char* entries, uint32_t n_entries,
uint16_t entry_stride) {
uint32_t mask = ht->ht_cap - 1;
enum { PF_DIST = 8 };
uint32_t pf_end = (n_entries < PF_DIST) ? n_entries : PF_DIST;
for (uint32_t j = 0; j < pf_end; j++) {
uint64_t h = *(const uint64_t*)(entries + (size_t)j * entry_stride);
__builtin_prefetch(&ht->slots[(uint32_t)(h & mask)], 0, 1);
}
for (uint32_t i = 0; i < n_entries; i++) {
if (i + PF_DIST < n_entries) {
uint64_t h = *(const uint64_t*)(entries + (size_t)(i + PF_DIST) * entry_stride);
__builtin_prefetch(&ht->slots[(uint32_t)(h & mask)], 0, 1);
}
const char* e = entries + (size_t)i * entry_stride;
mask = group_probe_entry(ht, e, key_types, mask);
}
}
typedef struct {
int8_t out_type;
bool src_f64;
uint16_t agg_op;
bool affine;
double bias_f64;
int64_t bias_i64;
void* dst;
ray_t* vec;
} agg_out_t;
#define grp_set_null par_set_null
#define grp_finalize_nulls par_finalize_nulls
typedef struct {
group_ht_t* part_hts;
uint32_t* part_offsets;
const uint32_t* group_out;
char** key_dsts;
int8_t* key_types;
uint8_t* key_attrs;
uint8_t* key_esizes;
uint32_t n_keys;
agg_out_t* agg_outs;
uint32_t n_aggs;
void** key_src_data;
} radix_phase3_ctx_t;
static void radix_phase3_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
(void)worker_id;
radix_phase3_ctx_t* c = (radix_phase3_ctx_t*)ctx;
uint32_t nk = c->n_keys;
uint32_t na = c->n_aggs;
for (int64_t p = start; p < end; p++) {
group_ht_t* ph = &c->part_hts[p];
uint32_t gc = ph->grp_count;
if (gc == 0) continue;
uint32_t off = c->part_offsets[p];
const ght_layout_t* ly = &ph->layout;
uint16_t rs = ly->row_stride;
const uint16_t* const koff = ly->key_off;
const uint8_t* const kflags = ly->key_flags;
const uint8_t* const aflags = ly->agg_flags;
const int8_t* const vslot = ly->agg_val_slot;
for (uint32_t gi = 0; gi < gc; gi++) {
const char* row = ph->rows + (size_t)gi * rs;
const char* rk = row + 8;
int64_t cnt = *(const int64_t*)(const void*)row;
const int64_t* nullw = ght_null_words_at(ly, rk);
const int64_t* nnbase = ly->off_nn
? (const int64_t*)(const void*)(row + ly->off_nn) : NULL;
uint32_t di = c->group_out ? c->group_out[off + gi] : off + gi;
for (uint32_t k = 0; k < nk; k++) {
char* dst = c->key_dsts[k];
uint8_t esz = c->key_esizes[k];
int8_t kt = c->key_types[k];
size_t doff = (size_t)di * esz;
if (nullw[k >> 6] & ((int64_t)((uint64_t)1 << (k & 63)))) {
switch (kt) {
case RAY_F64: {
double v = NULL_F64; memcpy(dst + doff, &v, 8); break;
}
case RAY_F32: {
float v = NULL_F32; memcpy(dst + doff, &v, 4); break;
}
case RAY_I64: case RAY_TIMESTAMP: {
int64_t v = NULL_I64; memcpy(dst + doff, &v, 8); break;
}
case RAY_I32: case RAY_DATE: case RAY_TIME: {
int32_t v = NULL_I32; memcpy(dst + doff, &v, 4); break;
}
case RAY_I16: {
int16_t v = NULL_I16; memcpy(dst + doff, &v, 2); break;
}
default:
memset(dst + doff, 0, esz);
break;
}
continue;
}
if (kflags[k] & GHT_KEYF_INLINE_STR) {
memcpy(dst + doff, rk + koff[k], sizeof(ray_str_t));
continue;
}
int64_t kv = *(const int64_t*)(const void*)(rk + koff[k]);
if (kflags[k] & GHT_KEYF_WIDE) {
const char* src = (const char*)c->key_src_data[k];
memcpy(dst + doff, src + (size_t)kv * esz, esz);
} else if (kt == RAY_F64) {
memcpy(dst + doff, &kv, 8);
} else {
write_col_i64(dst, di, kv, kt, c->key_attrs[k]);
}
}
for (uint32_t a = 0; a < na; a++) {
if (aflags[a] & GHT_AF_HOLISTIC) continue;
agg_out_t* ao = &c->agg_outs[a];
if (!ao->dst) continue;
uint16_t op = ao->agg_op;
bool sf = ao->src_f64;
int8_t s = vslot[a];
int64_t nn = (nnbase && s >= 0) ? nnbase[s] : cnt;
if (ao->out_type == RAY_F64) {
double v;
switch (op) {
case OP_SUM:
v = sf ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
if (ao->affine) v += ao->bias_f64 * cnt;
break;
case OP_PROD:
if (nn == 0) { v = NULL_F64; grp_set_null(ao->vec, di); break; }
v = sf ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
break;
case OP_AVG:
if (nn == 0) { v = NULL_F64; grp_set_null(ao->vec, di); break; }
v = sf ? ROW_RD_F64(row, ly->off_sum, s) / nn
: (ly->need_flags & GHT_NEED_SUM128)
? ray_i128_to_f64(ROW_RD_I64(row, ly->off_sum_hi, s),
(uint64_t)ROW_RD_I64(row, ly->off_sum, s)) / nn
: (double)ROW_RD_I64(row, ly->off_sum, s) / nn;
if (ao->affine) v += ao->bias_f64;
break;
case OP_MIN:
if (nn == 0) { v = NULL_F64; grp_set_null(ao->vec, di); break; }
v = sf ? ROW_RD_F64(row, ly->off_min, s)
: (double)ROW_RD_I64(row, ly->off_min, s);
break;
case OP_MAX:
if (nn == 0) { v = NULL_F64; grp_set_null(ao->vec, di); break; }
v = sf ? ROW_RD_F64(row, ly->off_max, s)
: (double)ROW_RD_I64(row, ly->off_max, s);
break;
case OP_FIRST: case OP_LAST:
if (nn == 0) { v = NULL_F64; grp_set_null(ao->vec, di); break; }
v = sf ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
break;
case OP_VAR: case OP_VAR_POP:
case OP_STDDEV: case OP_STDDEV_POP: {
bool insuf = (op == OP_VAR || op == OP_STDDEV) ? nn <= 1 : nn <= 0;
if (insuf) { v = NULL_F64; grp_set_null(ao->vec, di); break; }
double sum_val = sf ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
double sq_val = ly->off_sumsq ? ROW_RD_F64(row, ly->off_sumsq, s) : 0.0;
double mean = sum_val / nn;
double var_pop = sq_val / nn - mean * mean;
if (var_pop < 0) var_pop = 0;
if (op == OP_VAR_POP) v = var_pop;
else if (op == OP_VAR) v = var_pop * nn / (nn - 1);
else if (op == OP_STDDEV_POP) v = sqrt(var_pop);
else v = sqrt(var_pop * nn / (nn - 1));
break;
}
case OP_PEARSON_CORR:
case OP_COV:
case OP_SCOV:
case OP_WSUM:
case OP_WAVG: {
if ((op == OP_PEARSON_CORR || op == OP_SCOV) && nn < 2) { v = NULL_F64; grp_set_null(ao->vec, di); break; }
if ((op == OP_COV || op == OP_WAVG) && nn < 1) { v = NULL_F64; grp_set_null(ao->vec, di); break; }
double sx = sf ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
double sxx = ly->off_sumsq ? ROW_RD_F64(row, ly->off_sumsq, s) : 0.0;
double sy = ly->off_sum_y ? ROW_RD_F64(row, ly->off_sum_y, s) : 0.0;
double syy = ly->off_sumsq_y ? ROW_RD_F64(row, ly->off_sumsq_y, s) : 0.0;
double sxy = ly->off_sumxy ? ROW_RD_F64(row, ly->off_sumxy, s) : 0.0;
double dn = (double)nn;
double num = dn * sxy - sx * sy;
double dx = dn * sxx - sx * sx;
double dy = dn * syy - sy * sy;
if (op == OP_WSUM) { v = sxy; break; }
if (op == OP_WAVG) { if (sx == 0.0) { v = NULL_F64; break; } v = sxy / sx; break; }
if (op == OP_COV) { v = (sxy - sx * sy / dn) / dn; break; }
if (op == OP_SCOV) { v = (sxy - sx * sy / dn) / (dn - 1.0); break; }
if (dx <= 0.0 || dy <= 0.0) { v = NULL_F64; break; }
v = num / sqrt(dx * dy);
break;
}
default: v = 0.0; break;
}
((double*)(void*)ao->dst)[di] = ray_f64_fin(v);
} else {
int64_t v;
int64_t int_null = agg_int_null_sentinel_for(ao->out_type);
switch (op) {
case OP_SUM:
v = ROW_RD_I64(row, ly->off_sum, s);
if (ao->affine) v = wrap_add_i64(v, wrap_mul_i64(ao->bias_i64, cnt));
break;
case OP_PROD:
if (nn == 0) { v = int_null; grp_set_null(ao->vec, di); break; }
v = ROW_RD_I64(row, ly->off_sum, s); break;
case OP_COUNT: v = cnt; break;
case OP_ALL: v = (ROW_RD_I64(row, ly->off_sum, s) == nn) ? 1 : 0; break;
case OP_ANY: v = (ROW_RD_I64(row, ly->off_sum, s) > 0) ? 1 : 0; break;
case OP_MIN:
if (nn == 0) { v = int_null; grp_set_null(ao->vec, di); break; }
v = ROW_RD_I64(row, ly->off_min, s); break;
case OP_MAX:
if (nn == 0) { v = int_null; grp_set_null(ao->vec, di); break; }
v = ROW_RD_I64(row, ly->off_max, s); break;
case OP_FIRST: case OP_LAST:
if (nn == 0) { v = int_null; grp_set_null(ao->vec, di); break; }
v = ROW_RD_I64(row, ly->off_sum, s); break;
default: v = 0; break;
}
topk_write_i64(ao->dst, di, v,
ray_sym_elem_size(ao->out_type, ao->vec->attrs));
}
}
}
}
}
typedef struct {
int8_t* key_types;
uint32_t n_keys;
uint32_t n_workers;
uint32_t n_parts;
radix_buf_t* bufs;
group_ht_t* part_hts;
ght_layout_t layout;
void** key_data;
const void** key_pool;
topn_fuse_t fuse;
} radix_phase2_ctx_t;
static void radix_phase2_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
(void)worker_id;
radix_phase2_ctx_t* c = (radix_phase2_ctx_t*)ctx;
uint16_t estride = c->layout.entry_stride;
for (int64_t p = start; p < end; p++) {
if (c->fuse.items) c->fuse.item_cnt[p] = 0;
uint32_t total = 0;
for (uint32_t w = 0; w < c->n_workers; w++)
total += c->bufs[(size_t)w * c->n_parts + p].count;
if (total == 0) continue;
uint32_t fe_cap = 2;
uint32_t fe_target = total < 4096 ? total : 4096;
while (fe_cap < fe_target) fe_cap <<= 1;
uint32_t fe_blk = total < 256 ? total : 256;
if (!group_ht_init_sized(&c->part_hts[p], fe_cap, &c->layout, fe_blk))
continue;
if (c->layout.any_wide_key && c->key_data) {
group_ht_set_key_data(&c->part_hts[p], c->key_data);
group_ht_copy_key_pool(&c->part_hts[p], c->key_pool);
}
for (uint32_t w = 0; w < c->n_workers; w++) {
radix_buf_t* buf = &c->bufs[(size_t)w * c->n_parts + p];
if (buf->count == 0) continue;
group_rows_indirect(&c->part_hts[p], c->key_types,
buf->data, buf->count, estride);
}
if (c->fuse.items)
topn_fuse_partition(&c->fuse, &c->part_hts[p], (uint32_t)p);
}
}
static inline uint32_t group_merge_row(group_ht_t* ht,
const char* src_row, const int8_t* key_types, uint32_t mask)
{
const ght_layout_t* ly = &ht->layout;
int64_t src_count = *(const int64_t*)src_row;
const int64_t* skeys = (const int64_t*)(src_row + 8);
uint64_t h = hash_keys_inline(skeys, key_types, ly->n_keys,
ht->key_data, ly, ht->key_pool);
uint8_t salt = HT_SALT(h);
uint32_t slot = (uint32_t)(h & mask);
uint16_t na = ly->n_aggs;
const uint8_t* const aflags = ly->agg_flags;
const int8_t* const vslot = ly->agg_val_slot;
uint16_t off_sum = ly->off_sum;
uint16_t off_sum_hi = ly->off_sum_hi;
bool need_sum = (ly->need_flags & GHT_NEED_SUM) != 0;
bool need_sum128 = (ly->need_flags & GHT_NEED_SUM128) != 0;
for (;;) {
uint32_t sv = ht->slots[slot];
if (sv == HT_EMPTY) {
if (ht->grp_count >= ht->grp_cap) {
if (!group_ht_grow(ht)) { ht->oom = 1; return mask; }
}
uint32_t gid = ht->grp_count++;
char* row = ht->rows + (size_t)gid * ly->row_stride;
memcpy(row, src_row, ly->row_stride);
ht->slots[slot] = HT_PACK(salt, gid);
if (ht->grp_count * 2 > ht->ht_cap) {
group_ht_rehash(ht, key_types);
mask = ht->ht_cap - 1;
}
return mask;
}
if (HT_SALT_V(sv) == salt) {
uint32_t gid = HT_GID(sv);
char* row = ht->rows + (size_t)gid * ly->row_stride;
if (group_keys_equal((const int64_t*)(row + 8),
skeys, ly, ht->key_data, ht->key_pool)) {
*(int64_t*)row += src_count;
int64_t dst_first, src_first;
memcpy(&dst_first, row + ly->off_group_first, 8);
memcpy(&src_first, src_row + ly->off_group_first, 8);
if (src_first < dst_first)
memcpy(row + ly->off_group_first, &src_first, 8);
if (need_sum) {
for (uint32_t a = 0; a < na; a++) {
int8_t s = vslot[a];
if (s < 0) continue;
size_t off = (size_t)off_sum + (size_t)s * 8;
if (aflags[a] & GHT_AF_F64) {
double sv_f;
memcpy(&sv_f, src_row + off, 8);
*(double*)(row + off) += sv_f;
} else if (need_sum128) {
int64_t sv_i, sv_hi;
memcpy(&sv_i, src_row + off, 8);
memcpy(&sv_hi, src_row + off_sum_hi + (size_t)s * 8, 8);
ray_i128_add128((int64_t*)(row + off_sum_hi + (size_t)s * 8),
(uint64_t*)(row + off), sv_hi, (uint64_t)sv_i);
} else {
int64_t sv_i;
memcpy(&sv_i, src_row + off, 8);
*(int64_t*)(row + off) = wrap_add_i64(*(int64_t*)(row + off), sv_i);
}
}
}
return mask;
}
}
slot = (slot + 1) & mask;
}
}
typedef struct {
void** key_data;
int8_t* key_types;
uint8_t* key_attrs;
ray_t** key_vecs;
ray_t** agg_vecs;
ray_t** agg_vecs2;
uint8_t* agg_strlen;
uint8_t nullable_mask;
uint32_t n_workers;
uint32_t n_parts;
group_ht_t* wpart_hts;
ght_layout_t layout;
ray_t* rowsel;
const int64_t* match_idx;
uint32_t ht_init_cap;
uint32_t ht_init_block;
uint32_t ht_grow_cap;
_Atomic(int) oom;
} radix_v2_phase1_ctx_t;
static inline void radix_v2_pack_aggs(const radix_v2_phase1_ctx_t* c,
const ght_layout_t* ly,
int64_t* ev, int64_t row) {
const uint8_t* const aflags = ly->agg_flags;
uint16_t na = ly->n_aggs;
for (uint32_t a = 0; a < na; a++) {
uint8_t af = aflags[a];
int8_t vs = ly->agg_val_slot[a];
if (vs < 0) continue;
uint8_t vi = (uint8_t)vs;
ray_t* ac = c->agg_vecs ? c->agg_vecs[a] : NULL;
if (!ac) continue;
if (c->agg_strlen && c->agg_strlen[a])
ev[vi] = group_strlen_at(ac, row);
else if (af & GHT_AF_F64) {
double v = group_fp_type(ac->type)
? group_fp_at(ray_data(ac), ac->type, row)
: group_pack_i64_as_f64(ac, row);
memcpy(&ev[vi], &v, sizeof(v));
}
else
ev[vi] = read_col_i64(ray_data(ac), row, ac->type, ac->attrs);
vi++;
if ((af & GHT_AF_BINARY) && c->agg_vecs2 && c->agg_vecs2[a]) {
ray_t* ay = c->agg_vecs2[a];
if (af & GHT_AF_F64) {
double v = group_fp_type(ay->type)
? group_fp_at(ray_data(ay), ay->type, row)
: group_pack_i64_as_f64(ay, row);
memcpy(&ev[vi], &v, sizeof(v));
}
else
ev[vi] = group_pack_y_i64(ay, row);
vi++;
}
}
}
static void radix_v2_phase1_fn(void* ctx, uint32_t worker_id,
int64_t start, int64_t end) {
radix_v2_phase1_ctx_t* c = (radix_v2_phase1_ctx_t*)ctx;
if (atomic_load_explicit(&c->oom, memory_order_relaxed)) return;
const ght_layout_t* ly = &c->layout;
uint16_t nk = ly->n_keys;
const uint8_t* const kflags = ly->key_flags;
uint8_t wide_any = ly->any_wide_key;
uint8_t inline_str = ly->any_inline_str;
uint8_t nullable = c->nullable_mask;
const int64_t* match_idx = c->match_idx;
group_ht_t* my_hts = &c->wpart_hts[(size_t)worker_id * c->n_parts];
uint32_t masks[c->n_parts];
for (uint32_t p = 0; p < c->n_parts; p++)
masks[p] = my_hts[p].slots ? my_hts[p].ht_cap - 1 : 0;
char ebuf_stk[8 + 9 * 8 + 8 * 8 + 8];
const void* kpool_stk[8];
char* ebuf = ebuf_stk;
const void** kpool = kpool_stk;
ray_t* v2_stage_hdr = NULL;
if (ly->entry_stride > sizeof(ebuf_stk) || nk > 8) {
char* blk = (char*)scratch_alloc(&v2_stage_hdr,
(size_t)ly->entry_stride + (size_t)nk * sizeof(void*));
if (!blk) { atomic_store_explicit(&c->oom, 1, memory_order_relaxed); return; }
ebuf = blk;
kpool = (const void**)(blk + ly->entry_stride);
}
derive_key_pool(ly, c->key_vecs, kpool);
if (!wide_any && !inline_str && !nullable && ly->null_words == 0 &&
(ly->need_flags & ~(uint32_t)(GHT_NEED_SUM | GHT_NEED_SUM128)) == 0 &&
!(ly->agg_flags_any & (GHT_AF_FIRST | GHT_AF_LAST | GHT_AF_BINARY |
GHT_AF_HOLISTIC)) &&
nk >= 1 && nk <= 2 && ly->entry_stride <= 64) {
enum { V2M = 1024, V2MPF = 8 };
uint8_t v2_packed = ly->packed_key;
uint64_t mh[V2M];
int64_t mk[V2M * 2];
int64_t mrow[V2M];
uint32_t mpart[V2M];
uint16_t morder[V2M];
uint32_t plist[V2M];
ray_t* hist_hdr = NULL;
uint32_t* hist = (uint32_t*)scratch_calloc(&hist_hdr,
(size_t)c->n_parts * sizeof(uint32_t));
if (!hist) {
atomic_store_explicit(&c->oom, 1, memory_order_relaxed);
scratch_free(v2_stage_hdr);
return;
}
int64_t i = start;
while (i < end) {
if (ray_interrupted()) break;
uint32_t mn = 0;
for (; i < end && mn < V2M; i++) {
int64_t row = match_idx ? match_idx[i] : i;
if (!match_idx && c->rowsel &&
!group_rowsel_pass(c->rowsel, row))
continue;
uint64_t h = 0;
for (uint32_t k = 0; k < nk; k++) {
uint64_t kh = 0;
int8_t t = c->key_types[k];
if (t == RAY_F64) {
int64_t kv = group_key_f64_bits(c->key_data[k], row);
double dv; memcpy(&dv, &kv, 8);
mk[(size_t)mn * nk + k] = kv;
if (!v2_packed) kh = ray_hash_f64(dv);
} else {
int64_t kv = read_col_i64(c->key_data[k], row, t,
c->key_attrs[k]);
mk[(size_t)mn * nk + k] = kv;
if (!v2_packed) kh = ray_hash_i64(kv);
}
if (!v2_packed) h = (k == 0) ? kh : ray_hash_combine(h, kh);
}
if (v2_packed) h = ght_hash_lanes(&mk[(size_t)mn * nk], nk);
mh[mn] = h;
mrow[mn] = row;
mpart[mn] = group_radix_part(h, c->n_parts);
mn++;
}
if (mn == 0) continue;
uint32_t np = 0;
for (uint32_t j = 0; j < mn; j++) {
uint32_t p = mpart[j];
if (hist[p] == 0) plist[np++] = p;
hist[p]++;
}
uint32_t run = 0;
for (uint32_t t = 0; t < np; t++) {
uint32_t p = plist[t];
uint32_t cnt = hist[p];
hist[p] = run;
run += cnt;
}
for (uint32_t j = 0; j < mn; j++)
morder[hist[mpart[j]]++] = (uint16_t)j;
for (uint32_t t = 0; t < np; t++) hist[plist[t]] = 0;
for (uint32_t o = 0; o < mn; o++) {
uint32_t j = morder[o];
uint32_t p = mpart[j];
if (!my_hts[p].slots) {
if (!group_ht_init_sized(&my_hts[p], c->ht_init_cap, ly,
c->ht_init_block)) {
atomic_store_explicit(&c->oom, 1,
memory_order_relaxed);
goto v2_morsel_done;
}
my_hts[p].grow_cap = c->ht_grow_cap;
masks[p] = my_hts[p].ht_cap - 1;
}
if (o + V2MPF < mn) {
uint32_t jn = morder[o + V2MPF];
uint32_t pn = mpart[jn];
if (my_hts[pn].slots)
__builtin_prefetch(
&my_hts[pn].slots[(uint32_t)(mh[jn] & masks[pn])],
1, 1);
}
*(uint64_t*)ebuf = mh[j];
int64_t* ek = (int64_t*)(ebuf + 8);
for (uint32_t k = 0; k < nk; k++)
ek[k] = mk[(size_t)j * nk + k];
if (ly->need_flags)
radix_v2_pack_aggs(c, ly,
(int64_t*)(ebuf + 8 + (size_t)ly->key_region), mrow[j]);
memcpy(ebuf + ly->entry_stride - 8, &mrow[j], 8);
masks[p] = group_probe_entry(&my_hts[p], ebuf,
c->key_types, masks[p]);
if (my_hts[p].oom) {
atomic_store_explicit(&c->oom, 1, memory_order_relaxed);
goto v2_morsel_done;
}
}
}
v2_morsel_done:
scratch_free(hist_hdr);
scratch_free(v2_stage_hdr);
return;
}
for (int64_t i = start; i < end; i++) {
if (((i - start) & 65535) == 0 && ray_interrupted()) break;
int64_t row = match_idx ? match_idx[i] : i;
if (!match_idx && c->rowsel && !group_rowsel_pass(c->rowsel, row))
continue;
uint64_t h = 0;
int64_t* ek = (int64_t*)(ebuf + 8);
if (inline_str) {
h = inline_build_keys(ly, c->key_types, c->key_data, c->key_attrs,
kpool, c->key_vecs, nullable, row, ebuf + 8);
} else {
int64_t* nullw = ek + nk;
uint32_t null_words = ly->null_words;
for (uint32_t w = 0; w < null_words; w++) nullw[w] = 0;
if (ly->packed_key) {
for (uint32_t k = 0; k < nk; k++) {
if (__builtin_expect(nullable && ray_key_may_be_null(c->key_vecs[k])
&& ray_vec_is_null(c->key_vecs[k], row), 0)) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
ek[k] = 0;
} else {
ek[k] = read_col_i64(c->key_data[k], row, c->key_types[k], c->key_attrs[k]);
}
}
h = ght_hash_lanes(ek, (uint32_t)nk + null_words);
} else {
for (uint32_t k = 0; k < nk; k++) {
int8_t t = c->key_types[k];
uint64_t kh;
bool is_null = nullable && ray_key_may_be_null(c->key_vecs[k])
&& ray_vec_is_null(c->key_vecs[k], row);
if (is_null) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
ek[k] = 0;
kh = ray_hash_i64(0);
} else if (wide_any && (kflags[k] & GHT_KEYF_WIDE)) {
ek[k] = row;
kh = wide_key_hash_at(ly, k, c->key_data, kpool, row);
} else if (t == RAY_F64) {
int64_t kv = group_key_f64_bits(c->key_data[k], row);
double dv; memcpy(&dv, &kv, 8);
ek[k] = kv;
kh = ray_hash_f64(dv);
} else {
int64_t kv = read_col_i64(c->key_data[k], row, t, c->key_attrs[k]);
ek[k] = kv;
kh = ray_hash_i64(kv);
}
h = (k == 0) ? kh : ray_hash_combine(h, kh);
}
h = ght_hash_null_words(h, nullw, null_words);
}
}
*(uint64_t*)ebuf = h;
if (ly->need_flags)
radix_v2_pack_aggs(c, ly,
(int64_t*)(ebuf + 8 + (size_t)ly->key_region), row);
memcpy(ebuf + ly->entry_stride - 8, &row, 8);
uint32_t p = group_radix_part(h, c->n_parts);
if (!my_hts[p].slots) {
if (!group_ht_init_sized(&my_hts[p], c->ht_init_cap, ly,
c->ht_init_block)) {
atomic_store_explicit(&c->oom, 1, memory_order_relaxed);
break;
}
my_hts[p].grow_cap = c->ht_grow_cap;
if (wide_any && c->key_data) {
group_ht_set_key_data(&my_hts[p], c->key_data);
group_ht_set_key_pool(&my_hts[p], c->key_vecs);
}
masks[p] = my_hts[p].ht_cap - 1;
}
uint32_t new_mask = group_probe_entry(&my_hts[p], ebuf,
c->key_types, masks[p]);
if (my_hts[p].oom) {
atomic_store_explicit(&c->oom, 1, memory_order_relaxed);
break;
}
masks[p] = new_mask;
}
scratch_free(v2_stage_hdr);
}
typedef struct {
group_ht_t* wpart_hts;
group_ht_t* part_hts;
int8_t* key_types;
uint32_t n_workers;
uint32_t n_parts;
ght_layout_t layout;
void** key_data;
const void** key_pool;
topn_fuse_t fuse;
_Atomic(int) oom;
} radix_v2_phase2_ctx_t;
static void radix_v2_phase2_fn(void* ctx, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
radix_v2_phase2_ctx_t* c = (radix_v2_phase2_ctx_t*)ctx;
if (atomic_load_explicit(&c->oom, memory_order_relaxed)) return;
uint16_t row_stride = c->layout.row_stride;
for (int64_t p = start; p < end; p++) {
if (c->fuse.items) c->fuse.item_cnt[p] = 0;
uint32_t total_grps = 0;
for (uint32_t w = 0; w < c->n_workers; w++)
total_grps += c->wpart_hts[(size_t)w * c->n_parts + p].grp_count;
if (total_grps == 0) continue;
uint32_t merge_cap = 2;
while (merge_cap < total_grps && merge_cap < (1u << 30)) merge_cap <<= 1;
uint32_t merge_blk = total_grps < 256 ? total_grps : 256;
if (!group_ht_init_sized(&c->part_hts[p], merge_cap, &c->layout,
merge_blk)) {
atomic_store_explicit(&c->oom, 1, memory_order_relaxed);
return;
}
if (c->layout.any_wide_key && c->key_data) {
group_ht_set_key_data(&c->part_hts[p], c->key_data);
group_ht_copy_key_pool(&c->part_hts[p], c->key_pool);
}
uint32_t mask = c->part_hts[p].ht_cap - 1;
for (uint32_t w = 0; w < c->n_workers; w++) {
group_ht_t* src = &c->wpart_hts[(size_t)w * c->n_parts + p];
if (src->grp_count == 0) continue;
const char* rows = src->rows;
for (uint32_t gi = 0; gi < src->grp_count; gi++) {
mask = group_merge_row(&c->part_hts[p],
rows + (size_t)gi * row_stride,
c->key_types, mask);
if (c->part_hts[p].oom) {
atomic_store_explicit(&c->oom, 1, memory_order_relaxed);
return;
}
}
}
if (c->fuse.items)
topn_fuse_partition(&c->fuse, &c->part_hts[p], (uint32_t)p);
}
}
typedef struct {
const void* key_data;
int8_t key_type;
uint8_t key_attrs;
int64_t* per_worker_min;
int64_t* per_worker_max;
uint32_t n_workers;
const int64_t* match_idx;
ray_t* rowsel;
int64_t span_budget;
_Atomic(int)* abort_flag;
} minmax_ctx_t;
static void minmax_scan_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
minmax_ctx_t* c = (minmax_ctx_t*)ctx;
uint32_t wid = worker_id % c->n_workers;
const int64_t* match_idx = c->match_idx;
int64_t kmin = INT64_MAX, kmax = INT64_MIN;
int8_t t = c->key_type;
const int64_t span_budget = c->span_budget;
#define MINMAX_SEG_LOOP(TYPE, CAST) \
do { \
const TYPE* kd = (const TYPE*)c->key_data; \
for (int64_t i = start; i < end; i++) { \
if (((i - start) & 1023) == 0) { \
if (atomic_load_explicit(c->abort_flag, \
memory_order_relaxed)) \
goto minmax_done; \
if (kmax >= kmin && \
(uint64_t)(kmax - kmin) > (uint64_t)span_budget) { \
atomic_store_explicit(c->abort_flag, 1, \
memory_order_relaxed); \
goto minmax_done; \
} \
} \
int64_t r = match_idx ? match_idx[i] : i; \
if (!match_idx && c->rowsel && !group_rowsel_pass(c->rowsel, r)) continue; \
int64_t v = (int64_t)CAST kd[r]; \
if (t == RAY_SYM && v == 0) continue; \
if (v < kmin) kmin = v; \
if (v > kmax) kmax = v; \
} \
} while (0)
if (t == RAY_I64 || t == RAY_TIMESTAMP)
MINMAX_SEG_LOOP(int64_t, );
else if (RAY_IS_SYM(t)) {
uint8_t w = c->key_attrs & RAY_SYM_W_MASK;
if (w == RAY_SYM_W64) MINMAX_SEG_LOOP(int64_t, );
else if (w == RAY_SYM_W32) MINMAX_SEG_LOOP(uint32_t, );
else if (w == RAY_SYM_W16) MINMAX_SEG_LOOP(uint16_t, );
else MINMAX_SEG_LOOP(uint8_t, );
}
else if (t == RAY_BOOL || t == RAY_U8)
MINMAX_SEG_LOOP(uint8_t, );
else if (t == RAY_I16)
MINMAX_SEG_LOOP(int16_t, );
else
MINMAX_SEG_LOOP(int32_t, );
#undef MINMAX_SEG_LOOP
minmax_done:
if (kmin < c->per_worker_min[wid]) c->per_worker_min[wid] = kmin;
if (kmax > c->per_worker_max[wid]) c->per_worker_max[wid] = kmax;
}
typedef union { double f; int64_t i; } da_val_t;
typedef struct {
da_val_t* sum;
int64_t* sum_hi;
da_val_t* min_val;
da_val_t* max_val;
double* sumsq_f64;
int64_t* count;
int64_t* nn_count;
int64_t* first_row;
int64_t* last_row;
double* sum_y;
double* sumsq_y;
double* sumxy;
ray_t* _h_sum;
ray_t* _h_sum_hi;
ray_t* _h_min;
ray_t* _h_max;
ray_t* _h_sumsq;
ray_t* _h_count;
ray_t* _h_nn_count;
ray_t* _h_first_row;
ray_t* _h_last_row;
ray_t* _h_sum_y;
ray_t* _h_sumsq_y;
ray_t* _h_sumxy;
} da_accum_t;
static inline void da_accum_free(da_accum_t* a) {
scratch_free(a->_h_sum);
scratch_free(a->_h_sum_hi);
scratch_free(a->_h_min);
scratch_free(a->_h_max);
scratch_free(a->_h_sumsq);
scratch_free(a->_h_count);
scratch_free(a->_h_nn_count);
scratch_free(a->_h_first_row);
scratch_free(a->_h_last_row);
scratch_free(a->_h_sum_y);
scratch_free(a->_h_sumsq_y);
scratch_free(a->_h_sumxy);
}
static void emit_agg_columns(ray_t** result, ray_graph_t* g, const ray_op_ext_t* ext,
ray_t* const* agg_vecs, uint32_t grp_count,
uint32_t n_aggs,
const double* sum_f64, const int64_t* sum_i64,
const int64_t* sum_hi64,
const double* min_f64, const double* max_f64,
const int64_t* min_i64, const int64_t* max_i64,
const int64_t* counts,
const agg_affine_t* affine,
const agg_prod_t* prod,
const double* sumsq_f64,
const int64_t* nn_counts,
const double* sum_y_f64,
const double* sumsq_y_f64,
const double* sumxy_f64) {
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t agg_op = ext->agg_ops[a];
ray_t* agg_col = agg_vecs[a];
bool is_f64 = agg_col
? group_fp_type(agg_col->type)
: (prod && prod[a].enabled);
int8_t out_type;
switch (agg_op) {
case OP_AVG:
case OP_STDDEV: case OP_STDDEV_POP:
case OP_VAR: case OP_VAR_POP:
case OP_PEARSON_CORR:
case OP_COV: case OP_SCOV:
case OP_WSUM: case OP_WAVG:
out_type = RAY_F64; break;
case OP_ALL:
case OP_ANY:
out_type = RAY_BOOL; break;
case OP_COUNT: out_type = RAY_I64; break;
case OP_SUM: {
int8_t src_t = agg_col ? agg_col->type
: (op_node(g, ext->agg_ins[a]) ? op_node(g, ext->agg_ins[a])->out_type : 0);
out_type = is_f64 ? RAY_F64 : (src_t == RAY_TIME ? RAY_TIME : RAY_I64);
break;
}
case OP_PROD:
out_type = is_f64 ? RAY_F64 : RAY_I64; break;
default:
out_type = is_f64 ? RAY_F64
: (agg_col ? agg_col->type : RAY_I64); break;
}
ray_t* new_col = out_type == RAY_LIST
? ray_list_new((int64_t)grp_count)
: ray_vec_new(out_type, (int64_t)grp_count);
if (!new_col || RAY_IS_ERR(new_col)) continue;
if (out_type == RAY_SYM)
ray_sym_vec_adopt_domain(new_col, sym_domain_rep(agg_col));
new_col->len = (int64_t)grp_count;
for (uint32_t gi = 0; out_type != RAY_LIST && gi < grp_count; gi++) {
size_t idx = (size_t)gi * n_aggs + a;
int64_t nn = nn_counts ? nn_counts[idx] : counts[gi];
if (out_type == RAY_F64) {
double v;
switch (agg_op) {
case OP_SUM:
v = is_f64 ? sum_f64[idx] : (double)sum_i64[idx];
if (affine && affine[a].enabled)
v += affine[a].bias_f64 * counts[gi];
break;
case OP_PROD:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? sum_f64[idx] : (double)sum_i64[idx];
break;
case OP_AVG:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? sum_f64[idx] / nn
: sum_hi64 ? ray_i128_to_f64(sum_hi64[idx], (uint64_t)sum_i64[idx]) / nn
: (double)sum_i64[idx] / nn;
if (affine && affine[a].enabled)
v += affine[a].bias_f64;
break;
case OP_MIN:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? min_f64[idx] : (double)min_i64[idx]; break;
case OP_MAX:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? max_f64[idx] : (double)max_i64[idx]; break;
case OP_FIRST: case OP_LAST:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? sum_f64[idx] : (double)sum_i64[idx]; break;
case OP_VAR: case OP_VAR_POP:
case OP_STDDEV: case OP_STDDEV_POP: {
int64_t cnt = nn;
bool insuf = (agg_op == OP_VAR || agg_op == OP_STDDEV) ? cnt <= 1 : cnt <= 0;
if (insuf) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
double sum_val = is_f64 ? sum_f64[idx] : (double)sum_i64[idx];
double sq_val = sumsq_f64 ? sumsq_f64[idx] : 0.0;
double mean = sum_val / cnt;
double var_pop = sq_val / cnt - mean * mean;
if (var_pop < 0) var_pop = 0;
if (agg_op == OP_VAR_POP) v = var_pop;
else if (agg_op == OP_VAR) v = var_pop * cnt / (cnt - 1);
else if (agg_op == OP_STDDEV_POP) v = sqrt(var_pop);
else v = sqrt(var_pop * cnt / (cnt - 1));
break;
}
case OP_PEARSON_CORR:
case OP_COV:
case OP_SCOV:
case OP_WSUM:
case OP_WAVG: {
int64_t cnt = nn;
if ((agg_op == OP_PEARSON_CORR || agg_op == OP_SCOV) && cnt < 2) {
v = NULL_F64; ray_vec_set_null(new_col, gi, true); break;
}
if ((agg_op == OP_COV || agg_op == OP_WAVG) && cnt < 1) {
v = NULL_F64; ray_vec_set_null(new_col, gi, true); break;
}
double sx = sum_f64 ? sum_f64[idx] : 0.0;
double sy = sum_y_f64 ? sum_y_f64[idx] : 0.0;
double sxx = sumsq_f64 ? sumsq_f64[idx] : 0.0;
double syy = sumsq_y_f64 ? sumsq_y_f64[idx] : 0.0;
double sxy = sumxy_f64 ? sumxy_f64[idx] : 0.0;
if (agg_op == OP_WSUM) { v = sxy; break; }
if (agg_op == OP_WAVG) {
if (sx == 0.0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = sxy / sx; break;
}
double dn = (double)cnt;
if (agg_op == OP_COV) { v = (sxy - sx * sy / dn) / dn; break; }
if (agg_op == OP_SCOV) { v = (sxy - sx * sy / dn) / (dn - 1.0); break; }
double num = dn * sxy - sx * sy;
double dx = dn * sxx - sx * sx;
double dy = dn * syy - sy * sy;
if (dx <= 0.0 || dy <= 0.0) {
v = NULL_F64; ray_vec_set_null(new_col, gi, true); break;
}
v = num / sqrt(dx * dy);
break;
}
default: v = 0.0; break;
}
((double*)ray_data(new_col))[gi] = ray_f64_fin(v);
} else {
int64_t v;
int64_t int_null = agg_int_null_sentinel_for(out_type);
switch (agg_op) {
case OP_SUM:
v = sum_i64[idx];
if (affine && affine[a].enabled)
v = wrap_add_i64(v, wrap_mul_i64(affine[a].bias_i64, counts[gi]));
break;
case OP_PROD:
if (nn == 0) { v = int_null; ray_vec_set_null(new_col, gi, true); break; }
v = sum_i64[idx]; break;
case OP_COUNT: v = counts[gi]; break;
case OP_ALL: v = (sum_i64[idx] == nn) ? 1 : 0; break;
case OP_ANY: v = (sum_i64[idx] > 0) ? 1 : 0; break;
case OP_MIN:
if (nn == 0) { v = int_null; ray_vec_set_null(new_col, gi, true); break; }
v = min_i64[idx]; break;
case OP_MAX:
if (nn == 0) { v = int_null; ray_vec_set_null(new_col, gi, true); break; }
v = max_i64[idx]; break;
case OP_FIRST: case OP_LAST:
if (nn == 0) { v = int_null; ray_vec_set_null(new_col, gi, true); break; }
v = sum_i64[idx]; break;
default: v = 0; break;
}
topk_write_i64(ray_data(new_col), gi, v,
ray_sym_elem_size(out_type, new_col->attrs));
}
}
ray_op_ext_t* agg_ext = find_ext(g, ext->agg_ins[a]);
int64_t name_id;
if (agg_ext && agg_ext->base.opcode == OP_SCAN) {
name_id = agg_result_col_name(agg_ext->sym, agg_op);
} else {
char nbuf[32];
int np = 0;
nbuf[np++] = '_'; nbuf[np++] = 'e';
{ uint8_t v = a; char dig[3]; int nd = 0;
do { dig[nd++] = (char)('0' + v % 10); v /= 10; } while (v);
while (nd--) nbuf[np++] = dig[nd]; }
const char* nsfx = "";
size_t nslen = 0;
switch (agg_op) {
case OP_SUM: nsfx = "_sum"; nslen = 4; break;
case OP_PROD: nsfx = "_prod"; nslen = 5; break;
case OP_ALL: nsfx = "_all"; nslen = 4; break;
case OP_ANY: nsfx = "_any"; nslen = 4; break;
case OP_COUNT: nsfx = "_count"; nslen = 6; break;
case OP_AVG: nsfx = "_mean"; nslen = 5; break;
case OP_MIN: nsfx = "_min"; nslen = 4; break;
case OP_MAX: nsfx = "_max"; nslen = 4; break;
case OP_FIRST: nsfx = "_first"; nslen = 6; break;
case OP_LAST: nsfx = "_last"; nslen = 5; break;
case OP_STDDEV: nsfx = "_stddev"; nslen = 7; break;
case OP_STDDEV_POP: nsfx = "_stddev_pop"; nslen = 11; break;
case OP_VAR: nsfx = "_var"; nslen = 4; break;
case OP_VAR_POP: nsfx = "_var_pop"; nslen = 8; break;
case OP_MEDIAN: nsfx = "_median"; nslen = 7; break;
case OP_QUANTILE: nsfx = "_quantile"; nslen = 9; break;
case OP_MODE: nsfx = "_mode"; nslen = 5; break;
case OP_COV: nsfx = "_cov"; nslen = 4; break;
case OP_SCOV: nsfx = "_scov"; nslen = 5; break;
case OP_WSUM: nsfx = "_wsum"; nslen = 5; break;
case OP_WAVG: nsfx = "_wavg"; nslen = 5; break;
case OP_TOP_N: nsfx = "_top"; nslen = 4; break;
case OP_BOT_N: nsfx = "_bot"; nslen = 4; break;
}
memcpy(nbuf + np, nsfx, nslen);
name_id = ray_sym_intern(nbuf, (size_t)np + nslen);
}
if (new_col->type == RAY_F64) grp_finalize_nulls(new_col);
*result = ray_table_add_col(*result, name_id, new_col);
ray_release(new_col);
}
}
#define DA_NEED_SUM 0x01
#define DA_NEED_MIN 0x02
#define DA_NEED_MAX 0x04
#define DA_NEED_COUNT 0x08
#define DA_NEED_SUMSQ 0x10
#define DA_NEED_PAIR 0x20
typedef struct {
da_accum_t* accums;
uint32_t n_accums;
void** key_ptrs;
int8_t* key_types;
uint8_t* key_attrs;
uint8_t* key_esz;
int64_t* key_mins;
int64_t* key_strides;
uint8_t n_keys;
void** agg_ptrs;
int8_t* agg_types;
ray_t** agg_cols;
void** agg_ptrs2;
int8_t* agg_types2;
uint8_t* agg_strlen;
agg_prod_t* agg_prod;
uint16_t* agg_ops;
uint8_t n_aggs;
uint8_t need_flags;
uint64_t agg_f64_mask;
uint64_t agg_int_null_mask;
int64_t* agg_int_null_sentinel;
bool all_sum;
uint32_t n_slots;
const int64_t* match_idx;
ray_t* rowsel;
ray_t** sym_strings;
uint32_t sym_count;
int64_t da_n_scan;
ray_sym_domain_raw_t* agg_raw;
uint8_t* agg_raw_ok;
} da_ctx_t;
typedef struct {
uint8_t* used;
int64_t* keys;
int64_t* counts;
da_val_t* sums;
int64_t* sums_hi;
uint32_t cap;
uint32_t size;
ray_t* _h_used;
ray_t* _h_keys;
ray_t* _h_counts;
ray_t* _h_sums;
ray_t* _h_sums_hi;
} sparse_i64_ht_t;
static inline uint64_t sparse_i64_mix(uint64_t x) {
x ^= x >> 30;
x *= UINT64_C(0xbf58476d1ce4e5b9);
x ^= x >> 27;
x *= UINT64_C(0x94d049bb133111eb);
x ^= x >> 31;
return x;
}
static inline uint32_t sparse_i64_pow2(uint32_t x) {
if (x <= 1) return 1;
x--;
x |= x >> 1;
x |= x >> 2;
x |= x >> 4;
x |= x >> 8;
x |= x >> 16;
return x + 1;
}
static void sparse_i64_free(sparse_i64_ht_t* ht) {
if (!ht) return;
scratch_free(ht->_h_used);
scratch_free(ht->_h_keys);
scratch_free(ht->_h_counts);
scratch_free(ht->_h_sums);
scratch_free(ht->_h_sums_hi);
memset(ht, 0, sizeof(*ht));
}
static _Thread_local ray_group_emit_filter_t tl_group_emit_filter;
ray_group_emit_filter_t ray_group_emit_filter_get(void) {
return tl_group_emit_filter;
}
void ray_group_emit_filter_set(ray_group_emit_filter_t filter) {
tl_group_emit_filter = filter;
}
ray_group_emit_filter_t ray_group_emit_filter_active(void) {
ray_group_emit_filter_t f = ray_group_emit_filter_get();
if (!f.enabled) return f;
if (__VM && __VM->eval_depth != f.target_depth) { ray_group_emit_filter_t off = {0}; return off; }
return f;
}
static int64_t da_count_emit_keep_min(const int64_t* counts, uint32_t n_slots,
uint32_t group_count,
ray_group_emit_filter_t filter)
{
int64_t keep_min = filter.min_count_exclusive + 1;
int64_t k_take = filter.top_count_take;
if (k_take <= 0 || k_take >= (int64_t)group_count)
return keep_min;
ray_t* heap_hdr = NULL;
int64_t* heap = (int64_t*)scratch_alloc(&heap_hdr,
(size_t)k_take * sizeof(int64_t));
if (!heap)
return keep_min;
int64_t heap_n = 0;
for (uint32_t s = 0; s < n_slots; s++) {
int64_t cnt = counts[s];
if (cnt <= 0)
continue;
if (heap_n < k_take) {
int64_t j = heap_n++;
heap[j] = cnt;
while (j > 0) {
int64_t p = (j - 1) >> 1;
if (heap[p] <= heap[j]) break;
int64_t tmp = heap[p]; heap[p] = heap[j]; heap[j] = tmp;
j = p;
}
} else if (cnt > heap[0]) {
heap[0] = cnt;
int64_t j = 0;
for (;;) {
int64_t l = j * 2 + 1, r = l + 1, m = j;
if (l < heap_n && heap[l] < heap[m]) m = l;
if (r < heap_n && heap[r] < heap[m]) m = r;
if (m == j) break;
int64_t tmp = heap[m]; heap[m] = heap[j]; heap[j] = tmp;
j = m;
}
}
}
if (heap_n == k_take && heap[0] > keep_min)
keep_min = heap[0];
scratch_free(heap_hdr);
return keep_min;
}
static int64_t da_count_emit_keep_min_u32(const uint32_t* counts,
uint64_t n_slots,
uint32_t group_count,
ray_group_emit_filter_t filter)
{
int64_t keep_min = filter.min_count_exclusive + 1;
int64_t k_take = filter.top_count_take;
if (k_take <= 0 || k_take >= (int64_t)group_count)
return keep_min;
ray_t* heap_hdr = NULL;
int64_t* heap = (int64_t*)scratch_alloc(&heap_hdr,
(size_t)k_take * sizeof(int64_t));
if (!heap)
return keep_min;
int64_t heap_n = 0;
for (uint64_t s = 0; s < n_slots; s++) {
int64_t cnt = (int64_t)counts[s];
if (cnt <= 0)
continue;
if (heap_n < k_take) {
int64_t j = heap_n++;
heap[j] = cnt;
while (j > 0) {
int64_t p = (j - 1) >> 1;
if (heap[p] <= heap[j]) break;
int64_t tmp = heap[p]; heap[p] = heap[j]; heap[j] = tmp;
j = p;
}
} else if (cnt > heap[0]) {
heap[0] = cnt;
int64_t j = 0;
for (;;) {
int64_t l = j * 2 + 1, r = l + 1, m = j;
if (l < heap_n && heap[l] < heap[m]) m = l;
if (r < heap_n && heap[r] < heap[m]) m = r;
if (m == j) break;
int64_t tmp = heap[m]; heap[m] = heap[j]; heap[j] = tmp;
j = m;
}
}
}
if (heap_n == k_take && heap[0] > keep_min)
keep_min = heap[0];
scratch_free(heap_hdr);
return keep_min;
}
static bool sparse_i64_init(sparse_i64_ht_t* ht, uint32_t cap, uint8_t n_aggs,
bool need_sum, bool need_sum128) {
memset(ht, 0, sizeof(*ht));
if (cap < 1024) cap = 1024;
cap = sparse_i64_pow2(cap);
ht->used = (uint8_t*)scratch_calloc(&ht->_h_used, cap);
ht->keys = (int64_t*)scratch_alloc(&ht->_h_keys, (size_t)cap * sizeof(int64_t));
ht->counts = (int64_t*)scratch_calloc(&ht->_h_counts,
(size_t)cap * sizeof(int64_t));
if (need_sum) {
ht->sums = (da_val_t*)scratch_calloc(&ht->_h_sums,
(size_t)cap * n_aggs * sizeof(da_val_t));
}
if (need_sum && need_sum128) {
ht->sums_hi = (int64_t*)scratch_calloc(&ht->_h_sums_hi,
(size_t)cap * n_aggs * sizeof(int64_t));
}
if (!ht->used || !ht->keys || !ht->counts || (need_sum && !ht->sums) ||
(need_sum && need_sum128 && !ht->sums_hi)) {
sparse_i64_free(ht);
return false;
}
ht->cap = cap;
return true;
}
static int32_t sparse_i64_find_slot(const sparse_i64_ht_t* ht, int64_t key) {
uint32_t mask = ht->cap - 1;
uint32_t pos = (uint32_t)sparse_i64_mix((uint64_t)key) & mask;
while (ht->used[pos] && ht->keys[pos] != key)
pos = (pos + 1) & mask;
return (int32_t)pos;
}
static bool sparse_i64_rehash(sparse_i64_ht_t* ht, uint8_t n_aggs,
bool need_sum) {
sparse_i64_ht_t old = *ht;
sparse_i64_ht_t nw;
if (!sparse_i64_init(&nw, old.cap * 2u, n_aggs, need_sum, old.sums_hi != NULL))
return false;
for (uint32_t i = 0; i < old.cap; i++) {
if (!old.used[i]) continue;
int32_t s = sparse_i64_find_slot(&nw, old.keys[i]);
nw.used[s] = 1;
nw.keys[s] = old.keys[i];
nw.counts[s] = old.counts[i];
if (need_sum)
memcpy(&nw.sums[(size_t)s * n_aggs], &old.sums[(size_t)i * n_aggs],
(size_t)n_aggs * sizeof(da_val_t));
if (nw.sums_hi)
memcpy(&nw.sums_hi[(size_t)s * n_aggs], &old.sums_hi[(size_t)i * n_aggs],
(size_t)n_aggs * sizeof(int64_t));
nw.size++;
}
sparse_i64_free(&old);
*ht = nw;
return true;
}
static bool sparse_i64_touch(sparse_i64_ht_t* ht, int64_t key, uint8_t n_aggs,
bool need_sum, int32_t* out_slot) {
if ((uint64_t)(ht->size + 1) * 10u >= (uint64_t)ht->cap * 7u) {
if (!sparse_i64_rehash(ht, n_aggs, need_sum))
return false;
}
int32_t s = sparse_i64_find_slot(ht, key);
if (!ht->used[s]) {
ht->used[s] = 1;
ht->keys[s] = key;
ht->counts[s] = 0;
if (need_sum)
memset(&ht->sums[(size_t)s * n_aggs], 0,
(size_t)n_aggs * sizeof(da_val_t));
if (ht->sums_hi)
memset(&ht->sums_hi[(size_t)s * n_aggs], 0,
(size_t)n_aggs * sizeof(int64_t));
ht->size++;
}
*out_slot = s;
return true;
}
static inline int da_key_is_unsigned(int8_t t) {
return t == RAY_BOOL || t == RAY_U8 || RAY_IS_SYM(t);
}
static inline int32_t da_composite_gid(da_ctx_t* c, int64_t r) {
int32_t gid = 0;
for (uint32_t k = 0; k < c->n_keys; k++) {
int64_t val = read_signed_by_esz(c->key_ptrs[k], r, c->key_esz[k],
da_key_is_unsigned(c->key_types[k]));
if (c->key_types[k] == RAY_SYM && val == 0) val = c->key_mins[k];
gid += (int32_t)((val - c->key_mins[k]) * c->key_strides[k]);
}
return gid;
}
#define DEFINE_DA_COMPOSITE_GID_TYPED(SUFFIX, KTYPE) \
static inline int32_t da_composite_gid_##SUFFIX(da_ctx_t* c, int64_t r) { \
int32_t gid = 0; \
for (uint32_t k = 0; k < c->n_keys; k++) { \
int64_t val = (int64_t)((const KTYPE*)c->key_ptrs[k])[r]; \
if (c->key_types[k] == RAY_SYM && val == 0) val = c->key_mins[k]; \
gid += (int32_t)((val - c->key_mins[k]) * c->key_strides[k]); \
} \
return gid; \
}
DEFINE_DA_COMPOSITE_GID_TYPED(u8, uint8_t)
DEFINE_DA_COMPOSITE_GID_TYPED(u16, uint16_t)
DEFINE_DA_COMPOSITE_GID_TYPED(u32, uint32_t)
DEFINE_DA_COMPOSITE_GID_TYPED(i16, int16_t)
DEFINE_DA_COMPOSITE_GID_TYPED(i32, int32_t)
DEFINE_DA_COMPOSITE_GID_TYPED(i64, int64_t)
#undef DEFINE_DA_COMPOSITE_GID_TYPED
static inline void da_read_val(const void* ptr, int8_t type, uint8_t attrs,
int64_t r, double* out_f64, int64_t* out_i64) {
if (group_fp_type(type)) {
*out_f64 = group_fp_at(ptr, type, r);
*out_i64 = NULL_I64;
} else {
*out_i64 = read_col_i64(ptr, r, type, attrs);
*out_f64 = (double)*out_i64;
}
}
static inline int64_t group_strlen_at(const ray_t* col, int64_t row) {
if (!col || ray_vec_is_null((ray_t*)col, row)) return 0;
if (col->type == RAY_STR) {
const ray_str_t* elems;
const char* pool;
(void)pool;
str_resolve(col, &elems, &pool);
return (int64_t)elems[row].len;
}
{
int64_t sid = ray_read_sym(ray_data((ray_t*)col), row, col->type, col->attrs);
ray_sym_domain_raw_t raw;
if (ray_sym_domain_raw_pin(ray_sym_vec_domain((ray_t*)col), &raw) && sid >= 0 && sid < raw.count) {
size_t sl;
(void)ray_sym_domain_raw_str(&raw, sid, &sl);
return (int64_t)sl;
}
}
ray_t* s = ray_sym_vec_cell((ray_t*)col, row);
return s ? (int64_t)ray_str_len(s) : 0;
}
static inline int64_t group_strlen_at_cached(const ray_t* col, int64_t row,
ray_t** sym_strings,
uint32_t sym_count) {
if (!col || ray_vec_is_null((ray_t*)col, row)) return 0;
if (col->type == RAY_STR) {
const ray_str_t* elems;
const char* pool;
(void)pool;
str_resolve(col, &elems, &pool);
return (int64_t)elems[row].len;
}
if (col->type == RAY_SYM && sym_strings &&
ray_sym_vec_domain((ray_t*)col) == ray_sym_runtime_domain()) {
int64_t sym_id = ray_read_sym(ray_data((ray_t*)col), row,
col->type, col->attrs);
if (sym_id < 0 || (uint64_t)sym_id >= sym_count) return 0;
ray_t* atom = sym_strings[sym_id];
return atom ? (int64_t)ray_str_len(atom) : 0;
}
return group_strlen_at(col, row);
}
static bool try_strlen_sumavg_input(ray_graph_t* g, ray_t* tbl,
ray_op_t* input_op, ray_t** out_vec) {
if (!g || !tbl || !input_op || !out_vec) return false;
if (input_op->opcode != OP_STRLEN || input_op->arity != 1 || !op_child(g, input_op, 0))
return false;
ray_op_t* child = op_child(g, input_op, 0);
ray_op_ext_t* child_ext = find_ext(g, child->id);
if (!child_ext || child_ext->base.opcode != OP_SCAN) return false;
ray_t* col = ray_table_get_col(tbl, child_ext->sym);
if (!col || (col->type != RAY_STR && col->type != RAY_SYM)) return false;
*out_vec = col;
return true;
}
static ray_t* materialize_broadcast_input(ray_t* src, int64_t nrows) {
if (!src || RAY_IS_ERR(src) || nrows < 0) return NULL;
if (src->type == -RAY_STR) return broadcast_scalar(src, nrows);
int8_t out_type = ray_is_atom(src) ? (int8_t)-src->type : src->type;
if (out_type <= 0 || out_type >= RAY_TYPE_COUNT) return NULL;
ray_t* out = ray_vec_new(out_type, nrows);
if (!out || RAY_IS_ERR(out)) return out;
out->len = nrows;
if (nrows == 0) return out;
if (!ray_is_atom(src)) {
uint8_t esz = col_esz(src);
const char* s = (const char*)ray_data(src);
char* d = (char*)ray_data(out);
for (int64_t i = 0; i < nrows; i++)
memcpy(d + (size_t)i * esz, s, esz);
if (out_type == RAY_STR) col_propagate_str_pool(out, src);
return out;
}
switch (src->type) {
case -RAY_F64: {
double v = src->f64;
for (int64_t i = 0; i < nrows; i++) ((double*)ray_data(out))[i] = v;
return out;
}
case -RAY_I64:
case -RAY_SYM:
case -RAY_TIMESTAMP: {
int64_t v = src->i64;
for (int64_t i = 0; i < nrows; i++) ((int64_t*)ray_data(out))[i] = v;
return out;
}
case -RAY_DATE:
case -RAY_TIME: {
int32_t v = (int32_t)src->i64;
for (int64_t i = 0; i < nrows; i++) ((int32_t*)ray_data(out))[i] = v;
return out;
}
case -RAY_I32: {
int32_t v = src->i32;
for (int64_t i = 0; i < nrows; i++) ((int32_t*)ray_data(out))[i] = v;
return out;
}
case -RAY_I16: {
int16_t v = src->i16;
for (int64_t i = 0; i < nrows; i++) ((int16_t*)ray_data(out))[i] = v;
return out;
}
case -RAY_U8:
case -RAY_BOOL: {
uint8_t v = src->u8;
for (int64_t i = 0; i < nrows; i++) ((uint8_t*)ray_data(out))[i] = v;
return out;
}
default:
ray_release(out);
return NULL;
}
}
static bool sum_atom_value(ray_t* x, bool* is_f64, double* out_f64, int64_t* out_i64) {
if (!x || RAY_IS_ERR(x) || !ray_is_atom(x) || RAY_ATOM_IS_NULL(x))
return false;
switch (x->type) {
case -RAY_F64:
*is_f64 = true;
*out_f64 = x->f64;
return true;
case -RAY_I64:
case -RAY_DATE:
case -RAY_TIME:
case -RAY_TIMESTAMP:
case -RAY_SYM:
*is_f64 = false;
*out_i64 = x->i64;
return true;
case -RAY_I32:
*is_f64 = false;
*out_i64 = x->i32;
return true;
case -RAY_I16:
*is_f64 = false;
*out_i64 = x->i16;
return true;
case -RAY_U8:
case -RAY_BOOL:
*is_f64 = false;
*out_i64 = x->u8;
return true;
default:
return false;
}
}
static inline double prod_val_f64(const agg_prod_t* p, int64_t r) {
double x = (p->ta == RAY_F64) ? ((const double*)p->pa)[r]
: (double)read_col_i64(p->pa, r, p->ta, p->aa);
double y = (p->tb == RAY_F64) ? ((const double*)p->pb)[r]
: (double)read_col_i64(p->pb, r, p->tb, p->ab);
return x * y;
}
typedef struct {
const agg_prod_t* prod;
double* dst;
} prod_fill_ctx_t;
static void prod_fill_fn(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid;
prod_fill_ctx_t* c = (prod_fill_ctx_t*)raw;
for (int64_t i = start; i < end; i++)
c->dst[i] = prod_val_f64(c->prod, i);
}
static bool group_materialize_prod_slots(agg_prod_t* prod, ray_t** agg_vecs,
uint8_t* agg_owned, uint8_t n_aggs,
int64_t nrows) {
for (uint32_t a = 0; a < n_aggs; a++) {
if (!prod[a].enabled) continue;
ray_t* v = ray_vec_new(RAY_F64, nrows > 0 ? nrows : 1);
if (!v || RAY_IS_ERR(v)) return false;
v->len = nrows;
prod_fill_ctx_t fc = { .prod = &prod[a], .dst = (double*)ray_data(v) };
ray_pool_t* pool = ray_pool_get();
if (pool && nrows >= RAY_PARALLEL_THRESHOLD)
ray_pool_dispatch(pool, prod_fill_fn, &fc, nrows);
else
prod_fill_fn(&fc, 0, 0, nrows);
agg_vecs[a] = v;
agg_owned[a] = 1;
prod[a].enabled = false;
}
return true;
}
typedef struct {
void** agg_ptrs;
int8_t* agg_types;
ray_t** agg_cols;
uint8_t* agg_strlen;
uint16_t* agg_ops;
agg_linear_t* agg_linear;
agg_prod_t* agg_prod;
uint32_t n_aggs;
uint8_t need_flags;
const int64_t* match_idx;
ray_t* rowsel;
da_accum_t* accums;
uint32_t n_accums;
const bool* agg_int_null_has;
int64_t* agg_int_null_sentinel;
} scalar_ctx_t;
static inline int64_t scalar_i64_at(const void* ptr, int8_t type, int64_t r) {
return read_col_i64(ptr, r, type, 0);
}
static inline void group_i128_sum_i64_range(const int64_t* restrict x, int64_t n,
int64_t* hi, uint64_t* lo) {
const int64_t blk = 1024;
const uint64_t bias = (uint64_t)1 << 53;
int mode = 0;
for (int64_t b = 0; b < n; b += blk) {
int64_t e = (n - b > blk) ? b + blk : n;
if (mode == 0) {
uint64_t s = 0, o = 0;
for (int64_t r = b; r < e; r++) { uint64_t u = (uint64_t)x[r]; s += u; o |= u; }
if ((o >> 53) == 0) { ray_i128_add(hi, lo, (int64_t)s); continue; }
mode = 1;
}
if (mode == 1) {
uint64_t s = 0, o = 0;
for (int64_t r = b; r < e; r++) { uint64_t u = (uint64_t)x[r]; s += u; o |= u + bias; }
if ((o >> 54) == 0) { ray_i128_add(hi, lo, (int64_t)s); continue; }
mode = 2;
}
uint64_t sl = 0; int64_t sh = 0;
for (int64_t r = b; r < e; r++) {
int64_t v = x[r];
sl += (uint64_t)v;
sh += v >> 32;
}
uint64_t low = sl - ((uint64_t)sh << 32);
ray_i128_add128(hi, lo, sh >> 32, (uint64_t)sh << 32);
ray_i128_add128(hi, lo, 0, low);
}
}
static void scalar_sum_i64_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
scalar_ctx_t* c = (scalar_ctx_t*)ctx;
da_accum_t* acc = &c->accums[worker_id];
const int64_t* restrict data = (const int64_t*)c->agg_ptrs[0];
if (acc->sum_hi) {
group_i128_sum_i64_range(data + start, end - start,
&acc->sum_hi[0], (uint64_t*)&acc->sum[0].i);
} else {
int64_t sum = 0;
for (int64_t r = start; r < end; r++)
sum = wrap_add_i64(sum, data[r]);
acc->sum[0].i = wrap_add_i64(acc->sum[0].i, sum);
}
acc->count[0] += end - start;
}
static void scalar_sum_f64_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
scalar_ctx_t* c = (scalar_ctx_t*)ctx;
da_accum_t* acc = &c->accums[worker_id];
const double* restrict data = (const double*)c->agg_ptrs[0];
double sum = 0.0;
for (int64_t r = start; r < end; r++)
sum += data[r];
acc->sum[0].f += sum;
acc->count[0] += end - start;
}
static void scalar_sum_linear_i64_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
scalar_ctx_t* c = (scalar_ctx_t*)ctx;
da_accum_t* acc = &c->accums[worker_id];
const agg_linear_t* lin = &c->agg_linear[0];
int64_t n = end - start;
if (acc->sum_hi) {
bool bare = lin->n_terms == 1 && lin->coeff_i64[0] == 1 &&
lin->bias_i64 == 0 && lin->term_ptrs[0];
int8_t t0 = lin->term_types[0];
if (bare && (t0 == RAY_I64 || t0 == RAY_TIMESTAMP)) {
group_i128_sum_i64_range((const int64_t*)lin->term_ptrs[0] + start, n,
&acc->sum_hi[0], (uint64_t*)&acc->sum[0].i);
} else if (bare && t0 != RAY_SYM) {
const int64_t blk = (int64_t)1 << 30;
for (int64_t b = start; b < end; b += blk) {
int64_t e = (end - b > blk) ? b + blk : end;
int64_t sum = 0;
for (int64_t r = b; r < e; r++)
sum += scalar_i64_at(lin->term_ptrs[0], t0, r);
ray_i128_add(&acc->sum_hi[0], (uint64_t*)&acc->sum[0].i, sum);
}
} else {
for (int64_t r = start; r < end; r++) {
int64_t iv = lin->bias_i64;
for (uint8_t t = 0; t < lin->n_terms; t++)
iv = wrap_add_i64(iv, wrap_mul_i64(lin->coeff_i64[t],
scalar_i64_at(lin->term_ptrs[t], lin->term_types[t], r)));
ray_i128_add(&acc->sum_hi[0], (uint64_t*)&acc->sum[0].i, iv);
}
}
acc->count[0] += n;
return;
}
int64_t sum = wrap_mul_i64(lin->bias_i64, n);
for (uint8_t t = 0; t < lin->n_terms; t++) {
int64_t coeff = lin->coeff_i64[t];
if (coeff == 0) continue;
const void* ptr = lin->term_ptrs[t];
int8_t type = lin->term_types[t];
int64_t term_sum = 0;
for (int64_t r = start; r < end; r++)
term_sum = wrap_add_i64(term_sum, scalar_i64_at(ptr, type, r));
sum = wrap_add_i64(sum, wrap_mul_i64(coeff, term_sum));
}
acc->sum[0].i = wrap_add_i64(acc->sum[0].i, sum);
acc->count[0] += n;
}
static inline void scalar_accum_row(scalar_ctx_t* c, da_accum_t* acc, int64_t r) {
uint32_t n_aggs = c->n_aggs;
acc->count[0]++;
int64_t* nn = acc->nn_count;
for (uint32_t a = 0; a < n_aggs; a++) {
double fv; int64_t iv;
if (c->agg_prod && c->agg_prod[a].enabled) {
fv = prod_val_f64(&c->agg_prod[a], r);
iv = (int64_t)fv;
} else if (c->agg_linear && c->agg_linear[a].enabled) {
const agg_linear_t* lin = &c->agg_linear[a];
iv = lin->bias_i64;
for (uint8_t t = 0; t < lin->n_terms; t++) {
iv = wrap_add_i64(
iv,
wrap_mul_i64(lin->coeff_i64[t],
scalar_i64_at(lin->term_ptrs[t], lin->term_types[t], r)));
}
fv = (double)iv;
} else {
if (!c->agg_ptrs[a]) continue;
if (c->agg_strlen && c->agg_strlen[a]) {
iv = group_strlen_at(c->agg_cols[a], r);
fv = (double)iv;
} else {
uint8_t attrs = c->agg_cols[a] ? c->agg_cols[a]->attrs : 0;
da_read_val(c->agg_ptrs[a], c->agg_types[a], attrs, r, &fv, &iv);
}
}
uint16_t op = c->agg_ops[a];
bool is_f = group_fp_type(c->agg_types[a]);
bool int_null = !is_f && c->agg_int_null_has[a] &&
iv == c->agg_int_null_sentinel[a];
bool is_null = is_f ? !(fv == fv) : int_null;
if (op == OP_ALL || op == OP_ANY) {
if (!is_null) {
acc->sum[a].i += (is_f ? (fv != 0.0) : (iv != 0)) ? 1 : 0;
if (nn) nn[a]++;
}
} else if (op == OP_SUM || op == OP_AVG || op == OP_STDDEV || op == OP_STDDEV_POP || op == OP_VAR || op == OP_VAR_POP) {
if (is_f) {
if (RAY_LIKELY(fv == fv)) {
acc->sum[a].f += fv;
if (acc->sumsq_f64) acc->sumsq_f64[a] += fv * fv;
if (nn) nn[a]++;
}
} else if (RAY_LIKELY(!int_null)) {
if (acc->sum_hi) ray_i128_add(&acc->sum_hi[a], (uint64_t*)&acc->sum[a].i, iv);
else acc->sum[a].i = wrap_add_i64(acc->sum[a].i, iv);
if (acc->sumsq_f64) acc->sumsq_f64[a] += fv * fv;
if (nn) nn[a]++;
}
} else if (op == OP_PROD) {
bool first_seen = nn ? (nn[a] == 0) : (acc->count[0] == 1);
if (is_f) {
if (fv == fv) {
if (first_seen) acc->sum[a].f = fv;
else acc->sum[a].f *= fv;
if (nn) nn[a]++;
}
} else if (RAY_LIKELY(!int_null)) {
if (first_seen) acc->sum[a].i = iv;
else acc->sum[a].i = (int64_t)((uint64_t)acc->sum[a].i * (uint64_t)iv);
if (nn) nn[a]++;
}
} else if (op == OP_FIRST) {
if (!is_null) {
bool first_seen = nn ? (nn[a] == 0) : (acc->count[0] == 1);
if (first_seen) {
if (is_f) acc->sum[a].f = fv;
else acc->sum[a].i = iv;
}
if (nn) nn[a]++;
}
} else if (op == OP_LAST) {
if (!is_null) {
if (is_f) acc->sum[a].f = fv;
else acc->sum[a].i = iv;
if (nn) nn[a]++;
}
} else if (op == OP_MIN) {
if (is_f) { if (fv == fv && fv < acc->min_val[a].f) acc->min_val[a].f = fv; }
else if (c->agg_types[a] == RAY_SYM && !int_null) {
if (acc->min_val[a].i == INT64_MAX ||
sym_lex_lt(ray_sym_vec_domain(c->agg_cols[a]), iv, acc->min_val[a].i))
acc->min_val[a].i = iv;
}
else if (!int_null) { if (iv < acc->min_val[a].i) acc->min_val[a].i = iv; }
if (!is_null && nn) nn[a]++;
} else if (op == OP_MAX) {
if (is_f) { if (fv == fv && fv > acc->max_val[a].f) acc->max_val[a].f = fv; }
else if (c->agg_types[a] == RAY_SYM && !int_null) {
if (acc->max_val[a].i == INT64_MIN ||
sym_lex_gt(ray_sym_vec_domain(c->agg_cols[a]), iv, acc->max_val[a].i))
acc->max_val[a].i = iv;
}
else if (!int_null) { if (iv > acc->max_val[a].i) acc->max_val[a].i = iv; }
if (!is_null && nn) nn[a]++;
}
}
}
static void scalar_accum_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
scalar_ctx_t* c = (scalar_ctx_t*)ctx;
da_accum_t* acc = &c->accums[worker_id];
const int64_t* match_idx = c->match_idx;
for (int64_t i = start; i < end; i++) {
int64_t r = match_idx ? match_idx[i] : i;
if (!match_idx && c->rowsel && !group_rowsel_pass(c->rowsel, r)) continue;
scalar_accum_row(c, acc, r);
}
}
static inline int64_t da_strlen_at(const da_ctx_t* c, uint32_t a, int64_t r) {
const ray_t* col = c->agg_cols[a];
if (col->type == RAY_STR) {
const ray_str_t* elems; const char* pool; (void)pool;
str_resolve(col, &elems, &pool);
return (int64_t)elems[r].len;
}
if (col->type == RAY_SYM) {
int64_t sid = ray_read_sym(ray_data((ray_t*)col), r, RAY_SYM, col->attrs);
if (sid == 0) return 0;
if (c->agg_raw_ok && c->agg_raw_ok[a] && sid > 0 && sid < c->agg_raw[a].count) {
size_t sl;
(void)ray_sym_domain_raw_str(&c->agg_raw[a], sid, &sl);
return (int64_t)sl;
}
}
return group_strlen_at_cached(col, r, c->sym_strings, c->sym_count);
}
static inline bool da_sym_lex_lt(const da_ctx_t* c, uint32_t a_idx, int64_t x, int64_t y) {
if (x == y) return false;
if (c->agg_raw_ok && c->agg_raw_ok[a_idx] && x >= 0 && y >= 0 &&
x < c->agg_raw[a_idx].count && y < c->agg_raw[a_idx].count) {
size_t lx, ly;
const char* px = ray_sym_domain_raw_str(&c->agg_raw[a_idx], x, &lx);
const char* py = ray_sym_domain_raw_str(&c->agg_raw[a_idx], y, &ly);
size_t m = lx < ly ? lx : ly;
int r = m ? memcmp(px, py, m) : 0;
if (r != 0) return r < 0;
return lx < ly;
}
return sym_lex_lt(ray_sym_vec_domain(c->agg_cols[a_idx]), x, y);
}
static inline void da_accum_row(da_ctx_t* c, da_accum_t* acc, int32_t gid, int64_t r) {
uint8_t n_aggs = c->n_aggs;
acc->count[gid]++;
size_t base = (size_t)gid * n_aggs;
if (RAY_LIKELY(c->all_sum)) {
if (!acc->sum) return;
uint64_t f64m = c->agg_f64_mask;
uint64_t inm = c->agg_int_null_mask;
int64_t* nn = acc->nn_count;
for (uint32_t a = 0; a < n_aggs; a++) {
size_t idx = base + a;
if (c->agg_prod && c->agg_prod[a].enabled) {
acc->sum[idx].f += prod_val_f64(&c->agg_prod[a], r);
if (nn) nn[idx]++;
continue;
}
if (!c->agg_ptrs[a]) continue;
if (c->agg_strlen && c->agg_strlen[a]) {
acc->sum[idx].i = wrap_add_i64(acc->sum[idx].i, da_strlen_at(c, a, r));
if (nn) nn[idx]++;
} else if (f64m & ((uint64_t)1 << a)) {
double v = group_fp_at(c->agg_ptrs[a], c->agg_types[a], r);
if (RAY_LIKELY(v == v)) { acc->sum[idx].f += v; if (nn) nn[idx]++; }
} else {
uint8_t v_attrs = c->agg_cols[a] ? c->agg_cols[a]->attrs : 0;
int64_t v = read_col_i64(c->agg_ptrs[a], r, c->agg_types[a], v_attrs);
if (RAY_LIKELY(!((inm >> a) & 1) || v != c->agg_int_null_sentinel[a])) {
if (acc->sum_hi) ray_i128_add(&acc->sum_hi[idx], (uint64_t*)&acc->sum[idx].i, v);
else acc->sum[idx].i = wrap_add_i64(acc->sum[idx].i, v);
if (nn) nn[idx]++;
}
}
}
return;
}
bool fl_take_first = (acc->first_row && r < acc->first_row[gid]);
bool fl_take_last = (acc->last_row && r > acc->last_row[gid]);
bool first_advanced = false, last_advanced = false;
int64_t* nn = acc->nn_count;
for (uint32_t a = 0; a < n_aggs; a++) {
if (!c->agg_ptrs[a] && !(c->agg_prod && c->agg_prod[a].enabled))
continue;
size_t idx = base + a;
double fv; int64_t iv;
if (c->agg_prod && c->agg_prod[a].enabled) {
fv = prod_val_f64(&c->agg_prod[a], r);
iv = (int64_t)fv;
} else if (c->agg_strlen && c->agg_strlen[a]) {
iv = da_strlen_at(c, a, r);
fv = (double)iv;
} else {
uint8_t attrs = c->agg_cols[a] ? c->agg_cols[a]->attrs : 0;
da_read_val(c->agg_ptrs[a], c->agg_types[a], attrs, r, &fv, &iv);
}
uint16_t op = c->agg_ops[a];
bool is_f = group_fp_type(c->agg_types[a]);
bool int_null = (c->agg_int_null_mask & ((uint64_t)1 << a)) &&
iv == c->agg_int_null_sentinel[a];
bool is_null = is_f ? !(fv == fv) : int_null;
if (op == OP_ALL || op == OP_ANY) {
if (!is_null) {
acc->sum[idx].i += (is_f ? (fv != 0.0) : (iv != 0)) ? 1 : 0;
if (nn) nn[idx]++;
}
} else if (op == OP_SUM || op == OP_AVG || op == OP_STDDEV || op == OP_STDDEV_POP || op == OP_VAR || op == OP_VAR_POP) {
if (is_f) {
if (RAY_LIKELY(fv == fv)) {
acc->sum[idx].f += fv;
if (acc->sumsq_f64) acc->sumsq_f64[idx] += fv * fv;
if (nn) nn[idx]++;
}
} else if (RAY_LIKELY(!int_null)) {
if (acc->sum_hi) ray_i128_add(&acc->sum_hi[idx], (uint64_t*)&acc->sum[idx].i, iv);
else acc->sum[idx].i = (int64_t)((uint64_t)acc->sum[idx].i + (uint64_t)iv);
if (acc->sumsq_f64) acc->sumsq_f64[idx] += fv * fv;
if (nn) nn[idx]++;
}
} else if (op == OP_PEARSON_CORR || op == OP_COV || op == OP_SCOV ||
op == OP_WSUM || op == OP_WAVG) {
double xv = is_f ? fv : (double)iv;
double yv; int64_t yi;
da_read_val(c->agg_ptrs2[a], c->agg_types2[a], 0, r, &yv, &yi);
bool y_is_f = group_fp_type(c->agg_types2[a]);
double yval = y_is_f ? yv : (double)yi;
bool y_null = y_is_f ? !(yval == yval) : false;
if (RAY_LIKELY(!is_null && !y_null)) {
acc->sum[idx].f += xv;
if (acc->sumsq_f64) acc->sumsq_f64[idx] += xv * xv;
acc->sum_y[idx] += yval;
acc->sumsq_y[idx] += yval * yval;
acc->sumxy[idx] += xv * yval;
if (nn) nn[idx]++;
}
} else if (op == OP_PROD) {
bool first_seen = nn ? (nn[idx] == 0) : (acc->count[gid] == 1);
if (is_f) {
if (fv == fv) {
if (first_seen) acc->sum[idx].f = fv;
else acc->sum[idx].f *= fv;
if (nn) nn[idx]++;
}
} else if (RAY_LIKELY(!int_null)) {
if (first_seen) acc->sum[idx].i = iv;
else acc->sum[idx].i = (int64_t)((uint64_t)acc->sum[idx].i * (uint64_t)iv);
if (nn) nn[idx]++;
}
} else if (op == OP_FIRST) {
if (fl_take_first && !is_null) {
if (is_f) acc->sum[idx].f = fv;
else acc->sum[idx].i = iv;
first_advanced = true;
if (nn) nn[idx]++;
}
} else if (op == OP_LAST) {
if (fl_take_last && !is_null) {
if (is_f) acc->sum[idx].f = fv;
else acc->sum[idx].i = iv;
last_advanced = true;
if (nn) nn[idx]++;
}
} else if (op == OP_MIN) {
if (is_f) {
if (fv == fv && fv < acc->min_val[idx].f) acc->min_val[idx].f = fv;
} else if (c->agg_types[a] == RAY_SYM && !int_null) {
if (acc->min_val[idx].i == INT64_MAX ||
da_sym_lex_lt(c, a, iv, acc->min_val[idx].i))
acc->min_val[idx].i = iv;
} else if (!int_null) {
if (iv < acc->min_val[idx].i) acc->min_val[idx].i = iv;
}
if (!is_null && nn) nn[idx]++;
} else if (op == OP_MAX) {
if (is_f) {
if (fv == fv && fv > acc->max_val[idx].f) acc->max_val[idx].f = fv;
} else if (c->agg_types[a] == RAY_SYM && !int_null) {
if (acc->max_val[idx].i == INT64_MIN ||
da_sym_lex_lt(c, a, acc->max_val[idx].i, iv))
acc->max_val[idx].i = iv;
} else if (!int_null) {
if (iv > acc->max_val[idx].i) acc->max_val[idx].i = iv;
}
if (!is_null && nn) nn[idx]++;
}
}
if (first_advanced) acc->first_row[gid] = r;
if (last_advanced) acc->last_row[gid] = r;
}
static void da_accum_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
da_ctx_t* c = (da_ctx_t*)ctx;
da_accum_t* acc = &c->accums[worker_id];
uint8_t n_aggs = c->n_aggs;
uint8_t n_keys = c->n_keys;
const int64_t* match_idx = c->match_idx;
#define DA_PF_DIST 8
#define DA_SINGLE_KEY_LOOP(KTYPE, KCAST) \
do { \
const KTYPE* kp = (const KTYPE*)c->key_ptrs[0]; \
int64_t kmin = c->key_mins[0]; \
bool da_pf = c->n_slots >= 4096; \
for (int64_t i = start; i < end; i++) { \
int64_t r = match_idx ? match_idx[i] : i; \
if (!match_idx && c->rowsel && !group_rowsel_pass(c->rowsel, r)) continue; \
if (da_pf && RAY_LIKELY(i + DA_PF_DIST < end)) { \
int64_t pf_r = match_idx ? match_idx[i + DA_PF_DIST] : (i + DA_PF_DIST); \
int64_t pfk = (int64_t)KCAST kp[pf_r]; \
if (c->key_types[0] == RAY_SYM && pfk == 0) pfk = kmin; \
__builtin_prefetch(&acc->count[(int32_t)(pfk - kmin)], 1, 1); \
if (acc->sum) __builtin_prefetch( \
&acc->sum[(size_t)(int32_t)(pfk - kmin) * n_aggs], 1, 1); \
} \
int64_t kv = (int64_t)KCAST kp[r]; \
if (c->key_types[0] == RAY_SYM && kv == 0) kv = kmin; \
da_accum_row(c, acc, (int32_t)(kv - kmin), r); \
} \
} while (0)
if (n_keys == 1) {
if (da_key_is_unsigned(c->key_types[0])) {
switch (c->key_esz[0]) {
case 1: DA_SINGLE_KEY_LOOP(uint8_t, ); break;
case 2: DA_SINGLE_KEY_LOOP(uint16_t, ); break;
case 4: DA_SINGLE_KEY_LOOP(uint32_t, (int64_t)); break;
default: DA_SINGLE_KEY_LOOP(int64_t, ); break;
}
} else {
switch (c->key_esz[0]) {
case 1: DA_SINGLE_KEY_LOOP(int8_t, ); break;
case 2: DA_SINGLE_KEY_LOOP(int16_t, ); break;
case 4: DA_SINGLE_KEY_LOOP(int32_t, (int64_t)); break;
default: DA_SINGLE_KEY_LOOP(int64_t, ); break;
}
}
#undef DA_SINGLE_KEY_LOOP
return;
}
#define DA_MULTI_KEY_LOOP(GID_FN) \
do { \
bool _da_pf = c->n_slots >= 4096; \
for (int64_t i = start; i < end; i++) { \
int64_t r = match_idx ? match_idx[i] : i; \
if (!match_idx && c->rowsel && !group_rowsel_pass(c->rowsel, r)) continue; \
if (_da_pf && RAY_LIKELY(i + DA_PF_DIST < end)) { \
int64_t pf_r = match_idx ? match_idx[i + DA_PF_DIST] : (i + DA_PF_DIST); \
int32_t pf_gid = GID_FN(pf_r); \
__builtin_prefetch(&acc->count[pf_gid], 1, 1); \
if (acc->sum) __builtin_prefetch(&acc->sum[(size_t)pf_gid * n_aggs], 1, 1); \
} \
da_accum_row(c, acc, GID_FN(r), r); \
} \
} while (0)
bool uniform_esz = true;
bool uniform_sign = true;
for (uint32_t k = 1; k < n_keys; k++) {
if (c->key_esz[k] != c->key_esz[0]) { uniform_esz = false; break; }
if (da_key_is_unsigned(c->key_types[k]) !=
da_key_is_unsigned(c->key_types[0])) { uniform_sign = false; break; }
}
if (uniform_esz && uniform_sign) {
bool k_uns = da_key_is_unsigned(c->key_types[0]);
switch (c->key_esz[0]) {
case 1:
if (k_uns) {
#define GID_FN(R) da_composite_gid_u8(c, (R))
DA_MULTI_KEY_LOOP(GID_FN);
#undef GID_FN
} else {
#define GID_FN(R) da_composite_gid(c, (R))
DA_MULTI_KEY_LOOP(GID_FN);
#undef GID_FN
}
break;
case 2:
if (k_uns) {
#define GID_FN(R) da_composite_gid_u16(c, (R))
DA_MULTI_KEY_LOOP(GID_FN);
#undef GID_FN
} else {
#define GID_FN(R) da_composite_gid_i16(c, (R))
DA_MULTI_KEY_LOOP(GID_FN);
#undef GID_FN
}
break;
case 4:
if (k_uns) {
#define GID_FN(R) da_composite_gid_u32(c, (R))
DA_MULTI_KEY_LOOP(GID_FN);
#undef GID_FN
} else {
#define GID_FN(R) da_composite_gid_i32(c, (R))
DA_MULTI_KEY_LOOP(GID_FN);
#undef GID_FN
}
break;
default:
#define GID_FN(R) da_composite_gid_i64(c, (R))
DA_MULTI_KEY_LOOP(GID_FN);
#undef GID_FN
break;
}
} else {
#define GID_FN(R) da_composite_gid(c, (R))
DA_MULTI_KEY_LOOP(GID_FN);
#undef GID_FN
}
#undef DA_MULTI_KEY_LOOP
#undef DA_PF_DIST
}
typedef struct { int32_t* data; int64_t len, cap; ray_t* hdr; } da_bucket_t;
typedef struct {
da_ctx_t* c;
da_bucket_t* buckets;
uint32_t nw, k, span;
int64_t n_scan;
_Atomic(int) oom;
} da_part_ctx_t;
static bool da_bucket_push(da_bucket_t* b, int32_t v) {
if (b->len == b->cap) {
int64_t ncap = b->cap ? b->cap * 2 : 1024;
ray_t* nh = NULL;
int32_t* nd = (int32_t*)scratch_alloc(&nh, (size_t)ncap * sizeof(int32_t));
if (!nd) return false;
if (b->len) memcpy(nd, b->data, (size_t)b->len * sizeof(int32_t));
if (b->hdr) scratch_free(b->hdr);
b->data = nd; b->hdr = nh; b->cap = ncap;
}
b->data[b->len++] = v;
return true;
}
static void da_part_scatter_fn(void* raw, uint32_t wid, int64_t task, int64_t task_end) {
(void)wid; (void)task_end;
da_part_ctx_t* p = (da_part_ctx_t*)raw;
da_ctx_t* c = p->c;
da_bucket_t* mine = &p->buckets[(size_t)task * p->k];
int64_t start = p->n_scan * task / p->nw;
int64_t end = p->n_scan * (task + 1) / p->nw;
const int64_t* match_idx = c->match_idx;
for (int64_t i = start; i < end; i++) {
int64_t r = match_idx ? match_idx[i] : i;
if (!match_idx && c->rowsel && !group_rowsel_pass(c->rowsel, r)) continue;
uint32_t t = (uint32_t)da_composite_gid(c, r) / p->span;
if (!da_bucket_push(&mine[t], (int32_t)r)) {
atomic_store_explicit(&p->oom, 1, memory_order_relaxed);
return;
}
}
}
static void da_part_drain_fn(void* raw, uint32_t wid, int64_t start, int64_t end) {
(void)wid; (void)end;
da_part_ctx_t* p = (da_part_ctx_t*)raw;
da_ctx_t* c = p->c;
da_accum_t* acc = &c->accums[0];
uint32_t t = (uint32_t)start;
for (uint32_t w = 0; w < p->nw; w++) {
const da_bucket_t* b = &p->buckets[(size_t)w * p->k + t];
for (int64_t j = 0; j < b->len; j++) {
int64_t r = b->data[j];
da_accum_row(c, acc, da_composite_gid(c, r), r);
}
}
}
static bool da_accum_partitioned(da_ctx_t* c, ray_pool_t* pool, uint32_t k, int64_t n_scan) {
uint32_t nw = ray_pool_total_workers(pool);
uint32_t span = (c->n_slots + k - 1) / k;
if (span == 0) span = 1;
k = (c->n_slots + span - 1) / span;
ray_t* bh = NULL;
da_bucket_t* buckets = (da_bucket_t*)scratch_calloc(&bh, (size_t)nw * k * sizeof(da_bucket_t));
if (!buckets) return false;
da_part_ctx_t p = { .c = c, .buckets = buckets, .nw = nw, .k = k, .span = span,
.n_scan = n_scan };
atomic_store_explicit(&p.oom, 0, memory_order_relaxed);
ray_pool_dispatch_n(pool, da_part_scatter_fn, &p, nw);
bool ok = atomic_load_explicit(&p.oom, memory_order_relaxed) == 0;
if (ok) ray_pool_dispatch_n(pool, da_part_drain_fn, &p, k);
for (size_t i = 0; i < (size_t)nw * k; i++)
if (buckets[i].hdr) scratch_free(buckets[i].hdr);
scratch_free(bh);
return ok;
}
static void da_accum_task_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
(void)worker_id; (void)end;
da_ctx_t* c = (da_ctx_t*)ctx;
int64_t k = c->n_accums, n = c->da_n_scan;
int64_t base = n / k, rem = n % k;
int64_t lo = start * base + (start < rem ? start : rem);
int64_t hi = lo + base + (start < rem ? 1 : 0);
da_accum_fn(ctx, (uint32_t)start, lo, hi);
}
typedef struct {
da_accum_t* accums;
uint32_t n_src_workers;
uint8_t need_flags;
uint8_t n_aggs;
const int8_t* agg_types;
const uint16_t* agg_ops;
ray_t* const* agg_cols;
} da_merge_ctx_t;
static void da_merge_fn(void* ctx, uint32_t wid, int64_t start, int64_t end) {
(void)wid;
da_merge_ctx_t* c = (da_merge_ctx_t*)ctx;
da_accum_t* merged = &c->accums[0];
uint8_t n_aggs = c->n_aggs;
const int8_t* agg_types = c->agg_types;
for (uint32_t w = 1; w < c->n_src_workers; w++) {
da_accum_t* wa = &c->accums[w];
for (int64_t s = start; s < end; s++) {
size_t base = (size_t)s * n_aggs;
if (c->need_flags & DA_NEED_SUMSQ) {
for (uint32_t a = 0; a < n_aggs; a++)
merged->sumsq_f64[base + a] += wa->sumsq_f64[base + a];
}
if (c->need_flags & DA_NEED_PAIR) {
for (uint32_t a = 0; a < n_aggs; a++) {
merged->sum_y[base + a] += wa->sum_y[base + a];
merged->sumsq_y[base + a] += wa->sumsq_y[base + a];
merged->sumxy[base + a] += wa->sumxy[base + a];
}
}
if (c->need_flags & DA_NEED_SUM) {
for (uint32_t a = 0; a < n_aggs; a++) {
size_t idx = base + a;
uint16_t aop = c->agg_ops ? c->agg_ops[a] : OP_SUM;
int64_t mnn = merged->nn_count ? merged->nn_count[idx] : merged->count[s];
int64_t wnn = wa->nn_count ? wa->nn_count[idx] : wa->count[s];
if (aop == OP_FIRST) {
if (mnn == 0 && wnn > 0)
merged->sum[idx] = wa->sum[idx];
} else if (aop == OP_LAST) {
if (wnn > 0)
merged->sum[idx] = wa->sum[idx];
} else if (aop == OP_PROD) {
if (wnn > 0) {
if (mnn == 0)
merged->sum[idx] = wa->sum[idx];
else if (group_fp_type(agg_types[a]))
merged->sum[idx].f *= wa->sum[idx].f;
else
merged->sum[idx].i = (int64_t)((uint64_t)merged->sum[idx].i * (uint64_t)wa->sum[idx].i);
}
} else if (group_fp_type(agg_types[a]) || agg_is_binary_agg(aop))
merged->sum[idx].f += wa->sum[idx].f;
else if (merged->sum_hi)
ray_i128_add128(&merged->sum_hi[idx], (uint64_t*)&merged->sum[idx].i,
wa->sum_hi[idx], (uint64_t)wa->sum[idx].i);
else
merged->sum[idx].i = wrap_add_i64(merged->sum[idx].i, wa->sum[idx].i);
}
}
if (c->need_flags & DA_NEED_MIN) {
for (uint32_t a = 0; a < n_aggs; a++) {
size_t idx = base + a;
if (group_fp_type(agg_types[a])) {
if (wa->min_val[idx].f < merged->min_val[idx].f)
merged->min_val[idx].f = wa->min_val[idx].f;
} else if (agg_types[a] == RAY_SYM) {
if (wa->min_val[idx].i != INT64_MAX &&
(merged->min_val[idx].i == INT64_MAX ||
sym_lex_lt(ray_sym_vec_domain(c->agg_cols[a]),
wa->min_val[idx].i, merged->min_val[idx].i)))
merged->min_val[idx].i = wa->min_val[idx].i;
} else {
if (wa->min_val[idx].i < merged->min_val[idx].i)
merged->min_val[idx].i = wa->min_val[idx].i;
}
}
}
if (c->need_flags & DA_NEED_MAX) {
for (uint32_t a = 0; a < n_aggs; a++) {
size_t idx = base + a;
if (group_fp_type(agg_types[a])) {
if (wa->max_val[idx].f > merged->max_val[idx].f)
merged->max_val[idx].f = wa->max_val[idx].f;
} else if (agg_types[a] == RAY_SYM) {
if (wa->max_val[idx].i != INT64_MIN &&
(merged->max_val[idx].i == INT64_MIN ||
sym_lex_gt(ray_sym_vec_domain(c->agg_cols[a]),
wa->max_val[idx].i, merged->max_val[idx].i)))
merged->max_val[idx].i = wa->max_val[idx].i;
} else {
if (wa->max_val[idx].i > merged->max_val[idx].i)
merged->max_val[idx].i = wa->max_val[idx].i;
}
}
}
if (merged->nn_count && wa->nn_count) {
for (uint32_t a = 0; a < n_aggs; a++)
merged->nn_count[base + a] += wa->nn_count[base + a];
}
merged->count[s] += wa->count[s];
}
}
}
static inline uint32_t group_ht_lookup_gid(const group_ht_t* ht,
uint64_t hash,
const int64_t* ekeys,
const int8_t* key_types) {
(void)key_types;
const ght_layout_t* ly = &ht->layout;
uint32_t mask = ht->ht_cap - 1;
uint8_t salt = HT_SALT(hash);
uint32_t slot = (uint32_t)(hash & mask);
uint16_t rs = ly->row_stride;
for (;;) {
uint32_t sv = ht->slots[slot];
if (sv == HT_EMPTY) return UINT32_MAX;
if (HT_SALT_V(sv) == salt) {
uint32_t gid = HT_GID(sv);
const char* row = ht->rows + (size_t)gid * rs;
if (group_keys_equal((const int64_t*)(const void*)(row + 8),
ekeys, ly, ht->key_data, ht->key_pool))
return gid;
}
slot = (slot + 1) & mask;
}
}
typedef struct {
void** key_data;
int8_t* key_types;
uint8_t* key_attrs;
ray_t** key_vecs;
uint32_t n_keys;
uint8_t nullable_mask;
const ght_layout_t* layout;
const void** key_pool;
group_ht_t* part_hts;
const uint32_t* part_offsets;
const uint32_t* group_out;
uint32_t n_parts;
int64_t* row_gid;
const int64_t* match_idx;
ray_t* rowsel;
} reprobe_ctx_t;
static inline int64_t reprobe_flat_gid(const reprobe_ctx_t* c, int64_t row,
int64_t* ek_buf, char* keybuf) {
uint32_t nk = c->n_keys;
int8_t* key_types = c->key_types;
void** key_data = c->key_data;
uint8_t* key_attrs = c->key_attrs;
ray_t** key_vecs = c->key_vecs;
uint8_t nullable = c->nullable_mask;
const uint8_t* const kflags = c->layout->key_flags;
uint8_t wide_any = c->layout->any_wide_key;
uint64_t h = 0;
const int64_t* lookup_keys;
if (c->layout->any_inline_str) {
h = inline_build_keys(c->layout, key_types, key_data, key_attrs,
c->key_pool, key_vecs, nullable, row, keybuf);
lookup_keys = (const int64_t*)keybuf;
} else {
int64_t* nullw = ek_buf + nk;
uint32_t null_words = c->layout->null_words;
for (uint32_t w = 0; w < null_words; w++) nullw[w] = 0;
if (c->layout->packed_key) {
for (uint32_t k = 0; k < nk; k++) {
if (__builtin_expect(nullable && ray_key_may_be_null(key_vecs[k])
&& ray_vec_is_null(key_vecs[k], row), 0)) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
ek_buf[k] = 0;
} else {
ek_buf[k] = read_col_i64(key_data[k], row, key_types[k], key_attrs[k]);
}
}
h = ght_hash_lanes(ek_buf, (uint32_t)nk + null_words);
} else {
for (uint32_t k = 0; k < nk; k++) {
int8_t t = key_types[k];
uint64_t kh;
bool is_null = nullable && ray_key_may_be_null(key_vecs[k])
&& ray_vec_is_null(key_vecs[k], row);
if (is_null) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
ek_buf[k] = 0;
kh = ray_hash_i64(0);
} else if (wide_any && (kflags[k] & GHT_KEYF_WIDE)) {
ek_buf[k] = row;
kh = wide_key_hash_at(c->layout, k, key_data, c->key_pool, row);
} else if (t == RAY_F64) {
int64_t kv = group_key_f64_bits(key_data[k], row);
double dv; memcpy(&dv, &kv, 8);
ek_buf[k] = kv;
kh = ray_hash_f64(dv);
} else {
int64_t kv = read_col_i64(key_data[k], row, t, key_attrs[k]);
ek_buf[k] = kv;
kh = ray_hash_i64(kv);
}
h = (k == 0) ? kh : ray_hash_combine(h, kh);
}
h = ght_hash_null_words(h, nullw, null_words);
}
lookup_keys = ek_buf;
}
uint32_t part = group_radix_part(h, c->n_parts);
uint32_t local = group_ht_lookup_gid(&c->part_hts[part], h,
lookup_keys, key_types);
if (local == UINT32_MAX || local >= c->part_hts[part].grp_count)
return -1;
return (int64_t)c->part_offsets[part] + local;
}
static void reprobe_rows_fn(void* vctx, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id;
reprobe_ctx_t* c = (reprobe_ctx_t*)vctx;
int64_t ek_buf_stk[9];
char keybuf_stk[136];
int64_t* ek_buf = ek_buf_stk;
char* keybuf = keybuf_stk;
ray_t* rp_stage_hdr = NULL;
if ((size_t)c->layout->key_region > sizeof(ek_buf_stk)) {
size_t kb = c->layout->key_region;
char* blk = (char*)scratch_alloc(&rp_stage_hdr, kb + kb);
if (!blk) return;
keybuf = blk;
ek_buf = (int64_t*)(blk + kb);
}
const int64_t* match_idx = c->match_idx;
ray_t* rowsel = c->rowsel;
for (int64_t i = start; i < end; i++) {
if (((i - start) & 65535) == 0 && ray_interrupted()) break;
int64_t row = match_idx ? match_idx[i] : i;
if (!match_idx && rowsel && !group_rowsel_pass(rowsel, row)) {
c->row_gid[row] = -1;
continue;
}
int64_t flat = reprobe_flat_gid(c, row, ek_buf, keybuf);
if (flat < 0) {
c->row_gid[row] = -1;
} else {
c->row_gid[row] = c->group_out ? (int64_t)c->group_out[flat]
: flat;
}
}
scratch_free(rp_stage_hdr);
}
typedef struct {
const int64_t* row_gid;
int64_t* hist;
int64_t* cursor;
int64_t* idx_buf;
int64_t n_groups;
int64_t grain;
int64_t nrows;
} med_idx_ctx_t;
static void med_idx_hist_fn(void* vctx, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id; (void)end;
med_idx_ctx_t* c = (med_idx_ctx_t*)vctx;
int64_t task_id = start;
int64_t r_lo = task_id * c->grain;
int64_t r_hi = r_lo + c->grain;
if (r_hi > c->nrows) r_hi = c->nrows;
int64_t* hist = c->hist + task_id * c->n_groups;
const int64_t* row_gid = c->row_gid;
for (int64_t r = r_lo; r < r_hi; r++) {
int64_t gi = row_gid[r];
if (gi >= 0 && gi < c->n_groups) hist[gi]++;
}
}
static void med_idx_scat_fn(void* vctx, uint32_t worker_id,
int64_t start, int64_t end) {
(void)worker_id; (void)end;
med_idx_ctx_t* c = (med_idx_ctx_t*)vctx;
int64_t task_id = start;
int64_t r_lo = task_id * c->grain;
int64_t r_hi = r_lo + c->grain;
if (r_hi > c->nrows) r_hi = c->nrows;
int64_t* cur = c->cursor + task_id * c->n_groups;
const int64_t* row_gid = c->row_gid;
int64_t* idx_buf = c->idx_buf;
for (int64_t r = r_lo; r < r_hi; r++) {
int64_t gi = row_gid[r];
if (gi >= 0 && gi < c->n_groups) idx_buf[cur[gi]++] = r;
}
}
ray_t* exec_group(ray_graph_t* g, ray_op_t* op, ray_t* tbl,
int64_t group_limit);
static ray_t* exec_group_per_partition(ray_graph_t* g, ray_t* parted_tbl,
ray_op_ext_t* ext, int32_t n_parts,
const int64_t* key_syms,
const int64_t* agg_syms,
const uint8_t* key_is_expr,
const uint8_t* agg_is_expr,
const int64_t* borrow_syms, int n_borrow,
int has_avg, int has_stddev,
int64_t group_limit);
static inline bool group_expr_rowlocal_opcode(uint16_t opcode) {
return (opcode >= OP_ROUND && opcode <= OP_IDIV) ||
(opcode >= OP_SIN && opcode <= OP_SIGNUM);
}
static int group_expr_src_syms(ray_graph_t* g, uint32_t root_id,
int64_t* out_syms, int max) {
uint32_t nc = g->node_count;
if (nc > 4096 || root_id >= nc) return -1;
uint32_t stack[64];
int sp = 0, n = 0;
bool visited[4096];
memset(visited, 0, nc * sizeof(bool));
stack[sp++] = root_id;
while (sp > 0) {
uint32_t nid = stack[--sp];
if (nid >= nc || visited[nid]) continue;
visited[nid] = true;
ray_op_t* node = &g->nodes[nid];
if (node->flags & OP_FLAG_DEAD) continue;
if (node->opcode == OP_SCAN) {
ray_op_ext_t* sx = find_ext(g, nid);
if (!sx) return -1;
uint16_t tbl_id = 0;
memcpy(&tbl_id, sx->base.pad, sizeof(uint16_t));
if (tbl_id != 0) return -1;
bool dup = false;
for (int j = 0; j < n; j++)
if (out_syms[j] == sx->sym) { dup = true; break; }
if (!dup) {
if (n >= max) return -1;
out_syms[n++] = sx->sym;
}
continue;
}
if (node->opcode == OP_CONST) continue;
if (!group_expr_rowlocal_opcode(node->opcode)) return -1;
for (int i = 0; i < node->arity && i < 2; i++) {
if (node->in_id[i] == RAY_OP_NONE) continue;
if (sp >= 64) return -1;
stack[sp++] = node->in_id[i];
}
ray_op_ext_t* nx = find_ext(g, nid);
if (nx && nx->third_in != RAY_OP_NONE && nx->third_in != 0) {
if (sp >= 64) return -1;
stack[sp++] = nx->third_in;
}
}
return n;
}
static ray_t* group_eval_part_expr(ray_graph_t* g, ray_t* parted_tbl,
const int64_t* src_syms, int n_src,
uint32_t root, int32_t p) {
ray_t* sub = ray_table_new(n_src > 0 ? (int64_t)n_src : 1);
if (!sub || RAY_IS_ERR(sub)) return NULL;
for (int j = 0; j < n_src; j++) {
ray_t* pcol = ray_table_get_col(parted_tbl, src_syms[j]);
if (!pcol || !RAY_IS_PARTED(pcol->type)) { ray_release(sub); return NULL; }
ray_t* seg = ((ray_t**)ray_data(pcol))[p];
if (!seg) { ray_release(sub); return NULL; }
ray_retain(seg);
sub = ray_table_add_col(sub, src_syms[j], seg);
ray_release(seg);
if (!sub || RAY_IS_ERR(sub)) return NULL;
}
ray_t* saved = g->table;
g->table = sub;
ray_t* vec = exec_node(g, op_node(g, root));
g->table = saved;
ray_release(sub);
if (!vec || RAY_IS_ERR(vec) || !ray_is_vec(vec)) {
if (vec && !RAY_IS_ERR(vec)) ray_release(vec);
return NULL;
}
return vec;
}
static ray_t* exec_group_parted(ray_graph_t* g, ray_op_t* op, ray_t* parted_tbl,
int64_t group_limit) {
int64_t ncols = ray_table_ncols(parted_tbl);
if (ncols <= 0) return ray_error("nyi", NULL);
ray_op_ext_t* ext = find_ext(g, op->id);
if (!ext) return ray_error("nyi", NULL);
uint32_t n_keys = ext->n_keys;
uint32_t n_aggs = ext->n_aggs;
int32_t n_parts = 0;
int64_t total_rows = 0;
for (int64_t c = 0; c < ncols; c++) {
ray_t* col = ray_table_get_col_idx(parted_tbl, c);
if (col && RAY_IS_PARTED(col->type)) {
n_parts = (int32_t)col->len;
total_rows = ray_parted_nrows(col);
break;
}
}
if (n_parts <= 0 || total_rows <= 0) return ray_error("nyi", NULL);
int can_partition = 1;
int has_avg = 0;
int has_stddev = 0;
int64_t key_max = n_keys > 0 ? (int64_t)n_keys : 1;
ray_t* key_syms_hdr = NULL;
int64_t* key_syms = (int64_t*)scratch_alloc(&key_syms_hdr,
(size_t)key_max * sizeof(int64_t));
if (!key_syms) return ray_error("oom", NULL);
int64_t agg_max = n_aggs > 0 ? (int64_t)n_aggs : 1;
ray_t* agg_syms_hdr = NULL;
int64_t* agg_syms = (int64_t*)scratch_alloc(&agg_syms_hdr,
(size_t)agg_max * sizeof(int64_t));
if (!agg_syms) { scratch_free(key_syms_hdr); return ray_error("oom", NULL); }
ray_t *key_is_expr_hdr = NULL, *agg_is_expr_hdr = NULL,
*borrow_hdr = NULL, *esrc_hdr = NULL;
uint8_t* key_is_expr = (uint8_t*)scratch_calloc(&key_is_expr_hdr, (size_t)key_max);
uint8_t* agg_is_expr = (uint8_t*)scratch_calloc(&agg_is_expr_hdr, (size_t)agg_max);
int64_t bsz = ncols > 0 ? ncols : 1;
int64_t* borrow_syms = (int64_t*)scratch_alloc(&borrow_hdr, (size_t)bsz * sizeof(int64_t));
int64_t* esrc = (int64_t*)scratch_alloc(&esrc_hdr, (size_t)bsz * sizeof(int64_t));
if (!key_is_expr || !agg_is_expr || !borrow_syms || !esrc) {
scratch_free(key_syms_hdr); scratch_free(agg_syms_hdr);
scratch_free(key_is_expr_hdr); scratch_free(agg_is_expr_hdr);
scratch_free(borrow_hdr); scratch_free(esrc_hdr);
return ray_error("oom", NULL);
}
int n_borrow = 0;
#define GP_BORROW_ADD(sym) do { \
ray_t* _pc = ray_table_get_col(parted_tbl, (sym)); \
if (!_pc || !RAY_IS_PARTED(_pc->type)) { can_partition = 0; break; } \
int _dup = 0; \
for (int _i = 0; _i < n_borrow; _i++) \
if (borrow_syms[_i] == (sym)) { _dup = 1; break; } \
if (!_dup && n_borrow < bsz) borrow_syms[n_borrow++] = (sym); \
} while (0)
for (uint32_t k = 0; k < n_keys && can_partition; k++) {
ray_op_ext_t* ke = find_ext(g, ext->keys[k]);
if (ke && ke->base.opcode == OP_SCAN) {
key_syms[k] = ke->sym;
key_is_expr[k] = 0;
ray_t* pc = ray_table_get_col(parted_tbl, ke->sym);
if (pc && pc->type == RAY_MAPCOMMON) continue;
GP_BORROW_ADD(ke->sym);
} else {
int ns = group_expr_src_syms(g, ext->keys[k], esrc, (int)bsz);
if (ns <= 0) { can_partition = 0; break; }
char nbuf[24];
int nl = snprintf(nbuf, sizeof(nbuf), "_gpk%u", (unsigned)k);
int64_t synth = ray_sym_intern(nbuf, (size_t)nl);
if (ray_table_get_col(parted_tbl, synth)) { can_partition = 0; break; }
for (int s = 0; s < ns && can_partition; s++) GP_BORROW_ADD(esrc[s]);
key_syms[k] = synth;
key_is_expr[k] = 1;
}
}
for (uint32_t a = 0; a < n_aggs && can_partition; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop == OP_MEDIAN || aop == OP_QUANTILE || aop == OP_MODE ||
aop == OP_TOP_N || aop == OP_BOT_N) {
can_partition = 0; break;
}
if (aop != OP_SUM && aop != OP_COUNT && aop != OP_MIN &&
aop != OP_MAX && aop != OP_AVG && aop != OP_FIRST &&
aop != OP_LAST && aop != OP_STDDEV && aop != OP_STDDEV_POP &&
aop != OP_VAR && aop != OP_VAR_POP) { can_partition = 0; break; }
if (aop == OP_AVG) has_avg = 1;
if (aop == OP_STDDEV || aop == OP_STDDEV_POP ||
aop == OP_VAR || aop == OP_VAR_POP) has_stddev = 1;
ray_op_ext_t* ae = find_ext(g, ext->agg_ins[a]);
if (ae && ae->base.opcode == OP_SCAN) {
agg_syms[a] = ae->sym;
agg_is_expr[a] = 0;
GP_BORROW_ADD(ae->sym);
} else {
int ns = group_expr_src_syms(g, ext->agg_ins[a], esrc, (int)bsz);
if (ns <= 0) { can_partition = 0; break; }
char nbuf[24];
int nl = snprintf(nbuf, sizeof(nbuf), "_gpa%u", (unsigned)a);
int64_t synth = ray_sym_intern(nbuf, (size_t)nl);
if (ray_table_get_col(parted_tbl, synth)) { can_partition = 0; break; }
for (int s = 0; s < ns && can_partition; s++) GP_BORROW_ADD(esrc[s]);
agg_syms[a] = synth;
agg_is_expr[a] = 1;
}
}
#undef GP_BORROW_ADD
if (g->selection && n_borrow == 0) can_partition = 0;
if (can_partition) {
int64_t rows_per_part = total_rows / n_parts;
int64_t est_groups = 1;
for (uint32_t k = 0; k < n_keys; k++) {
if (key_is_expr[k]) {
ray_t* v0 = group_eval_part_expr(g, parted_tbl, borrow_syms,
n_borrow, ext->keys[k], 0);
if (!v0) { est_groups = rows_per_part; break; }
ray_t* cd = ray_count_distinct_approx(v0);
ray_release(v0);
if (!cd || RAY_IS_ERR(cd) || !ray_is_atom(cd)) {
if (cd && !RAY_IS_ERR(cd)) ray_release(cd);
est_groups = rows_per_part; break;
}
int64_t card = cd->i64;
ray_release(cd);
if (card < 1) card = 1;
est_groups *= card;
if (est_groups > rows_per_part) { est_groups = rows_per_part; break; }
continue;
}
ray_t* pcol = ray_table_get_col(parted_tbl, key_syms[k]);
if (!pcol) { est_groups = rows_per_part; break; }
if (pcol->type == RAY_MAPCOMMON) { continue; }
if (!RAY_IS_PARTED(pcol->type)) { est_groups = rows_per_part; break; }
ray_t* seg0 = ((ray_t**)ray_data(pcol))[0];
if (!seg0 || seg0->len <= 0) { est_groups = rows_per_part; break; }
int8_t bt = RAY_PARTED_BASETYPE(pcol->type);
int64_t card;
if (RAY_IS_SYM(bt)) {
int64_t dom_n = ray_sym_domain_count(ray_sym_vec_domain(seg0));
uint32_t sym_n = (dom_n > 0 && dom_n <= UINT32_MAX) ? (uint32_t)dom_n : 0;
if (sym_n == 0 || sym_n > 4194304) { est_groups = rows_per_part; break; }
size_t bwords = ((size_t)sym_n + 63) / 64;
ray_t* bits_hdr = NULL;
uint64_t* bits = (uint64_t*)scratch_calloc(&bits_hdr, bwords * 8);
if (!bits) { est_groups = rows_per_part; break; }
for (int64_t r = 0; r < seg0->len; r++) {
uint32_t id = (uint32_t)ray_read_sym(ray_data(seg0), r, seg0->type, seg0->attrs);
bits[id / 64] |= 1ULL << (id % 64);
}
card = 0;
for (size_t i = 0; i < bwords; i++)
card += __builtin_popcountll(bits[i]);
scratch_free(bits_hdr);
} else if (bt == RAY_I64) {
const int64_t* v = (const int64_t*)ray_data(seg0);
int64_t lo = v[0], hi = v[0];
for (int64_t r = 1; r < seg0->len; r++) {
if (v[r] < lo) lo = v[r];
if (v[r] > hi) hi = v[r];
}
card = hi - lo + 1;
} else if (bt == RAY_I32) {
const int32_t* v = (const int32_t*)ray_data(seg0);
int32_t lo = v[0], hi = v[0];
for (int64_t r = 1; r < seg0->len; r++) {
if (v[r] < lo) lo = v[r];
if (v[r] > hi) hi = v[r];
}
card = (int64_t)(hi - lo + 1);
} else {
card = seg0->len;
}
est_groups *= card;
if (est_groups > rows_per_part) { est_groups = rows_per_part; break; }
}
if (g->selection) {
ray_rowsel_t* sm = ray_rowsel_meta(g->selection);
int64_t pass_per_part = sm->total_pass / (n_parts > 0 ? n_parts : 1);
if (est_groups > pass_per_part) est_groups = pass_per_part;
}
int64_t concat_bytes = total_rows * 8LL * (int64_t)(n_keys + n_aggs);
if (est_groups * 100 > rows_per_part &&
concat_bytes < 4LL * 1024 * 1024 * 1024)
can_partition = 0;
}
if (can_partition) {
ray_t* result = exec_group_per_partition(g, parted_tbl, ext, n_parts,
key_syms, agg_syms,
key_is_expr, agg_is_expr,
borrow_syms, n_borrow,
has_avg, has_stddev, group_limit);
if (result) {
scratch_free(key_syms_hdr);
scratch_free(agg_syms_hdr);
scratch_free(key_is_expr_hdr);
scratch_free(agg_is_expr_hdr);
scratch_free(borrow_hdr);
scratch_free(esrc_hdr);
return result;
}
}
scratch_free(key_syms_hdr);
scratch_free(agg_syms_hdr);
scratch_free(key_is_expr_hdr);
scratch_free(agg_is_expr_hdr);
scratch_free(borrow_hdr);
scratch_free(esrc_hdr);
{
int64_t needed_max = (int64_t)n_keys + (int64_t)n_aggs + ncols;
if (needed_max < 1) needed_max = 1;
ray_t* needed_hdr = NULL;
int64_t* needed = (int64_t*)scratch_alloc(&needed_hdr,
(size_t)needed_max * sizeof(int64_t));
int64_t esrc_max = ncols > 0 ? ncols : 1;
ray_t* c_esrc_hdr = NULL;
int64_t* c_esrc = (int64_t*)scratch_alloc(&c_esrc_hdr,
(size_t)esrc_max * sizeof(int64_t));
if (!needed || !c_esrc) {
scratch_free(needed_hdr); scratch_free(c_esrc_hdr);
return ray_error("oom", NULL);
}
int n_needed = 0;
int copy_all = 0;
for (uint32_t k = 0; k < n_keys && !copy_all; k++) {
ray_op_ext_t* ke = find_ext(g, ext->keys[k]);
if (ke && ke->base.opcode == OP_SCAN) {
int dup = 0;
for (int i = 0; i < n_needed; i++)
if (needed[i] == ke->sym) { dup = 1; break; }
if (!dup) needed[n_needed++] = ke->sym;
} else {
int ns = group_expr_src_syms(g, ext->keys[k], c_esrc, (int)esrc_max);
if (ns < 0) { copy_all = 1; break; }
for (int s = 0; s < ns; s++) {
int dup = 0;
for (int i = 0; i < n_needed; i++)
if (needed[i] == c_esrc[s]) { dup = 1; break; }
if (!dup) needed[n_needed++] = c_esrc[s];
}
}
}
for (uint32_t a = 0; !copy_all && a < n_aggs; a++) {
ray_op_ext_t* ae = find_ext(g, ext->agg_ins[a]);
if (ae && ae->base.opcode == OP_SCAN) {
int dup = 0;
for (int i = 0; i < n_needed; i++)
if (needed[i] == ae->sym) { dup = 1; break; }
if (!dup) needed[n_needed++] = ae->sym;
} else {
int ns = group_expr_src_syms(g, ext->agg_ins[a], c_esrc, (int)esrc_max);
if (ns < 0) { copy_all = 1; break; }
for (int s = 0; s < ns; s++) {
int dup = 0;
for (int i = 0; i < n_needed; i++)
if (needed[i] == c_esrc[s]) { dup = 1; break; }
if (!dup) needed[n_needed++] = c_esrc[s];
}
}
}
if (copy_all) n_needed = 0;
scratch_free(c_esrc_hdr);
ray_t* flat_tbl = ray_table_new(n_needed > 0 ? (int64_t)n_needed : ncols);
if (!flat_tbl || RAY_IS_ERR(flat_tbl)) { scratch_free(needed_hdr); return flat_tbl; }
int64_t cols_to_iter = n_needed > 0 ? (int64_t)n_needed : ncols;
for (int64_t ci = 0; ci < cols_to_iter; ci++) {
ray_t* col;
int64_t name_id;
if (n_needed > 0) {
col = ray_table_get_col(parted_tbl, needed[ci]);
name_id = needed[ci];
} else {
col = ray_table_get_col_idx(parted_tbl, ci);
name_id = ray_table_col_name(parted_tbl, ci);
}
if (!col) continue;
if (col->type == RAY_MAPCOMMON) {
ray_t* mc_flat = materialize_mapcommon(col);
if (mc_flat && !RAY_IS_ERR(mc_flat)) {
flat_tbl = ray_table_add_col(flat_tbl, name_id, mc_flat);
ray_release(mc_flat);
}
continue;
}
if (!RAY_IS_PARTED(col->type)) {
ray_retain(col);
flat_tbl = ray_table_add_col(flat_tbl, name_id, col);
ray_release(col);
continue;
}
int8_t base_type = (int8_t)RAY_PARTED_BASETYPE(col->type);
ray_t** segs = (ray_t**)ray_data(col);
ray_t* flat;
if (base_type == RAY_STR) {
flat = parted_flatten_str(segs, col->len, total_rows);
} else if (base_type == RAY_LIST) {
flat = parted_flatten_list(segs, col->len, total_rows);
} else {
uint8_t base_attrs = (base_type == RAY_SYM)
? parted_sym_max_attrs(segs, col->len) : 0;
flat = typed_vec_new(base_type, base_attrs, total_rows);
if (!flat || RAY_IS_ERR(flat)) {
ray_release(flat_tbl);
scratch_free(needed_hdr);
return flat ? flat : ray_error("oom", NULL);
}
if (base_type == RAY_SYM)
ray_sym_vec_adopt_domain(flat, sym_domain_rep(col));
flat->len = total_rows;
int64_t offset = 0;
uint8_t seg_nulls = 0;
for (int32_t p = 0; p < n_parts; p++) {
ray_t* seg = segs[p];
if (!seg || seg->len <= 0) continue;
parted_copy_cells(ray_data(flat), base_type, base_attrs,
offset, seg, 0, seg->len);
seg_nulls |= seg->attrs & RAY_ATTR_HAS_NULLS;
offset += seg->len;
}
flat->attrs |= seg_nulls;
}
if (!flat || RAY_IS_ERR(flat)) {
ray_release(flat_tbl);
scratch_free(needed_hdr);
return flat ? flat : ray_error("oom", NULL);
}
flat_tbl = ray_table_add_col(flat_tbl, name_id, flat);
ray_release(flat);
}
ray_t* saved = g->table;
g->table = flat_tbl;
ray_t* result = exec_group(g, op, flat_tbl, 0);
g->table = saved;
ray_release(flat_tbl);
scratch_free(needed_hdr);
return result;
}
}
static ray_t* exec_group_run(ray_graph_t* g, ray_op_t* op, ray_t* tbl,
int64_t group_limit);
static ray_t* dict_codes_to_str(const ray_t* codes_col, ray_t* src_col,
const ray_index_t* dix) {
int64_t n = codes_col->len, ndist = dix->u.dict.n_distinct;
const int32_t* cd = (const int32_t*)ray_data((ray_t*)codes_col);
const int32_t* focc = (const int32_t*)ray_data(dix->u.dict.first_occ);
if (n <= 0) return ray_vec_new(RAY_STR, 1);
const char** ptrs = (const char**)ray_alloc_raw((size_t)n * sizeof(const char*));
uint32_t* lens = (uint32_t*)ray_alloc_raw((size_t)n * sizeof(uint32_t));
if (ptrs && lens) {
for (int64_t i = 0; i < n; i++) {
int32_t code = cd[i];
if (code < 0 || code >= ndist) { ptrs[i] = ""; lens[i] = 0; }
else {
size_t vl; const char* vp = ray_str_vec_get(src_col, focc[code], &vl);
ptrs[i] = vp ? vp : "";
lens[i] = vp ? (uint32_t)vl : 0;
}
}
ray_t* out = ray_str_vec_from_parts(ptrs, lens, NULL, n);
ray_free_raw(ptrs);
ray_free_raw(lens);
return out;
}
if (ptrs) ray_free_raw(ptrs);
if (lens) ray_free_raw(lens);
ray_t* out = ray_vec_new(RAY_STR, n);
if (!out || RAY_IS_ERR(out)) return out;
for (int64_t i = 0; i < n; i++) {
int32_t code = cd[i];
if (code < 0 || code >= ndist) { out = ray_str_vec_append(out, "", 0); }
else {
size_t vl; const char* vp = ray_str_vec_get(src_col, focc[code], &vl);
out = ray_str_vec_append(out, vp ? vp : "", vp ? vl : 0);
}
if (!out || RAY_IS_ERR(out)) return out;
}
return out;
}
static _Thread_local ray_dict_cd_t tl_dict_cd;
static void ray_dict_cd_stash(ray_t* codes_col, int64_t n_groups, int64_t key_sym,
const ray_t* tbl, int64_t n_distinct) {
if (tl_dict_cd.result_codes) { ray_free_raw(tl_dict_cd.result_codes); tl_dict_cd.result_codes = NULL; }
tl_dict_cd.valid = 0;
if (!codes_col || n_groups <= 0) return;
int32_t* rc = ray_alloc_raw((size_t)n_groups * sizeof(int32_t));
if (!rc) return;
memcpy(rc, ray_data(codes_col), (size_t)n_groups * sizeof(int32_t));
tl_dict_cd.result_codes = rc;
tl_dict_cd.n_groups = n_groups;
tl_dict_cd.key_sym = key_sym;
tl_dict_cd.tbl = tbl;
tl_dict_cd.n_distinct = n_distinct;
tl_dict_cd.valid = 1;
}
ray_dict_cd_t ray_dict_cd_get(void) { return tl_dict_cd; }
void ray_dict_cd_clear(void) {
if (tl_dict_cd.result_codes) ray_free_raw(tl_dict_cd.result_codes);
memset(&tl_dict_cd, 0, sizeof(tl_dict_cd));
}
static void ray_grp_card_add(int64_t sym, int64_t n_distinct) {
if (__VM && __VM->grp_card_n < 8) {
__VM->grp_card_sym[__VM->grp_card_n] = sym;
__VM->grp_card_val[__VM->grp_card_n] = n_distinct;
__VM->grp_card_n++;
}
}
static int64_t ray_grp_card_lookup(int64_t sym) {
if (!__VM) return 0;
for (uint8_t i = 0; i < __VM->grp_card_n; i++)
if (__VM->grp_card_sym[i] == sym) return __VM->grp_card_val[i];
return 0;
}
ray_t* exec_group(ray_graph_t* g, ray_op_t* op, ray_t* tbl, int64_t group_limit) {
ray_dict_cd_clear();
ray_vm_t* gvm = __VM;
int64_t sv_grp_sym[8], sv_grp_val[8]; uint8_t sv_grp_n = 0;
if (gvm) {
sv_grp_n = gvm->grp_card_n;
memcpy(sv_grp_sym, gvm->grp_card_sym, sizeof sv_grp_sym);
memcpy(sv_grp_val, gvm->grp_card_val, sizeof sv_grp_val);
gvm->grp_card_n = 0;
}
ray_t* result;
ray_t* wrap_hdr = NULL;
ray_t** key_col = NULL;
int64_t* dict_sym = NULL;
if (!tbl || RAY_IS_ERR(tbl) || tbl->type != RAY_TABLE) {
result = exec_group_run(g, op, tbl, group_limit); goto grp_done;
}
ray_op_ext_t* ext = find_ext(g, op->id);
if (!ext || ext->n_keys == 0) {
result = exec_group_run(g, op, tbl, group_limit); goto grp_done;
}
uint32_t nk = ext->n_keys;
key_col = (ray_t**)scratch_alloc(&wrap_hdr,
(size_t)nk * (sizeof(ray_t*) + sizeof(int64_t)));
if (!key_col) { result = exec_group_run(g, op, tbl, group_limit); goto grp_done; }
dict_sym = (int64_t*)(key_col + nk);
for (uint32_t k = 0; k < nk; k++) key_col[k] = NULL;
bool any = false;
for (uint32_t k = 0; k < nk; k++) {
ray_op_t* key_op = op_node(g, ext->keys[k]);
ray_op_ext_t* ke = key_op ? find_ext(g, key_op->id) : NULL;
if (!ke || ke->base.opcode != OP_SCAN) continue;
int64_t sym = ke->sym;
ray_t* col = ray_table_get_col(tbl, sym);
if (!col || col->type != RAY_STR) continue;
if (ray_index_kind(col) != RAY_IDX_DICT) continue;
bool unsafe = false;
for (uint32_t a = 0; a < ext->n_aggs; a++) {
if (ext->agg_ops[a] == OP_COUNT) continue;
ray_op_t* ain = op_node(g, ext->agg_ins[a]);
ray_op_ext_t* ae = ain ? find_ext(g, ain->id) : NULL;
if (ae && ae->base.opcode == OP_SCAN && ae->sym == sym) { unsafe = true; break; }
}
if (unsafe) continue;
key_col[k] = col;
dict_sym[k] = sym;
any = true;
}
if (!any) { result = exec_group_run(g, op, tbl, group_limit); goto grp_done; }
int64_t ncols = ray_table_ncols(tbl);
ray_t* sub = ray_table_new(ncols);
if (!sub || RAY_IS_ERR(sub)) { result = exec_group_run(g, op, tbl, group_limit); goto grp_done; }
for (int64_t c = 0; c < ncols; c++) {
int64_t name = ray_table_col_name(tbl, c);
ray_t* use = ray_table_get_col_idx(tbl, c);
for (uint32_t k = 0; k < nk; k++)
if (key_col[k] && dict_sym[k] == name) {
const ray_index_t* dix = ray_index_payload(key_col[k]->index);
use = dix->u.dict.codes;
ray_grp_card_add(name, dix->u.dict.n_distinct);
break;
}
sub = ray_table_add_col(sub, name, use);
if (!sub || RAY_IS_ERR(sub)) { result = exec_group_run(g, op, tbl, group_limit); goto grp_done; }
}
result = exec_group_run(g, op, sub, group_limit);
ray_release(sub);
if (result && !RAY_IS_ERR(result) && result->type == RAY_TABLE) {
for (uint32_t k = 0; k < nk; k++) {
if (!key_col[k]) continue;
ray_t* codes_col = ray_table_get_col_idx(result, k);
if (codes_col && codes_col->type == RAY_I32) {
if (nk == 1) {
ray_dict_cd_stash(codes_col, ray_table_nrows(result),
dict_sym[k], tbl,
ray_index_payload(key_col[k]->index)->u.dict.n_distinct);
}
ray_t* str_col = dict_codes_to_str(
codes_col, key_col[k], ray_index_payload(key_col[k]->index));
if (str_col && !RAY_IS_ERR(str_col)) {
ray_table_set_col_idx(result, k, str_col);
ray_release(str_col);
} else if (str_col) ray_release(str_col);
}
}
}
grp_done:
scratch_free(wrap_hdr);
if (gvm) {
gvm->grp_card_n = sv_grp_n;
memcpy(gvm->grp_card_sym, sv_grp_sym, sizeof sv_grp_sym);
memcpy(gvm->grp_card_val, sv_grp_val, sizeof sv_grp_val);
}
return result;
}
#define SG_CHUNK_ROWS 32768
typedef struct { int64_t gi, lo, hi; } sg_task_t;
typedef struct {
const ray_idx_slice_t* slices;
const sg_task_t* tasks;
ray_t* const* agg_vecs;
ray_t* const* agg_vecs2;
const agg_prod_t* prod;
const uint16_t* agg_ops;
uint8_t n_aggs;
da_val_t* partials;
int64_t* partials_hi;
double* partial_sumsq;
double* partial_sum_y;
double* partial_sumsq_y;
double* partial_sumxy;
int8_t pair_sum[16];
int8_t fused_by[16];
} sg_ctx_t;
static inline double sg_prod_range(const agg_prod_t* p, int64_t r0, int64_t n,
int64_t* int_side_sum) {
const void* pa = p->pa; const void* pb = p->pb;
int8_t ta = p->ta, tb = p->tb;
if (tb == RAY_F64 && ta != RAY_F64) {
const void* tmp = pa; pa = pb; pb = tmp;
int8_t tt = ta; ta = tb; tb = tt;
}
double a0 = 0.0, a1 = 0.0, a2 = 0.0, a3 = 0.0;
int64_t j = 0;
if (ta == RAY_F64 && tb == RAY_F64) {
const double* restrict x = (const double*)pa + r0;
const double* restrict y = (const double*)pb + r0;
for (; j + 4 <= n; j += 4) {
a0 += x[j] * y[j]; a1 += x[j + 1] * y[j + 1];
a2 += x[j + 2] * y[j + 2]; a3 += x[j + 3] * y[j + 3];
}
for (; j < n; j++) a0 += x[j] * y[j];
} else if (ta == RAY_F64 && (tb == RAY_I64 || tb == RAY_TIMESTAMP)) {
const double* restrict x = (const double*)pa + r0;
const int64_t* restrict y = (const int64_t*)pb + r0;
uint64_t s0 = 0;
if (int_side_sum) {
for (; j + 4 <= n; j += 4) {
a0 += x[j] * (double)y[j]; a1 += x[j + 1] * (double)y[j + 1];
a2 += x[j + 2] * (double)y[j + 2]; a3 += x[j + 3] * (double)y[j + 3];
s0 += (uint64_t)y[j] + (uint64_t)y[j + 1] +
(uint64_t)y[j + 2] + (uint64_t)y[j + 3];
}
for (; j < n; j++) { a0 += x[j] * (double)y[j]; s0 += (uint64_t)y[j]; }
*int_side_sum = (int64_t)s0;
} else {
for (; j + 4 <= n; j += 4) {
a0 += x[j] * (double)y[j]; a1 += x[j + 1] * (double)y[j + 1];
a2 += x[j + 2] * (double)y[j + 2]; a3 += x[j + 3] * (double)y[j + 3];
}
for (; j < n; j++) a0 += x[j] * (double)y[j];
}
} else if (ta == RAY_F64 && tb == RAY_I32) {
const double* restrict x = (const double*)pa + r0;
const int32_t* restrict y = (const int32_t*)pb + r0;
uint64_t s0 = 0;
if (int_side_sum) {
for (; j + 4 <= n; j += 4) {
a0 += x[j] * (double)y[j]; a1 += x[j + 1] * (double)y[j + 1];
a2 += x[j + 2] * (double)y[j + 2]; a3 += x[j + 3] * (double)y[j + 3];
s0 += (uint64_t)(int64_t)y[j] + (uint64_t)(int64_t)y[j + 1] +
(uint64_t)(int64_t)y[j + 2] + (uint64_t)(int64_t)y[j + 3];
}
for (; j < n; j++) { a0 += x[j] * (double)y[j]; s0 += (uint64_t)(int64_t)y[j]; }
*int_side_sum = (int64_t)s0;
} else {
for (; j + 4 <= n; j += 4) {
a0 += x[j] * (double)y[j]; a1 += x[j + 1] * (double)y[j + 1];
a2 += x[j + 2] * (double)y[j + 2]; a3 += x[j + 3] * (double)y[j + 3];
}
for (; j < n; j++) a0 += x[j] * (double)y[j];
}
} else {
if (int_side_sum) *int_side_sum = 0;
for (; j < n; j++) a0 += prod_val_f64(p, r0 + j);
}
return (a0 + a1) + (a2 + a3);
}
static inline double sg_num_at(ray_t* v, int64_t row) {
if (group_fp_type(v->type))
return group_fp_at(ray_data(v), v->type, row);
return (double)read_col_i64(ray_data(v), row, v->type, v->attrs);
}
static inline void sg_pair_accum(ray_t* x, ray_t* y, uint16_t op,
const int64_t* rows,
int64_t n, bool contig, int64_t r0,
double* sx, double* sy, double* sxx,
double* syy, double* sxy) {
(void)op;
double ax = 0.0, ay = 0.0, axx = 0.0, ayy = 0.0, axy = 0.0;
if (contig) {
for (int64_t j = 0; j < n; j++) {
int64_t r = r0 + j;
double xv = sg_num_at(x, r);
double yv = sg_num_at(y, r);
ax += xv; ay += yv;
axx += xv * xv; ayy += yv * yv; axy += xv * yv;
}
} else {
for (int64_t j = 0; j < n; j++) {
int64_t r = rows[j];
double xv = sg_num_at(x, r);
double yv = sg_num_at(y, r);
ax += xv; ay += yv;
axx += xv * xv; ayy += yv * yv; axy += xv * yv;
}
}
*sx = ax; *sy = ay; *sxx = axx; *syy = ayy; *sxy = axy;
}
static void sg_accum_fn(void* raw, uint32_t wid, int64_t tstart, int64_t tend) {
(void)wid;
sg_ctx_t* c = (sg_ctx_t*)raw;
for (int64_t ti = tstart; ti < tend; ti++) {
const sg_task_t* tk = &c->tasks[ti];
const ray_idx_slice_t* sl = &c->slices[tk->gi];
const int64_t* restrict rows = sl->rows ? sl->rows + tk->lo : NULL;
int64_t n = tk->hi - tk->lo;
bool contig = (n > 0 && (!rows || rows[n - 1] - rows[0] + 1 == n));
int64_t r0 = n > 0 ? (rows ? rows[0] : sl->first + tk->lo) : 0;
for (uint32_t a = 0; a < c->n_aggs; a++) {
size_t idx = (size_t)ti * c->n_aggs + a;
uint16_t op = c->agg_ops[a];
if (agg_is_binary_agg(op)) {
double sx = 0.0, sy = 0.0, sxx = 0.0, syy = 0.0, sxy = 0.0;
sg_pair_accum(c->agg_vecs[a], c->agg_vecs2[a], op, rows, n,
contig, r0, &sx, &sy, &sxx, &syy, &sxy);
c->partials[idx].f = sx;
if (c->partial_sumsq) c->partial_sumsq[idx] = sxx;
if (c->partial_sum_y) c->partial_sum_y[idx] = sy;
if (c->partial_sumsq_y) c->partial_sumsq_y[idx] = syy;
if (c->partial_sumxy) c->partial_sumxy[idx] = sxy;
} else if (c->prod[a].enabled) {
double acc = 0.0;
if (contig) {
int8_t pb_slot = c->pair_sum[a];
if (pb_slot >= 0) {
int64_t iss = 0;
acc = sg_prod_range(&c->prod[a], r0, n, &iss);
c->partials[(size_t)ti * c->n_aggs + pb_slot].i = iss;
} else {
acc = sg_prod_range(&c->prod[a], r0, n, NULL);
}
} else {
for (int64_t j = 0; j < n; j++)
acc += prod_val_f64(&c->prod[a], rows[j]);
}
c->partials[idx].f = acc;
} else if (contig && c->fused_by[a] >= 0) {
} else if (op == OP_COUNT) {
} else if (group_fp_type(c->agg_vecs[a]->type)) {
double acc = 0.0;
double ssq = 0.0;
bool need_sq = c->partial_sumsq &&
(op == OP_STDDEV || op == OP_STDDEV_POP ||
op == OP_VAR || op == OP_VAR_POP);
if (contig) {
for (int64_t j = 0; j < n; j++) {
double v = group_fp_at(ray_data(c->agg_vecs[a]),
c->agg_vecs[a]->type, r0 + j);
if (RAY_LIKELY(v == v)) {
acc += v;
if (need_sq) ssq += v * v;
}
}
} else {
for (int64_t j = 0; j < n; j++) {
double v = group_fp_at(ray_data(c->agg_vecs[a]),
c->agg_vecs[a]->type, rows[j]);
if (RAY_LIKELY(v == v)) {
acc += v;
if (need_sq) ssq += v * v;
}
}
}
c->partials[idx].f = acc;
if (need_sq) c->partial_sumsq[idx] = ssq;
} else {
ray_t* av = c->agg_vecs[a];
const void* p = ray_data(av);
int8_t t = av->type;
uint8_t at = av->attrs;
uint64_t acc = 0;
int64_t hi = 0;
double ssq = 0.0;
bool need_sq = c->partial_sumsq &&
(op == OP_STDDEV || op == OP_STDDEV_POP ||
op == OP_VAR || op == OP_VAR_POP);
if (contig && (t == RAY_I64 || t == RAY_TIMESTAMP)) {
const int64_t* restrict x = (const int64_t*)p + r0;
if (c->partials_hi) {
group_i128_sum_i64_range(x, n, &hi, &acc);
if (need_sq)
for (int64_t j = 0; j < n; j++) { double d = (double)x[j]; ssq += d * d; }
} else {
for (int64_t j = 0; j < n; j++) {
int64_t v = x[j];
acc += (uint64_t)v;
if (need_sq) { double d = (double)v; ssq += d * d; }
}
}
} else if (contig && (t == RAY_I32 || t == RAY_DATE || t == RAY_TIME)) {
const int32_t* restrict x = (const int32_t*)p + r0;
for (int64_t j = 0; j < n; j++) {
int64_t v = (int64_t)x[j];
acc += (uint64_t)v;
if (need_sq) { double d = (double)v; ssq += d * d; }
}
hi = (int64_t)acc < 0 ? -1 : 0;
} else {
for (int64_t j = 0; j < n; j++) {
int64_t v = read_col_i64(p, rows[j], t, at);
ray_i128_add(&hi, &acc, v);
if (need_sq) { double d = (double)v; ssq += d * d; }
}
}
c->partials[idx].i = (int64_t)acc;
if (c->partials_hi) c->partials_hi[idx] = hi;
if (need_sq) c->partial_sumsq[idx] = ssq;
}
}
}
}
static int sg_cmp_i64(const void* a, const void* b) {
int64_t x = *(const int64_t*)a, y = *(const int64_t*)b;
return (x > y) - (x < y);
}
static inline void sg_hint_release(ray_graph_t* g) {
if (g->sg_col) { ray_release(g->sg_col); g->sg_col = NULL; }
if (g->sg_slices_hdr) { ray_free(g->sg_slices_hdr); g->sg_slices_hdr = NULL; }
g->sg_nslices = 0;
}
static ray_t* sg_hint_to_selection(ray_graph_t* g, ray_t* tbl) {
int64_t nrows = ray_table_nrows(tbl);
const ray_idx_slice_t* sl = g->sg_slices_hdr
? (const ray_idx_slice_t*)ray_data(g->sg_slices_hdr) : NULL;
int64_t K = g->sg_nslices;
int64_t total = 0;
for (int64_t i = 0; i < K; i++) total += sl[i].n;
ray_t* sel = NULL;
if (total == 0) {
sel = ray_index_rowsel_from_ids(nrows, NULL, 0);
} else {
ray_t* hdr = ray_alloc((size_t)total * (int64_t)sizeof(int64_t));
if (!hdr) { sg_hint_release(g); return ray_error("oom", NULL); }
int64_t* ids = (int64_t*)ray_data(hdr);
int64_t w = 0;
for (int64_t i = 0; i < K; i++) {
if (sl[i].rows) {
memcpy(ids + w, sl[i].rows,
(size_t)sl[i].n * sizeof(int64_t));
} else {
for (int64_t j = 0; j < sl[i].n; j++)
ids[w + j] = sl[i].first + j;
}
w += sl[i].n;
}
qsort(ids, (size_t)total, sizeof(int64_t), sg_cmp_i64);
sel = ray_index_rowsel_from_ids(nrows, ids, total);
ray_free(hdr);
}
sg_hint_release(g);
if (!sel) return ray_error("oom", NULL);
g->selection = sel;
return NULL;
}
static bool sg_shape_eligible(ray_graph_t* g, ray_op_t* op, ray_t* tbl,
int64_t group_limit,
ray_t** agg_vecs, ray_t** agg_vecs2,
agg_prod_t* prod) {
if (group_limit < 0) return false;
if (ray_group_emit_filter_active().enabled) return false;
ray_op_ext_t* ext = find_ext(g, op->id);
if (!ext || ext->n_keys != 1 || ext->n_aggs < 1 || ext->n_aggs > 16)
return false;
if (ext->agg_k) return false;
ray_op_ext_t* ke = find_ext(g, ext->keys[0]);
if (!ke || ke->base.opcode != OP_SCAN) return false;
ray_t* key_col = ray_table_get_col(tbl, ke->sym);
if (!key_col || key_col != g->sg_col) return false;
for (uint32_t a = 0; a < ext->n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop == OP_COUNT) continue;
bool pair = agg_is_binary_agg(aop);
if (!pair &&
aop != OP_SUM && aop != OP_AVG &&
aop != OP_STDDEV && aop != OP_STDDEV_POP &&
aop != OP_VAR && aop != OP_VAR_POP)
return false;
ray_op_t* in = op_node(g, ext->agg_ins[a]);
if (!in) return false;
if (!pair && (aop == OP_SUM || aop == OP_AVG) &&
try_prod_sumavg_input_f64(g, tbl, in, &prod[a]))
continue;
ray_op_ext_t* ae = find_ext(g, in->id);
if (!ae || ae->base.opcode != OP_SCAN) return false;
ray_t* col = ray_table_get_col(tbl, ae->sym);
if (!col || ray_is_atom(col)) return false;
if (ray_vec_has_nulls(col)) return false;
if (!agg_type_admitted(aop, col->type)) return false;
switch (col->type) {
case RAY_U8: case RAY_I16: case RAY_I32: case RAY_I64:
case RAY_TIME: case RAY_F32: case RAY_F64: break;
default: return false;
}
agg_vecs[a] = col;
if (pair) {
if (!ext->agg_ins2 || ext->agg_ins2[a] == RAY_OP_NONE) return false;
ray_op_t* in2 = op_node(g, ext->agg_ins2[a]);
if (!in2) return false;
ray_op_ext_t* ae2 = find_ext(g, in2->id);
if (!ae2 || ae2->base.opcode != OP_SCAN) return false;
ray_t* col2 = ray_table_get_col(tbl, ae2->sym);
if (!col2 || ray_is_atom(col2)) return false;
if (ray_vec_has_nulls(col2)) return false;
if (!agg_type_admitted(aop, col2->type)) return false;
switch (col2->type) {
case RAY_U8: case RAY_I16: case RAY_I32: case RAY_I64:
case RAY_F32: case RAY_F64: break;
default: return false;
}
agg_vecs2[a] = col2;
} else if (ext->agg_ins2 && ext->agg_ins2[a] != RAY_OP_NONE) {
return false;
}
}
return true;
}
ray_t* ray_group_slice_hint_settle(ray_graph_t* g, ray_op_t* op, ray_t* tbl,
int64_t group_limit) {
if (!g->sg_col || !tbl || tbl->type != RAY_TABLE) return NULL;
ray_t* agg_vecs[16] = {0};
ray_t* agg_vecs2[16] = {0};
agg_prod_t prod[16];
memset(prod, 0, sizeof(prod));
if (sg_shape_eligible(g, op, tbl, group_limit, agg_vecs, agg_vecs2, prod))
return NULL;
return sg_hint_to_selection(g, tbl);
}
static ray_t* exec_group_slices(ray_graph_t* g, ray_op_t* op, ray_t* tbl,
int64_t group_limit) {
ray_t* agg_vecs[16] = {0};
ray_t* agg_vecs2[16] = {0};
agg_prod_t prod[16];
memset(prod, 0, sizeof(prod));
if (!sg_shape_eligible(g, op, tbl, group_limit, agg_vecs, agg_vecs2, prod))
return NULL;
ray_op_ext_t* ext = find_ext(g, op->id);
ray_op_ext_t* ke = find_ext(g, ext->keys[0]);
ray_t* key_col = ray_table_get_col(tbl, ke->sym);
uint32_t n_aggs = ext->n_aggs;
int64_t K = g->sg_nslices;
const ray_idx_slice_t* slices = (K > 0)
? (const ray_idx_slice_t*)ray_data(g->sg_slices_hdr) : NULL;
bool need_sumsq = false, need_pair = false, need_sum128 = false;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop == OP_STDDEV || aop == OP_STDDEV_POP ||
aop == OP_VAR || aop == OP_VAR_POP ||
agg_is_binary_agg(aop))
need_sumsq = true;
if (agg_is_binary_agg(aop))
need_pair = true;
if (aop == OP_AVG && !prod[a].enabled && agg_vecs[a] &&
!group_fp_type(agg_vecs[a]->type) &&
group_avg_needs_hi(agg_vecs[a]->type, ray_table_nrows(tbl)))
need_sum128 = true;
}
ray_t *sum_hdr = NULL, *cnt_hdr = NULL, *task_hdr = NULL, *part_hdr = NULL;
ray_t *sumsq_hdr = NULL, *sum_y_hdr = NULL, *sumsq_y_hdr = NULL, *sumxy_hdr = NULL;
ray_t *part_sumsq_hdr = NULL, *part_sum_y_hdr = NULL, *part_sumsq_y_hdr = NULL, *part_sumxy_hdr = NULL;
ray_t *sum_hi_hdr = NULL, *part_hi_hdr = NULL;
da_val_t* sums = NULL;
int64_t* sums_hi = NULL;
double *sumsq = NULL, *sum_y = NULL, *sumsq_y = NULL, *sumxy = NULL;
int64_t* counts = NULL;
if (K > 0) {
sums = (da_val_t*)scratch_calloc(&sum_hdr,
(size_t)K * n_aggs * sizeof(da_val_t));
counts = (int64_t*)scratch_calloc(&cnt_hdr,
(size_t)K * sizeof(int64_t));
if (need_sum128)
sums_hi = (int64_t*)scratch_calloc(&sum_hi_hdr,
(size_t)K * n_aggs * sizeof(int64_t));
if (need_sumsq)
sumsq = (double*)scratch_calloc(&sumsq_hdr,
(size_t)K * n_aggs * sizeof(double));
if (need_pair) {
sum_y = (double*)scratch_calloc(&sum_y_hdr,
(size_t)K * n_aggs * sizeof(double));
sumsq_y = (double*)scratch_calloc(&sumsq_y_hdr,
(size_t)K * n_aggs * sizeof(double));
sumxy = (double*)scratch_calloc(&sumxy_hdr,
(size_t)K * n_aggs * sizeof(double));
}
int64_t n_tasks = 0;
for (int64_t i = 0; i < K; i++)
n_tasks += (slices[i].n + SG_CHUNK_ROWS - 1) / SG_CHUNK_ROWS;
sg_task_t* tasks = NULL;
da_val_t* partials = NULL;
int64_t* part_hi = NULL;
double *part_sumsq = NULL, *part_sum_y = NULL, *part_sumsq_y = NULL, *part_sumxy = NULL;
if (sums && counts) {
task_hdr = ray_alloc((size_t)n_tasks * (int64_t)sizeof(sg_task_t));
tasks = task_hdr ? (sg_task_t*)ray_data(task_hdr) : NULL;
partials = (da_val_t*)scratch_calloc(&part_hdr,
(size_t)n_tasks * n_aggs * sizeof(da_val_t));
if (need_sum128)
part_hi = (int64_t*)scratch_calloc(&part_hi_hdr,
(size_t)n_tasks * n_aggs * sizeof(int64_t));
if (need_sumsq)
part_sumsq = (double*)scratch_calloc(&part_sumsq_hdr,
(size_t)n_tasks * n_aggs * sizeof(double));
if (need_pair) {
part_sum_y = (double*)scratch_calloc(&part_sum_y_hdr,
(size_t)n_tasks * n_aggs * sizeof(double));
part_sumsq_y = (double*)scratch_calloc(&part_sumsq_y_hdr,
(size_t)n_tasks * n_aggs * sizeof(double));
part_sumxy = (double*)scratch_calloc(&part_sumxy_hdr,
(size_t)n_tasks * n_aggs * sizeof(double));
}
}
if (!sums || !counts || !tasks || !partials ||
(need_sum128 && (!sums_hi || !part_hi)) ||
(need_sumsq && (!sumsq || !part_sumsq)) ||
(need_pair && (!sum_y || !sumsq_y || !sumxy ||
!part_sum_y || !part_sumsq_y || !part_sumxy))) {
scratch_free(sum_hdr); scratch_free(cnt_hdr);
scratch_free(sumsq_hdr); scratch_free(sum_y_hdr);
scratch_free(sumsq_y_hdr); scratch_free(sumxy_hdr);
scratch_free(sum_hi_hdr); scratch_free(part_hi_hdr);
if (task_hdr) ray_free(task_hdr);
scratch_free(part_hdr);
scratch_free(part_sumsq_hdr); scratch_free(part_sum_y_hdr);
scratch_free(part_sumsq_y_hdr); scratch_free(part_sumxy_hdr);
return NULL;
}
int64_t total_rows = 0, tw = 0;
for (int64_t i = 0; i < K; i++) {
total_rows += slices[i].n;
for (int64_t lo = 0; lo < slices[i].n; lo += SG_CHUNK_ROWS) {
int64_t hi = lo + SG_CHUNK_ROWS;
if (hi > slices[i].n) hi = slices[i].n;
tasks[tw].gi = i; tasks[tw].lo = lo; tasks[tw].hi = hi; tw++;
}
}
sg_ctx_t ctx = { slices, tasks, agg_vecs, agg_vecs2, prod, ext->agg_ops,
n_aggs, partials, part_hi, part_sumsq, part_sum_y,
part_sumsq_y, part_sumxy, {0}, {0} };
for (uint8_t a = 0; a < 16; a++) { ctx.pair_sum[a] = -1; ctx.fused_by[a] = -1; }
for (uint32_t a = 0; a < n_aggs; a++) {
if (!prod[a].enabled) continue;
const void* ip; int8_t it;
if (prod[a].ta != RAY_F64) { ip = prod[a].pa; it = prod[a].ta; }
else if (prod[a].tb != RAY_F64) { ip = prod[a].pb; it = prod[a].tb; }
else continue;
if (it != RAY_I64 && it != RAY_I32) continue;
for (uint32_t b = 0; b < n_aggs; b++) {
if (b == a || !agg_vecs[b] || ctx.fused_by[b] >= 0) continue;
if (ext->agg_ops[b] != OP_SUM && ext->agg_ops[b] != OP_AVG) continue;
if (ext->agg_ops[b] == OP_AVG && part_hi) continue;
if (ray_data(agg_vecs[b]) != ip || agg_vecs[b]->type != it) continue;
ctx.pair_sum[a] = (int8_t)b;
ctx.fused_by[b] = (int8_t)a;
break;
}
}
ray_pool_t* pool = ray_pool_get();
if (pool && n_tasks > 1 && total_rows >= RAY_PARALLEL_THRESHOLD)
ray_pool_dispatch_n(pool, sg_accum_fn, &ctx, (uint32_t)n_tasks);
else
sg_accum_fn(&ctx, 0, 0, n_tasks);
for (int64_t ti = 0; ti < n_tasks; ti++) {
int64_t gi = tasks[ti].gi;
counts[gi] += tasks[ti].hi - tasks[ti].lo;
for (uint32_t a = 0; a < n_aggs; a++) {
size_t di = (size_t)gi * n_aggs + a;
size_t si = (size_t)ti * n_aggs + a;
bool pair = agg_is_binary_agg(ext->agg_ops[a]);
if (pair || prod[a].enabled || ext->agg_ops[a] == OP_COUNT ||
(agg_vecs[a] && group_fp_type(agg_vecs[a]->type)))
sums[di].f += partials[si].f;
else if (sums_hi)
ray_i128_add128(&sums_hi[di], (uint64_t*)&sums[di].i,
part_hi[si], (uint64_t)partials[si].i);
else
sums[di].i = (int64_t)((uint64_t)sums[di].i +
(uint64_t)partials[si].i);
if (sumsq) sumsq[di] += part_sumsq[si];
if (pair) {
sum_y[di] += part_sum_y[si];
sumsq_y[di] += part_sumsq_y[si];
sumxy[di] += part_sumxy[si];
}
}
}
ray_free(task_hdr);
scratch_free(part_hdr); scratch_free(part_hi_hdr);
scratch_free(part_sumsq_hdr); scratch_free(part_sum_y_hdr);
scratch_free(part_sumsq_y_hdr); scratch_free(part_sumxy_hdr);
}
ray_t* result = ray_table_new(1 + n_aggs);
if (!result || RAY_IS_ERR(result)) {
scratch_free(sum_hdr); scratch_free(cnt_hdr);
scratch_free(sumsq_hdr); scratch_free(sum_y_hdr);
scratch_free(sumsq_y_hdr); scratch_free(sumxy_hdr);
scratch_free(sum_hi_hdr);
return NULL;
}
ray_t* kc = col_vec_new(key_col, K > 0 ? K : 1);
if (!kc || RAY_IS_ERR(kc)) {
if (kc) ray_release(kc);
ray_release(result);
scratch_free(sum_hdr); scratch_free(cnt_hdr);
scratch_free(sumsq_hdr); scratch_free(sum_y_hdr);
scratch_free(sumsq_y_hdr); scratch_free(sumxy_hdr);
scratch_free(sum_hi_hdr);
return NULL;
}
if (kc->type == RAY_SYM)
ray_sym_vec_adopt_domain(kc, sym_domain_rep(key_col));
kc->len = K;
for (int64_t gi = 0; gi < K; gi++)
write_col_i64(ray_data(kc), gi, slices[gi].dom,
key_col->type, kc->attrs);
result = ray_table_add_col(result, ke->sym, kc);
ray_release(kc);
if (!result || RAY_IS_ERR(result)) {
scratch_free(sum_hdr); scratch_free(cnt_hdr);
scratch_free(sumsq_hdr); scratch_free(sum_y_hdr);
scratch_free(sumsq_y_hdr); scratch_free(sumxy_hdr);
scratch_free(sum_hi_hdr);
return result;
}
emit_agg_columns(&result, g, ext, agg_vecs, (uint32_t)K, n_aggs,
(double*)sums, (int64_t*)sums, sums_hi,
NULL, NULL, NULL, NULL,
counts, NULL, prod, sumsq, NULL, sum_y, sumsq_y, sumxy);
scratch_free(sum_hdr); scratch_free(cnt_hdr);
scratch_free(sumsq_hdr); scratch_free(sum_y_hdr);
scratch_free(sumsq_y_hdr); scratch_free(sumxy_hdr);
scratch_free(sum_hi_hdr);
return result;
}
static ray_t* exec_group_v2_exprs(ray_graph_t* g, ray_op_t* op, ray_t* tbl,
int64_t group_limit) {
if (!tbl || tbl->type != RAY_TABLE) return NULL;
ray_op_ext_t* ext = find_ext(g, op->id);
if (!ext || ext->n_keys < 1 || ext->n_aggs < 1) return NULL;
int64_t nrows = ray_table_nrows(tbl);
if (nrows <= 0) return NULL;
uint32_t na = ext->n_aggs, nkk = ext->n_keys;
bool any_expr = false;
for (uint32_t a = 0; a < na; a++) {
if (ext->agg_ops[a] == OP_COUNT) continue;
ray_op_t* in = op_node(g, ext->agg_ins[a]);
if (!in) return NULL;
if (in->opcode == OP_SCAN || in->opcode == OP_CONST) continue;
any_expr = true;
}
if (!any_expr) return NULL;
ray_t* synth[na];
char names[na][16];
for (uint32_t a = 0; a < na; a++) synth[a] = NULL;
ray_t* sub = NULL;
bool ok = true;
for (uint32_t a = 0; a < na && ok; a++) {
if (ext->agg_ops[a] == OP_COUNT) continue;
ray_op_t* in = op_node(g, ext->agg_ins[a]);
if (in->opcode == OP_SCAN || in->opcode == OP_CONST) continue;
snprintf(names[a], sizeof(names[a]), "_e%u", (unsigned)a);
if (ray_table_get_col(tbl, ray_sym_intern(names[a], strlen(names[a])))) {
ok = false; break;
}
ray_expr_t ex;
if (!expr_compile(g, tbl, in, &ex)) { ok = false; break; }
ray_t* v = expr_eval_full(&ex, nrows);
if (!v || RAY_IS_ERR(v) || !ray_is_vec(v) || v->len != nrows) {
if (v && !RAY_IS_ERR(v)) ray_release(v);
else if (v) ray_release(v);
ok = false; break;
}
synth[a] = v;
}
if (ok) {
int64_t ncols = ray_table_ncols(tbl);
sub = ray_table_new(ncols + na);
if (!sub || RAY_IS_ERR(sub)) { ok = false; }
for (int64_t c = 0; c < ncols && ok; c++) {
sub = ray_table_add_col(sub, ray_table_col_name(tbl, c),
ray_table_get_col_idx(tbl, c));
if (!sub || RAY_IS_ERR(sub)) ok = false;
}
for (uint32_t a = 0; a < na && ok; a++) {
if (!synth[a]) continue;
sub = ray_table_add_col(sub,
ray_sym_intern(names[a], strlen(names[a])), synth[a]);
if (!sub || RAY_IS_ERR(sub)) ok = false;
}
}
ray_op_t* op2 = NULL;
if (ok) {
uint32_t in_ids[na];
bool has2 = false, hask = false;
for (uint32_t a = 0; a < na && ok; a++) {
if (synth[a]) {
ray_op_t* sc = ray_scan(g, names[a]);
if (!sc) { ok = false; break; }
in_ids[a] = sc->id;
} else {
in_ids[a] = ext->agg_ins[a];
}
if (ext->agg_ins2 && ext->agg_ins2[a] != RAY_OP_NONE) has2 = true;
if (ext->agg_k && ext->agg_k[a]) hask = true;
}
if (ok) {
ray_op_t* keys[nkk];
ray_op_t* ins[na];
ray_op_t* ins2[na];
for (uint32_t k = 0; k < nkk; k++)
keys[k] = op_node(g, ext->keys[k]);
for (uint32_t a = 0; a < na; a++) {
ins[a] = op_node(g, in_ids[a]);
ins2[a] = (ext->agg_ins2 && ext->agg_ins2[a] != RAY_OP_NONE)
? op_node(g, ext->agg_ins2[a]) : NULL;
}
if (hask || has2)
op2 = ray_group_build(g, keys, ext->n_keys, ext->agg_ops, ins,
has2 ? ins2 : NULL,
hask ? ext->agg_k : NULL,
ext->n_aggs);
else
op2 = ray_group(g, keys, ext->n_keys, ext->agg_ops, ins,
ext->n_aggs);
if (!op2) ok = false;
}
}
ray_t* result = NULL;
if (ok && agg_v2_can_handle(g, op2, sub))
result = exec_group_v2(g, op2, sub, group_limit);
for (uint32_t a = 0; a < na; a++)
if (synth[a]) ray_release(synth[a]);
if (sub) ray_release(sub);
return result;
}
typedef struct {
ray_graph_t* g;
ray_op_ext_t* ext;
void** key_data;
int8_t* key_types;
ray_t** key_vecs;
uint8_t* key_owned;
void** agg_ptrs;
int8_t* agg_types;
uint8_t* agg_strlen;
ray_t** agg_vecs;
uint8_t* agg_owned;
agg_affine_t* agg_affine;
agg_prod_t* agg_prod;
ray_t** strlen_sym_strings;
uint32_t strlen_sym_count;
uint64_t agg_f64_mask;
uint32_t n_aggs;
uint32_t n_keys;
int64_t n_scan;
uint8_t key_esz;
bool sp_need_sum;
bool sp_need_sum128;
const int64_t* match_idx;
ray_t* rowsel;
ray_t* match_idx_block;
ray_t* vla_hdr;
ray_group_emit_filter_t emit_filter;
} sp_dyn_ctx_t;
typedef struct {
const void* key_data;
uint8_t key_esz;
uint32_t* counts;
uint64_t bound;
uint32_t n_workers;
const int64_t* match_idx;
ray_t* rowsel;
_Atomic(int) fail;
} sp_dyn_pcount_ctx_t;
static void sp_dyn_pcount_fn(void* vctx, uint32_t wid, int64_t start,
int64_t end) {
sp_dyn_pcount_ctx_t* c = (sp_dyn_pcount_ctx_t*)vctx;
if (atomic_load_explicit(&c->fail, memory_order_relaxed)) return;
uint32_t* my = c->counts + (size_t)(wid % c->n_workers) * c->bound;
const int64_t* match_idx = c->match_idx;
const uint64_t bound = c->bound;
const uint8_t esz = c->key_esz;
const void* kd = c->key_data;
for (int64_t i = start; i < end; i++) {
int64_t r = match_idx ? match_idx[i] : i;
if (!match_idx && c->rowsel && !group_rowsel_pass(c->rowsel, r))
continue;
int64_t key = read_by_esz(kd, r, esz);
if ((uint64_t)key >= bound) {
atomic_store_explicit(&c->fail, 1, memory_order_relaxed);
return;
}
if (my[key] != UINT32_MAX) my[key]++;
}
}
static ray_t* __attribute__((noinline))
exec_group_sp_dyn_emit(const sp_dyn_ctx_t* c) {
ray_graph_t* g = c->g;
ray_op_ext_t* ext = c->ext;
void** key_data = c->key_data;
int8_t* key_types = c->key_types;
ray_t** key_vecs = c->key_vecs;
uint8_t* key_owned = c->key_owned;
void** agg_ptrs = c->agg_ptrs;
int8_t* agg_types = c->agg_types;
uint8_t* agg_strlen = c->agg_strlen;
ray_t** agg_vecs = c->agg_vecs;
uint8_t* agg_owned = c->agg_owned;
agg_affine_t* agg_affine = c->agg_affine;
agg_prod_t* agg_prod = c->agg_prod;
ray_t** strlen_sym_strings = c->strlen_sym_strings;
uint32_t strlen_sym_count = c->strlen_sym_count;
uint64_t agg_f64_mask = c->agg_f64_mask;
uint32_t n_aggs = c->n_aggs;
uint32_t n_keys = c->n_keys;
int64_t n_scan = c->n_scan;
uint8_t key_esz = c->key_esz;
bool sp_need_sum = c->sp_need_sum;
bool sp_need_sum128 = c->sp_need_sum128;
const int64_t* match_idx = c->match_idx;
ray_t* rowsel = c->rowsel;
ray_t* match_idx_block = c->match_idx_block;
ray_t* vla_hdr = c->vla_hdr;
ray_group_emit_filter_t emit_filter = c->emit_filter;
uint64_t cap = key_esz == 1 ? 256u
: key_esz == 2 ? (1u << 16)
: (1u << 20);
const uint64_t max_dense_cap = 1u << 24;
bool count_only_first = (key_types[0] == RAY_SYM);
ray_t *cnt_hdr = NULL, *range_sum_hdr = NULL, *range_hi_hdr = NULL;
uint32_t* range_count = (uint32_t*)scratch_calloc(
&cnt_hdr, (size_t)cap * sizeof(uint32_t));
da_val_t* range_sum = NULL;
int64_t* range_sum_hi = NULL;
bool dyn_ok = range_count != NULL;
if (dyn_ok && sp_need_sum && !count_only_first) {
range_sum = (da_val_t*)scratch_calloc(
&range_sum_hdr,
(size_t)cap * n_aggs * sizeof(da_val_t));
dyn_ok = range_sum != NULL;
if (dyn_ok && sp_need_sum128) {
range_sum_hi = (int64_t*)scratch_calloc(
&range_hi_hdr,
(size_t)cap * n_aggs * sizeof(int64_t));
dyn_ok = range_sum_hi != NULL;
}
}
uint64_t max_seen = 0;
bool have_dyn_key = false;
bool dyn_par_done = false;
if (dyn_ok && !range_sum && !rowsel && n_scan >= 262144) {
uint64_t bound = 0;
if (key_types[0] == RAY_SYM) {
int64_t dc = ray_sym_domain_count(
ray_sym_vec_domain(key_vecs[0]));
if (dc > 0 && (uint64_t)dc <= max_dense_cap)
bound = (uint64_t)dc;
} else if (key_esz == 1) bound = 256u;
else if (key_esz == 2) bound = 1u << 16;
ray_pool_t* dp = ray_pool_get();
uint32_t dnw = dp ? ray_pool_total_workers(dp) : 1;
if (bound > 0 && dp && dnw >= 2 &&
(uint64_t)dnw * bound * sizeof(uint32_t) <= (512u << 20)) {
ray_t* pc_hdr = NULL;
uint32_t* pc = (uint32_t*)scratch_calloc(&pc_hdr,
(size_t)dnw * bound * sizeof(uint32_t));
if (pc) {
sp_dyn_pcount_ctx_t pctx = {
.key_data = key_data[0],
.key_esz = key_esz,
.counts = pc,
.bound = bound,
.n_workers = dnw,
.match_idx = match_idx,
.rowsel = rowsel,
.fail = 0,
};
ray_pool_dispatch(dp, sp_dyn_pcount_fn, &pctx,
n_scan);
if (!atomic_load_explicit(&pctx.fail,
memory_order_relaxed)) {
if (bound > cap) {
uint32_t* nc = (uint32_t*)scratch_realloc(
&cnt_hdr, (size_t)cap * sizeof(uint32_t),
(size_t)bound * sizeof(uint32_t));
if (nc) {
range_count = nc;
memset(range_count + cap, 0,
(size_t)(bound - cap) * sizeof(uint32_t));
cap = bound;
}
}
if (bound <= cap) {
for (uint32_t w = 0; w < dnw; w++) {
const uint32_t* src = pc + (size_t)w * bound;
for (uint64_t o = 0; o < bound; o++) {
uint64_t s = (uint64_t)range_count[o] + src[o];
range_count[o] = s > UINT32_MAX
? UINT32_MAX : (uint32_t)s;
}
}
for (uint64_t o = 0; o < bound; o++)
if (range_count[o] > 0) {
have_dyn_key = true;
max_seen = o;
}
dyn_par_done = true;
}
}
scratch_free(pc_hdr);
}
}
}
#define DYN_DENSE_ACCUM_ROW(row_expr) \
do { \
int64_t dyn_row = (row_expr); \
int64_t key = read_by_esz(key_data[0], dyn_row, key_esz); \
if (key < 0 || (uint64_t)key >= max_dense_cap) { \
dyn_ok = false; \
goto dyn_dense_done; \
} \
uint64_t off = (uint64_t)key; \
if (off >= cap) { \
uint64_t old_cap = cap; \
while (off >= cap) cap <<= 1; \
uint32_t* new_count = (uint32_t*)scratch_realloc( \
&cnt_hdr, (size_t)old_cap * sizeof(uint32_t), \
(size_t)cap * sizeof(uint32_t)); \
if (!new_count) { \
dyn_ok = false; \
goto dyn_dense_done; \
} \
range_count = new_count; \
memset(range_count + old_cap, 0, \
(size_t)(cap - old_cap) * sizeof(uint32_t)); \
if (sp_need_sum && !count_only_first) { \
da_val_t* new_sum = (da_val_t*)scratch_realloc( \
&range_sum_hdr, \
(size_t)old_cap * n_aggs * sizeof(da_val_t), \
(size_t)cap * n_aggs * sizeof(da_val_t)); \
if (!new_sum) { \
dyn_ok = false; \
goto dyn_dense_done; \
} \
range_sum = new_sum; \
memset(range_sum + (size_t)old_cap * n_aggs, 0, \
(size_t)(cap - old_cap) * n_aggs * sizeof(da_val_t)); \
} \
if (range_sum_hi) { \
int64_t* new_hi = (int64_t*)scratch_realloc( \
&range_hi_hdr, \
(size_t)old_cap * n_aggs * sizeof(int64_t), \
(size_t)cap * n_aggs * sizeof(int64_t)); \
if (!new_hi) { \
dyn_ok = false; \
goto dyn_dense_done; \
} \
range_sum_hi = new_hi; \
memset(range_sum_hi + (size_t)old_cap * n_aggs, 0, \
(size_t)(cap - old_cap) * n_aggs * sizeof(int64_t)); \
} \
} \
have_dyn_key = true; \
if (off > max_seen) max_seen = off; \
if (range_count[off] != UINT32_MAX) range_count[off]++; \
if (range_sum) { \
da_val_t* sums = &range_sum[(size_t)off * n_aggs]; \
for (uint32_t a = 0; a < n_aggs; a++) { \
if (ext->agg_ops[a] == OP_COUNT || !agg_ptrs[a]) continue; \
if (agg_strlen[a]) \
sums[a].i = wrap_add_i64( \
sums[a].i, \
group_strlen_at_cached( \
agg_vecs[a], dyn_row, strlen_sym_strings, strlen_sym_count)); \
else if (agg_f64_mask & ((uint64_t)1 << a)) \
sums[a].f += group_fp_at(agg_ptrs[a], agg_types[a], dyn_row); \
else if (range_sum_hi) \
ray_i128_add(&range_sum_hi[(size_t)off * n_aggs + a], \
(uint64_t*)&sums[a].i, \
read_col_i64(agg_ptrs[a], dyn_row, agg_types[a], 0)); \
else \
sums[a].i = wrap_add_i64( \
sums[a].i, \
read_col_i64(agg_ptrs[a], dyn_row, agg_types[a], 0)); \
} \
} \
} while (0)
if (dyn_par_done) {
} else if (dyn_ok && match_idx) {
for (int64_t i = 0; i < n_scan; i++)
DYN_DENSE_ACCUM_ROW(match_idx[i]);
} else if (dyn_ok && rowsel) {
ray_rowsel_t* m = ray_rowsel_meta(rowsel);
const uint8_t* flags = ray_rowsel_flags(rowsel);
const uint32_t* offs = ray_rowsel_offsets(rowsel);
const uint16_t* idx = ray_rowsel_idx(rowsel);
uint32_t nseg = (uint32_t)((m->nrows + RAY_MORSEL_ELEMS - 1) /
RAY_MORSEL_ELEMS);
for (uint32_t seg = 0; seg < nseg; seg++) {
int64_t base = (int64_t)seg * RAY_MORSEL_ELEMS;
if (flags[seg] == RAY_SEL_NONE) continue;
if (flags[seg] == RAY_SEL_ALL) {
int64_t end = base + RAY_MORSEL_ELEMS;
if (end > m->nrows) end = m->nrows;
for (int64_t r = base; r < end; r++)
DYN_DENSE_ACCUM_ROW(r);
} else {
for (uint32_t p = offs[seg]; p < offs[seg + 1]; p++)
DYN_DENSE_ACCUM_ROW(base + idx[p]);
}
}
} else if (dyn_ok) {
for (int64_t r = 0; r < n_scan; r++)
DYN_DENSE_ACCUM_ROW(r);
}
dyn_dense_done:
#undef DYN_DENSE_ACCUM_ROW
if (dyn_ok && have_dyn_key) {
uint32_t total_groups = 0;
for (uint64_t off = 0; off <= max_seen; off++)
if (range_count[off] > 0)
total_groups++;
int64_t keep_min = da_count_emit_keep_min_u32(
range_count, max_seen + 1, total_groups, emit_filter);
uint32_t grp_count = 0;
for (uint64_t off = 0; off <= max_seen; off++)
if ((int64_t)range_count[off] >= keep_min)
grp_count++;
ray_t* result = ray_table_new((int64_t)n_keys + n_aggs);
if (!result || RAY_IS_ERR(result)) {
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr); scratch_free(cnt_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result ? result : ray_error("oom", NULL);
}
ray_t* key_col = col_vec_new(key_vecs[0], (int64_t)grp_count);
out_col_adopt_str_pool(key_col, key_vecs[0]);
if (key_col && !RAY_IS_ERR(key_col) && key_col->type == RAY_SYM)
ray_sym_vec_adopt_domain(key_col, sym_domain_rep(key_vecs[0]));
if (!key_col || RAY_IS_ERR(key_col)) {
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr); scratch_free(cnt_hdr);
ray_release(result);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return key_col ? key_col : ray_error("oom", NULL);
}
key_col->len = (int64_t)grp_count;
ray_t *_h_sum = NULL, *_h_cnt = NULL, *_h_sum_hi = NULL;
da_val_t* dense_sum = sp_need_sum
? (da_val_t*)scratch_alloc(&_h_sum,
(size_t)grp_count * n_aggs * sizeof(da_val_t))
: NULL;
int64_t* dense_sum_hi = (sp_need_sum && sp_need_sum128)
? (int64_t*)scratch_alloc(&_h_sum_hi,
(size_t)grp_count * n_aggs * sizeof(int64_t))
: NULL;
int64_t* dense_count = (int64_t*)scratch_alloc(
&_h_cnt, (size_t)grp_count * sizeof(int64_t));
if ((sp_need_sum && !dense_sum) || !dense_count ||
(sp_need_sum && sp_need_sum128 && !dense_sum_hi)) {
scratch_free(_h_sum); scratch_free(_h_sum_hi); scratch_free(_h_cnt);
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr); scratch_free(cnt_hdr);
ray_release(key_col); ray_release(result);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return ray_error("oom", NULL);
}
if (sp_need_sum && !range_sum) {
memset(dense_sum, 0,
(size_t)grp_count * n_aggs * sizeof(da_val_t));
if (dense_sum_hi)
memset(dense_sum_hi, 0,
(size_t)grp_count * n_aggs * sizeof(int64_t));
}
uint32_t gi = 0;
for (uint64_t off = 0; off <= max_seen; off++) {
uint32_t cnt = range_count[off];
if ((int64_t)cnt < keep_min) {
if (!range_sum) range_count[off] = 0;
continue;
}
write_col_i64(ray_data(key_col), gi, (int64_t)off,
key_col->type, key_col->attrs);
dense_count[gi] = (int64_t)cnt;
if (range_sum) {
memcpy(&dense_sum[(size_t)gi * n_aggs],
&range_sum[(size_t)off * n_aggs],
(size_t)n_aggs * sizeof(da_val_t));
if (dense_sum_hi)
memcpy(&dense_sum_hi[(size_t)gi * n_aggs],
&range_sum_hi[(size_t)off * n_aggs],
(size_t)n_aggs * sizeof(int64_t));
}
if (!range_sum) range_count[off] = gi + 1u;
gi++;
}
if (sp_need_sum && !range_sum) {
#define DYN_DENSE_SUM_ROW(row_expr) \
do { \
int64_t dyn_row = (row_expr); \
int64_t key = read_by_esz(key_data[0], dyn_row, key_esz); \
if (key < 0 || (uint64_t)key > max_seen) break; \
uint32_t marker = range_count[(uint64_t)key]; \
if (!marker) break; \
da_val_t* sums = &dense_sum[(size_t)(marker - 1u) * n_aggs]; \
for (uint32_t a = 0; a < n_aggs; a++) { \
if (ext->agg_ops[a] == OP_COUNT || !agg_ptrs[a]) continue; \
if (agg_strlen[a]) \
sums[a].i = wrap_add_i64( \
sums[a].i, \
group_strlen_at_cached( \
agg_vecs[a], dyn_row, strlen_sym_strings, strlen_sym_count)); \
else if (agg_f64_mask & ((uint64_t)1 << a)) \
sums[a].f += group_fp_at(agg_ptrs[a], agg_types[a], dyn_row); \
else if (dense_sum_hi) \
ray_i128_add(&dense_sum_hi[(size_t)(marker - 1u) * n_aggs + a], \
(uint64_t*)&sums[a].i, \
read_col_i64(agg_ptrs[a], dyn_row, agg_types[a], 0)); \
else \
sums[a].i = wrap_add_i64( \
sums[a].i, \
read_col_i64(agg_ptrs[a], dyn_row, agg_types[a], 0)); \
} \
} while (0)
if (match_idx) {
for (int64_t i = 0; i < n_scan; i++)
DYN_DENSE_SUM_ROW(match_idx[i]);
} else if (rowsel) {
ray_rowsel_t* m = ray_rowsel_meta(rowsel);
const uint8_t* flags = ray_rowsel_flags(rowsel);
const uint32_t* offs = ray_rowsel_offsets(rowsel);
const uint16_t* idx = ray_rowsel_idx(rowsel);
uint32_t nseg = (uint32_t)((m->nrows + RAY_MORSEL_ELEMS - 1) /
RAY_MORSEL_ELEMS);
for (uint32_t seg = 0; seg < nseg; seg++) {
int64_t base = (int64_t)seg * RAY_MORSEL_ELEMS;
if (flags[seg] == RAY_SEL_NONE) continue;
if (flags[seg] == RAY_SEL_ALL) {
int64_t end = base + RAY_MORSEL_ELEMS;
if (end > m->nrows) end = m->nrows;
for (int64_t r = base; r < end; r++)
DYN_DENSE_SUM_ROW(r);
} else {
for (uint32_t p = offs[seg]; p < offs[seg + 1]; p++)
DYN_DENSE_SUM_ROW(base + idx[p]);
}
}
} else {
for (int64_t r = 0; r < n_scan; r++)
DYN_DENSE_SUM_ROW(r);
}
#undef DYN_DENSE_SUM_ROW
}
ray_op_ext_t* key_ext = find_ext(g, ext->keys[0]);
int64_t name_id = key_ext ? key_ext->sym : 0;
result = ray_table_add_col(result, name_id, key_col);
ray_release(key_col);
emit_agg_columns(&result, g, ext, agg_vecs, grp_count, n_aggs,
(double*)dense_sum, (int64_t*)dense_sum, dense_sum_hi,
NULL, NULL, NULL, NULL,
dense_count, agg_affine, agg_prod, NULL, NULL,
NULL, NULL, NULL);
scratch_free(_h_sum); scratch_free(_h_sum_hi); scratch_free(_h_cnt);
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr); scratch_free(cnt_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result;
}
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr);
scratch_free(cnt_hdr);
return NULL;
}
static ray_t* exec_group_run(ray_graph_t* g, ray_op_t* op, ray_t* tbl,
int64_t group_limit) {
if (!tbl || RAY_IS_ERR(tbl)) return tbl;
if (g->selection) {
ray_rowsel_t* sm = ray_rowsel_meta(g->selection);
int64_t tbl_nrows = ray_table_nrows(tbl);
if (sm->nrows != tbl_nrows)
return ray_error("domain",
"exec_group: selection nrows mismatch (sel=%lld tbl=%lld)",
(long long)sm->nrows, (long long)tbl_nrows);
}
if (g->sg_col) {
ray_t* sgr = exec_group_slices(g, op, tbl, group_limit);
if (sgr) { agg_route_record(AGG_ROUTE_SLICES); sg_hint_release(g); return sgr; }
ray_t* err = sg_hint_to_selection(g, tbl);
if (err) return err;
}
{
int64_t nc = ray_table_ncols(tbl);
for (int64_t c = 0; c < nc; c++) {
ray_t* col = ray_table_get_col_idx(tbl, c);
if (col && (RAY_IS_PARTED(col->type) || col->type == RAY_MAPCOMMON)) {
agg_route_record(AGG_ROUTE_PARTED);
return exec_group_parted(g, op, tbl, group_limit);
}
}
}
ray_op_ext_t* ext = find_ext(g, op->id);
if (!ext) return ray_error("nyi", NULL);
agg_v2_reason_t admission = !ray_agg_engine_v2 ? AGG_V2_DISABLED
: group_limit < 0 ? AGG_V2_SHAPE
: agg_v2_admission(g, op, tbl);
agg_route_reason(admission);
if (admission == AGG_V2_ADMITTED)
return exec_group_v2(g, op, tbl, group_limit);
if (ray_agg_engine_v2 && group_limit >= 0
&& !ray_group_emit_filter_active().enabled) {
ray_t* r = exec_group_v2_exprs(g, op, tbl, group_limit);
if (r) return r;
}
agg_route_record(AGG_ROUTE_LEGACY);
int64_t nrows = ray_table_nrows(tbl);
uint32_t n_keys = ext->n_keys;
uint32_t n_aggs = ext->n_aggs;
if (g->selection && n_keys == 1 && n_aggs > 0 && nrows > 0) {
int64_t cnt_sym_probe = ray_sym_intern("_count", 6);
ray_t* cnt_col_probe = ray_table_get_col(tbl, cnt_sym_probe);
ray_op_ext_t* key_ext_probe = find_ext(g, ext->keys[0]);
int64_t src_sym_probe = ray_sym_intern("_src", 4);
if (cnt_col_probe && cnt_col_probe->type == RAY_I64 &&
key_ext_probe && key_ext_probe->base.opcode == OP_SCAN &&
key_ext_probe->sym == src_sym_probe) {
ray_rowsel_t* sm = ray_rowsel_meta(g->selection);
if (sm->total_pass > 0) {
bool needs_weighting = false;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
ray_op_ext_t* agg_ext = find_ext(g, ext->agg_ins[a]);
if (aop == OP_COUNT) { needs_weighting = true; break; }
if (aop == OP_SUM && agg_ext &&
agg_ext->base.opcode == OP_SCAN &&
agg_ext->sym == cnt_sym_probe) {
needs_weighting = true; break;
}
}
if (needs_weighting)
return ray_error("nyi",
"GROUP BY with selection on factorized expand result "
"(COUNT/SUM(_count) semantics)");
}
}
}
if (!g->selection && n_keys == 1 && n_aggs > 0 && nrows > 0) {
int64_t cnt_sym = ray_sym_intern("_count", 6);
ray_t* cnt_col = ray_table_get_col(tbl, cnt_sym);
if (cnt_col && cnt_col->type == RAY_I64) {
ray_op_ext_t* key_ext = find_ext(g, ext->keys[0]);
int64_t src_sym = ray_sym_intern("_src", 4);
if (key_ext && key_ext->base.opcode == OP_SCAN &&
key_ext->sym == src_sym) {
bool all_compat = true;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
ray_op_ext_t* agg_ext = find_ext(g, ext->agg_ins[a]);
if (aop == OP_COUNT) continue;
if (aop == OP_SUM && agg_ext &&
agg_ext->base.opcode == OP_SCAN &&
agg_ext->sym == cnt_sym) continue;
all_compat = false;
break;
}
if (all_compat) {
ray_t* src_col = ray_table_get_col(tbl, src_sym);
if (src_col) {
int64_t out_nkeys = 1;
int64_t out_ncols = out_nkeys + n_aggs;
ray_t* result = ray_table_new((int64_t)out_ncols);
if (!result || RAY_IS_ERR(result))
return ray_error("oom", NULL);
ray_retain(src_col);
ray_t* tmp_r = ray_table_add_col(result, src_sym, src_col);
ray_release(src_col);
if (!tmp_r || RAY_IS_ERR(tmp_r)) {
ray_release(result);
return ray_error("oom", NULL);
}
result = tmp_r;
for (uint32_t a = 0; a < n_aggs; a++) {
ray_retain(cnt_col);
int64_t agg_name = ray_sym_intern("_agg", 4);
if (n_aggs > 1) {
char buf[16];
int n = snprintf(buf, sizeof(buf), "_agg%u", a);
agg_name = ray_sym_intern(buf, (size_t)n);
}
tmp_r = ray_table_add_col(result, agg_name, cnt_col);
ray_release(cnt_col);
if (!tmp_r || RAY_IS_ERR(tmp_r)) {
ray_release(result);
return ray_error("oom", NULL);
}
result = tmp_r;
}
return result;
}
}
}
}
}
ray_t* match_idx_block = NULL;
const int64_t* match_idx = NULL;
ray_t* rowsel = NULL;
int64_t n_scan = nrows;
if (g->selection) {
rowsel = g->selection;
}
uint32_t vla_keys = n_keys > 0 ? n_keys : 1;
ray_t* key_vecs[vla_keys];
memset(key_vecs, 0, vla_keys * sizeof(ray_t*));
uint8_t key_owned[vla_keys];
memset(key_owned, 0, vla_keys * sizeof(uint8_t));
int64_t key_scan_sym[vla_keys];
for (uint32_t k = 0; k < vla_keys; k++) key_scan_sym[k] = -1;
for (uint32_t k = 0; k < n_keys; k++) {
ray_op_t* key_op = op_node(g, ext->keys[k]);
ray_op_ext_t* key_ext = find_ext(g, key_op->id);
if (key_ext && key_ext->base.opcode == OP_SCAN) {
key_vecs[k] = ray_table_get_col(tbl, key_ext->sym);
key_scan_sym[k] = key_ext->sym;
} else {
ray_t* saved_table = g->table;
g->table = tbl;
ray_t* vec = exec_node(g, key_op);
g->table = saved_table;
if (vec && !RAY_IS_ERR(vec)) {
int64_t tn = ray_table_nrows(tbl);
bool scalar = ray_is_atom(vec) || (ray_is_vec(vec) && vec->len == 1 && tn != 1);
if (scalar) {
ray_t* wide;
if (vec->type == -RAY_STR) {
wide = broadcast_scalar(vec, tn);
} else {
ray_t* n_obj = make_i64(tn);
wide = n_obj ? ray_take_fn(vec, n_obj) : ray_error("oom", NULL);
if (n_obj) ray_release(n_obj);
}
ray_release(vec);
vec = wide;
} else if (ray_is_vec(vec) && vec->len != tn) {
int64_t got = vec->len;
ray_release(vec);
vec = ray_error("length", "by: key has %lld values for %lld rows",
(long long)got, (long long)tn);
}
if (!vec || RAY_IS_ERR(vec)) {
for (uint32_t j = 0; j < k; j++)
if (key_owned[j] && key_vecs[j]) ray_release(key_vecs[j]);
return vec ? vec : ray_error("oom", NULL);
}
key_vecs[k] = vec;
key_owned[k] = 1;
}
}
if (!key_vecs[k]) {
for (uint32_t j = 0; j < k; j++)
if (key_owned[j] && key_vecs[j]) ray_release(key_vecs[j]);
return ray_error("domain", "by: column not found in table");
}
}
uint32_t vla_aggs = n_aggs > 0 ? n_aggs : 1;
ray_t* vla_hdr = NULL;
size_t vla_bytes = (size_t)vla_aggs * (2 * sizeof(ray_t*)
+ sizeof(agg_affine_t) + sizeof(agg_linear_t)
+ sizeof(agg_prod_t) + 3 * sizeof(uint8_t));
uint8_t* vla_blk = (uint8_t*)scratch_calloc(&vla_hdr, vla_bytes);
if (!vla_blk) {
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
return ray_error("oom", NULL);
}
ray_t** agg_vecs = (ray_t**)vla_blk;
ray_t** agg_vecs2 = agg_vecs + vla_aggs;
agg_affine_t* agg_affine = (agg_affine_t*)(agg_vecs2 + vla_aggs);
agg_linear_t* agg_linear = (agg_linear_t*)(agg_affine + vla_aggs);
agg_prod_t* agg_prod = (agg_prod_t*)(agg_linear + vla_aggs);
uint8_t* agg_owned = (uint8_t*)(agg_prod + vla_aggs);
uint8_t* agg_owned2 = agg_owned + vla_aggs;
uint8_t* agg_strlen = agg_owned2 + vla_aggs;
for (uint32_t a = 0; a < n_aggs; a++) {
ray_op_t* agg_input_op = op_node(g, ext->agg_ins[a]);
ray_op_ext_t* agg_ext = find_ext(g, agg_input_op->id);
uint16_t agg_kind = ext->agg_ops[a];
if ((agg_kind == OP_SUM || agg_kind == OP_AVG) &&
try_affine_sumavg_input(g, tbl, agg_input_op, &agg_vecs[a], &agg_affine[a])) {
continue;
}
if ((agg_kind == OP_SUM || agg_kind == OP_AVG) &&
try_strlen_sumavg_input(g, tbl, agg_input_op, &agg_vecs[a])) {
agg_strlen[a] = 1;
continue;
}
if (n_keys == 0 && nrows > 0 &&
(agg_kind == OP_SUM || agg_kind == OP_AVG) &&
try_linear_sumavg_input_i64(g, tbl, agg_input_op, &agg_linear[a])) {
continue;
}
if ((agg_kind == OP_SUM || agg_kind == OP_AVG) && nrows > 0 &&
try_prod_sumavg_input_f64(g, tbl, agg_input_op, &agg_prod[a])) {
continue;
}
if (agg_ext && agg_ext->base.opcode == OP_SCAN) {
agg_vecs[a] = ray_table_get_col(tbl, agg_ext->sym);
} else if (agg_ext && agg_ext->base.opcode == OP_CONST && agg_ext->literal) {
agg_vecs[a] = agg_ext->literal;
} else {
ray_expr_t agg_expr;
if (expr_compile(g, tbl, agg_input_op, &agg_expr)) {
ray_t* vec = expr_eval_full(&agg_expr, nrows);
if (vec && !RAY_IS_ERR(vec)) {
agg_vecs[a] = vec;
agg_owned[a] = 1;
goto resolve_ins2;
}
if (vec && RAY_IS_ERR(vec)) ray_release(vec);
}
ray_t* saved_table = g->table;
g->table = tbl;
ray_t* vec = exec_node(g, agg_input_op);
g->table = saved_table;
if (vec && RAY_IS_ERR(vec)) {
for (uint32_t i = 0; i < a; i++)
{ if (agg_owned[i] && agg_vecs[i]) ray_release(agg_vecs[i]); if (agg_owned2[i] && agg_vecs2[i]) ray_release(agg_vecs2[i]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
scratch_free(vla_hdr);
return vec;
}
if (vec) {
agg_vecs[a] = vec;
agg_owned[a] = 1;
}
}
resolve_ins2:;
if (ext->agg_ins2 && ext->agg_ins2[a] != RAY_OP_NONE) {
ray_op_t* agg_input_op2 = op_node(g, ext->agg_ins2[a]);
ray_op_ext_t* agg_ext2 = find_ext(g, agg_input_op2->id);
if (agg_ext2 && agg_ext2->base.opcode == OP_SCAN) {
agg_vecs2[a] = ray_table_get_col(tbl, agg_ext2->sym);
} else if (agg_ext2 && agg_ext2->base.opcode == OP_CONST && agg_ext2->literal) {
agg_vecs2[a] = agg_ext2->literal;
} else {
ray_expr_t agg_expr2;
int compiled2 = 0;
if (expr_compile(g, tbl, agg_input_op2, &agg_expr2)) {
ray_t* vec = expr_eval_full(&agg_expr2, nrows);
if (vec && !RAY_IS_ERR(vec)) {
agg_vecs2[a] = vec;
agg_owned2[a] = 1;
compiled2 = 1;
} else if (vec) {
ray_release(vec);
}
}
if (!compiled2) {
ray_t* saved_table = g->table;
g->table = tbl;
ray_t* vec = exec_node(g, agg_input_op2);
g->table = saved_table;
if (vec && RAY_IS_ERR(vec)) {
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t i = 0; i < a; i++)
{ if (agg_owned[i] && agg_vecs[i]) ray_release(agg_vecs[i]); if (agg_owned2[i] && agg_vecs2[i]) ray_release(agg_vecs2[i]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
scratch_free(vla_hdr);
return vec;
}
if (vec) {
agg_vecs2[a] = vec;
agg_owned2[a] = 1;
}
}
}
}
}
for (uint32_t a = 0; a < n_aggs; a++) {
if (!agg_vecs[a] || RAY_IS_ERR(agg_vecs[a])) continue;
if (ext->agg_ops[a] == OP_COUNT) continue;
if (agg_strlen[a]) continue;
bool needs_broadcast = ray_is_atom(agg_vecs[a]) ||
(agg_vecs[a]->type > 0 && agg_vecs[a]->len == 1 && nrows > 1);
if (!needs_broadcast) continue;
ray_t* bcast = materialize_broadcast_input(agg_vecs[a], nrows);
if (!bcast || RAY_IS_ERR(bcast)) {
for (uint32_t i = 0; i < n_aggs; i++) {
{ if (agg_owned[i] && agg_vecs[i]) ray_release(agg_vecs[i]); if (agg_owned2[i] && agg_vecs2[i]) ray_release(agg_vecs2[i]); }
}
for (uint32_t k = 0; k < n_keys; k++) {
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
}
scratch_free(vla_hdr);
return bcast && RAY_IS_ERR(bcast) ? bcast : ray_error("oom", NULL);
}
if (agg_owned[a]) ray_release(agg_vecs[a]);
agg_vecs[a] = bcast;
agg_owned[a] = 1;
}
void* key_data[vla_keys];
int8_t key_types[vla_keys];
uint8_t key_attrs[vla_keys];
for (uint32_t k = 0; k < n_keys; k++) {
if (key_vecs[k]) {
key_data[k] = ray_data(key_vecs[k]);
key_types[k] = key_vecs[k]->type;
key_attrs[k] = key_vecs[k]->attrs;
} else {
key_data[k] = NULL;
key_types[k] = 0;
key_attrs[k] = 0;
}
}
ray_group_emit_filter_t emit_filter = ray_group_emit_filter_active();
bool use_emit_filter = emit_filter.enabled &&
emit_filter.agg_index < n_aggs &&
ext->agg_ops[emit_filter.agg_index] == OP_COUNT &&
(emit_filter.top_count_take <= 0 || emit_filter.desc);
bool use_topn_filter = emit_filter.enabled &&
emit_filter.top_count_take > 0 &&
emit_filter.agg_index < n_aggs &&
(ext->agg_ops[emit_filter.agg_index] == OP_COUNT ||
ext->agg_ops[emit_filter.agg_index] == OP_SUM ||
ext->agg_ops[emit_filter.agg_index] == OP_MIN ||
ext->agg_ops[emit_filter.agg_index] == OP_MAX);
bool sc_has_binary = false;
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_is_binary_agg(ext->agg_ops[a])) { sc_has_binary = true; break; }
if (n_keys == 0 && nrows > 0 && !sc_has_binary) {
uint8_t need_flags = DA_NEED_COUNT;
bool has_first_last = false;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop == OP_SUM || aop == OP_PROD || aop == OP_AVG || aop == OP_ALL || aop == OP_ANY || aop == OP_FIRST || aop == OP_LAST)
need_flags |= DA_NEED_SUM;
else if (aop == OP_STDDEV || aop == OP_STDDEV_POP || aop == OP_VAR || aop == OP_VAR_POP)
{ need_flags |= DA_NEED_SUM; need_flags |= DA_NEED_SUMSQ; }
else if (aop == OP_MIN) need_flags |= DA_NEED_MIN;
else if (aop == OP_MAX) need_flags |= DA_NEED_MAX;
if (aop == OP_FIRST || aop == OP_LAST) has_first_last = true;
}
ray_t* sc_vla_hdr = NULL;
size_t sc_vla_bytes = (size_t)vla_aggs * (sizeof(void*) + sizeof(int64_t)
+ sizeof(int8_t) + sizeof(bool));
uint8_t* sc_blk = (uint8_t*)scratch_alloc(&sc_vla_hdr, sc_vla_bytes);
if (!sc_blk) {
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return ray_error("oom", NULL);
}
void* *agg_ptrs = (void**)sc_blk;
int64_t *sc_int_null_sentinel = (int64_t*)(agg_ptrs + vla_aggs);
int8_t *agg_types = (int8_t*)(sc_int_null_sentinel + vla_aggs);
bool *sc_int_null_has = (bool*)(agg_types + vla_aggs);
bool sc_any_nullable = false;
bool sc_need_sum128 = false;
for (uint32_t a = 0; a < n_aggs; a++) {
if (agg_prod[a].enabled) {
agg_ptrs[a] = NULL;
agg_types[a] = RAY_F64;
sc_int_null_sentinel[a] = 0;
sc_int_null_has[a] = false;
continue;
}
if (agg_vecs[a]) {
agg_ptrs[a] = ray_data(agg_vecs[a]);
agg_types[a] = agg_vecs[a]->type;
sc_int_null_sentinel[a] = agg_int_null_sentinel_for(agg_vecs[a]->type);
int8_t t = agg_vecs[a]->type;
bool is_sentinel_typed = (t == RAY_I16 || t == RAY_I32 || t == RAY_I64 ||
t == RAY_DATE || t == RAY_TIME || t == RAY_TIMESTAMP || t == RAY_SYM);
sc_int_null_has[a] = is_sentinel_typed && ray_vec_may_have_nulls(agg_vecs[a]);
if (ray_vec_may_have_nulls(agg_vecs[a]) &&
(group_fp_type(agg_vecs[a]->type) || is_sentinel_typed))
sc_any_nullable = true;
} else {
agg_ptrs[a] = NULL;
agg_types[a] = 0;
sc_int_null_sentinel[a] = 0;
sc_int_null_has[a] = false;
}
if (ext->agg_ops[a] == OP_AVG && !group_fp_type(agg_types[a]) &&
!agg_strlen[a] && group_avg_needs_hi(agg_types[a], nrows))
sc_need_sum128 = true;
}
if (!match_idx && !rowsel && !sc_any_nullable && n_aggs > 1) {
ray_t* base_col = NULL;
const void* base_ptr = NULL;
int8_t base_type = 0;
int64_t base_len = 0;
bool one_base_input = true;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop != OP_SUM && aop != OP_AVG) { one_base_input = false; break; }
if (agg_prod[a].enabled || agg_strlen[a]) {
one_base_input = false;
break;
}
const void* slot_ptr = NULL;
int8_t slot_type = 0;
int64_t slot_len = nrows;
if (agg_vecs[a] && agg_ptrs[a]) {
if (!agg_type_admitted(OP_SUM, agg_vecs[a]->type)) {
one_base_input = false;
break;
}
slot_ptr = agg_ptrs[a];
slot_type = agg_vecs[a]->type;
slot_len = agg_vecs[a]->len;
if (!base_col) base_col = agg_vecs[a];
} else if (agg_linear[a].enabled &&
agg_linear[a].n_terms == 1 &&
agg_linear[a].coeff_i64[0] == 1 &&
agg_linear[a].bias_i64 == 0 &&
agg_linear[a].term_ptrs[0]) {
if (!agg_type_admitted(OP_SUM, agg_linear[a].term_types[0])) {
one_base_input = false;
break;
}
slot_ptr = agg_linear[a].term_ptrs[0];
slot_type = agg_linear[a].term_types[0];
} else {
one_base_input = false;
break;
}
if (!base_ptr) {
base_ptr = slot_ptr;
base_type = slot_type;
base_len = slot_len;
} else if (slot_ptr != base_ptr || slot_type != base_type ||
slot_len != base_len) {
one_base_input = false;
break;
}
}
int64_t zone_hi = 0, zone_nn = 0; uint64_t zone_lo = 0;
bool zone_128 = false;
if (one_base_input && base_col && sc_need_sum128 &&
!group_fp_type(base_type)) {
zone_128 = ray_zone_int_sum128(base_col, &zone_hi, &zone_lo, &zone_nn)
&& zone_nn == nrows;
if (!zone_128) one_base_input = false;
}
if (one_base_input && base_col) {
ray_t* base_sum_obj = zone_128 ? ray_i64((int64_t)zone_lo)
: ray_sum_fn(base_col);
if (!base_sum_obj || RAY_IS_ERR(base_sum_obj)) {
scratch_free(sc_vla_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) ray_release(match_idx_block);
scratch_free(vla_hdr);
return base_sum_obj ? base_sum_obj : ray_error("oom", NULL);
}
if (ray_is_lazy(base_sum_obj))
base_sum_obj = ray_lazy_materialize(base_sum_obj);
if (!base_sum_obj || RAY_IS_ERR(base_sum_obj)) {
scratch_free(sc_vla_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) ray_release(match_idx_block);
scratch_free(vla_hdr);
return base_sum_obj ? base_sum_obj : ray_error("oom", NULL);
}
bool base_is_f64 = false;
double base_sum_f64 = 0.0;
int64_t base_sum_i64 = 0;
if (sum_atom_value(base_sum_obj, &base_is_f64,
&base_sum_f64, &base_sum_i64)) {
ray_t *sum_hdr = NULL, *cnt_hdr = NULL;
double* sums_f64 = NULL;
int64_t* sums_i64 = NULL;
int64_t* sums_hi = NULL;
if (base_is_f64)
sums_f64 = (double*)scratch_alloc(&sum_hdr,
(size_t)n_aggs * sizeof(double));
else
sums_i64 = (int64_t*)scratch_alloc(&sum_hdr,
(size_t)2 * n_aggs * sizeof(int64_t));
int64_t* counts = (int64_t*)scratch_alloc(&cnt_hdr,
sizeof(int64_t));
if ((base_is_f64 ? (sums_f64 != NULL) : (sums_i64 != NULL)) &&
counts) {
if (base_is_f64) {
for (uint32_t a = 0; a < n_aggs; a++)
sums_f64[a] = base_sum_f64;
} else {
for (uint32_t a = 0; a < n_aggs; a++)
sums_i64[a] = base_sum_i64;
if (zone_128) {
sums_hi = sums_i64 + n_aggs;
for (uint32_t a = 0; a < n_aggs; a++)
sums_hi[a] = zone_hi;
}
}
counts[0] = nrows;
ray_t* result = ray_table_new(n_aggs);
if (!result || RAY_IS_ERR(result)) {
scratch_free(sum_hdr);
scratch_free(cnt_hdr);
ray_release(base_sum_obj);
scratch_free(sc_vla_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) ray_release(match_idx_block);
scratch_free(vla_hdr);
return result ? result : ray_error("oom", NULL);
}
emit_agg_columns(&result, g, ext, agg_vecs, 1, n_aggs,
sums_f64, sums_i64, sums_hi,
NULL, NULL, NULL, NULL,
counts, agg_affine, agg_prod,
NULL, NULL, NULL, NULL, NULL);
scratch_free(sum_hdr);
scratch_free(cnt_hdr);
ray_release(base_sum_obj);
scratch_free(sc_vla_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) ray_release(match_idx_block);
scratch_free(vla_hdr);
return result;
}
scratch_free(sum_hdr);
scratch_free(cnt_hdr);
}
ray_release(base_sum_obj);
}
}
ray_pool_t* sc_pool = ray_pool_get();
uint32_t sc_n = (sc_pool && nrows >= RAY_PARALLEL_THRESHOLD && !has_first_last)
? ray_pool_total_workers(sc_pool) : 1;
ray_t* sc_hdr;
da_accum_t* sc_acc = (da_accum_t*)scratch_calloc(&sc_hdr,
sc_n * sizeof(da_accum_t));
if (!sc_acc) { scratch_free(sc_vla_hdr); goto da_path; }
bool alloc_ok = true;
const size_t sc_line = 64;
size_t sc_words = 1;
if (need_flags & DA_NEED_SUM) sc_words += n_aggs;
if (sc_need_sum128) sc_words += n_aggs;
if (need_flags & DA_NEED_MIN) sc_words += n_aggs;
if (need_flags & DA_NEED_MAX) sc_words += n_aggs;
if (need_flags & DA_NEED_SUMSQ) sc_words += n_aggs;
if (sc_any_nullable) sc_words += n_aggs;
size_t sc_bytes = ((sc_words * sizeof(int64_t)) + sc_line - 1) / sc_line * sc_line;
for (uint32_t w = 0; w < sc_n; w++) {
void* raw = scratch_calloc(&sc_acc[w]._h_sum, sc_bytes + sc_line);
if (!raw) { alloc_ok = false; break; }
int64_t* blk = (int64_t*)(void*)(((uintptr_t)raw + sc_line - 1) & ~(uintptr_t)(sc_line - 1));
size_t off = 0;
sc_acc[w].count = blk + off; off += 1;
if (need_flags & DA_NEED_SUM) {
sc_acc[w].sum = (da_val_t*)(void*)(blk + off); off += n_aggs;
}
if (sc_need_sum128) {
sc_acc[w].sum_hi = blk + off; off += n_aggs;
}
if (need_flags & DA_NEED_MIN) {
sc_acc[w].min_val = (da_val_t*)(void*)(blk + off); off += n_aggs;
for (uint32_t a = 0; a < n_aggs; a++) {
if (group_fp_type(agg_types[a])) sc_acc[w].min_val[a].f = DBL_MAX;
else sc_acc[w].min_val[a].i = INT64_MAX;
}
}
if (need_flags & DA_NEED_MAX) {
sc_acc[w].max_val = (da_val_t*)(void*)(blk + off); off += n_aggs;
for (uint32_t a = 0; a < n_aggs; a++) {
if (group_fp_type(agg_types[a])) sc_acc[w].max_val[a].f = -DBL_MAX;
else sc_acc[w].max_val[a].i = INT64_MIN;
}
}
if (need_flags & DA_NEED_SUMSQ) {
sc_acc[w].sumsq_f64 = (double*)(void*)(blk + off); off += n_aggs;
}
if (sc_any_nullable) {
sc_acc[w].nn_count = blk + off; off += n_aggs;
}
}
if (!alloc_ok) {
for (uint32_t w = 0; w < sc_n; w++) da_accum_free(&sc_acc[w]);
scratch_free(sc_hdr);
scratch_free(sc_vla_hdr);
goto da_path;
}
scalar_ctx_t sc_ctx = {
.agg_ptrs = agg_ptrs,
.agg_types = agg_types,
.agg_cols = agg_vecs,
.agg_strlen = agg_strlen,
.agg_ops = ext->agg_ops,
.agg_linear = agg_linear,
.agg_prod = agg_prod,
.n_aggs = n_aggs,
.need_flags = need_flags,
.match_idx = match_idx,
.rowsel = rowsel,
.accums = sc_acc,
.n_accums = sc_n,
.agg_int_null_has = sc_int_null_has,
.agg_int_null_sentinel = sc_int_null_sentinel,
};
typedef void (*scalar_fn_t)(void*, uint32_t, int64_t, int64_t);
scalar_fn_t sc_fn = scalar_accum_fn;
bool agg0_has_nulls = n_aggs > 0 &&
(sc_int_null_has[0] ||
(agg_vecs[0] && group_fp_type(agg_vecs[0]->type) &&
ray_vec_may_have_nulls(agg_vecs[0])));
if (n_aggs == 1 && !match_idx && !rowsel && agg_ptrs[0] != NULL && !agg0_has_nulls) {
uint16_t op0 = ext->agg_ops[0];
int8_t t0 = agg_types[0];
if ((op0 == OP_SUM || op0 == OP_AVG) &&
(t0 == RAY_I64 || t0 == RAY_SYM || t0 == RAY_TIMESTAMP))
sc_fn = scalar_sum_i64_fn;
else if ((op0 == OP_SUM || op0 == OP_AVG) && t0 == RAY_F64)
sc_fn = scalar_sum_f64_fn;
} else if (n_aggs == 1 && !match_idx && !rowsel && agg_linear[0].enabled) {
uint16_t op0 = ext->agg_ops[0];
if (op0 == OP_SUM || op0 == OP_AVG)
sc_fn = scalar_sum_linear_i64_fn;
}
if (sc_n > 1)
ray_pool_dispatch(sc_pool, sc_fn, &sc_ctx, n_scan);
else
sc_fn(&sc_ctx, 0, 0, n_scan);
da_accum_t* m = &sc_acc[0];
for (uint32_t w = 1; w < sc_n; w++) {
da_accum_t* wa = &sc_acc[w];
if (need_flags & DA_NEED_SUM) {
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t merge_op = ext->agg_ops[a];
int64_t mnn = m->nn_count ? m->nn_count[a] : m->count[0];
int64_t wnn = wa->nn_count ? wa->nn_count[a] : wa->count[0];
if (merge_op == OP_FIRST) {
if (mnn == 0 && wnn > 0)
m->sum[a] = wa->sum[a];
} else if (merge_op == OP_LAST) {
if (wnn > 0)
m->sum[a] = wa->sum[a];
} else if (merge_op == OP_PROD) {
if (wnn > 0) {
if (mnn == 0)
m->sum[a] = wa->sum[a];
else if (group_fp_type(agg_types[a]))
m->sum[a].f *= wa->sum[a].f;
else
m->sum[a].i = (int64_t)((uint64_t)m->sum[a].i * (uint64_t)wa->sum[a].i);
}
} else {
if (group_fp_type(agg_types[a]))
m->sum[a].f += wa->sum[a].f;
else if (m->sum_hi)
ray_i128_add128(&m->sum_hi[a], (uint64_t*)&m->sum[a].i,
wa->sum_hi[a], (uint64_t)wa->sum[a].i);
else
m->sum[a].i = wrap_add_i64(m->sum[a].i, wa->sum[a].i);
}
}
}
if (need_flags & DA_NEED_SUMSQ) {
for (uint32_t a = 0; a < n_aggs; a++)
m->sumsq_f64[a] += wa->sumsq_f64[a];
}
if (need_flags & DA_NEED_MIN) {
for (uint32_t a = 0; a < n_aggs; a++) {
if (group_fp_type(agg_types[a])) {
if (wa->min_val[a].f < m->min_val[a].f)
m->min_val[a].f = wa->min_val[a].f;
} else if (agg_types[a] == RAY_SYM) {
if (wa->min_val[a].i != INT64_MAX &&
(m->min_val[a].i == INT64_MAX ||
sym_lex_lt(ray_sym_vec_domain(agg_vecs[a]),
wa->min_val[a].i, m->min_val[a].i)))
m->min_val[a].i = wa->min_val[a].i;
} else {
if (wa->min_val[a].i < m->min_val[a].i)
m->min_val[a].i = wa->min_val[a].i;
}
}
}
if (need_flags & DA_NEED_MAX) {
for (uint32_t a = 0; a < n_aggs; a++) {
if (group_fp_type(agg_types[a])) {
if (wa->max_val[a].f > m->max_val[a].f)
m->max_val[a].f = wa->max_val[a].f;
} else if (agg_types[a] == RAY_SYM) {
if (wa->max_val[a].i != INT64_MIN &&
(m->max_val[a].i == INT64_MIN ||
sym_lex_gt(ray_sym_vec_domain(agg_vecs[a]),
wa->max_val[a].i, m->max_val[a].i)))
m->max_val[a].i = wa->max_val[a].i;
} else {
if (wa->max_val[a].i > m->max_val[a].i)
m->max_val[a].i = wa->max_val[a].i;
}
}
}
if (m->nn_count && wa->nn_count) {
for (uint32_t a = 0; a < n_aggs; a++)
m->nn_count[a] += wa->nn_count[a];
}
m->count[0] += wa->count[0];
}
for (uint32_t w = 1; w < sc_n; w++) da_accum_free(&sc_acc[w]);
ray_t* result = ray_table_new(n_aggs);
if (!result || RAY_IS_ERR(result)) {
da_accum_free(&sc_acc[0]); scratch_free(sc_hdr);
scratch_free(sc_vla_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result ? result : ray_error("oom", NULL);
}
emit_agg_columns(&result, g, ext, agg_vecs, 1, n_aggs,
(double*)m->sum, (int64_t*)m->sum, m->sum_hi,
(double*)m->min_val, (double*)m->max_val,
(int64_t*)m->min_val, (int64_t*)m->max_val,
m->count, agg_affine, agg_prod, m->sumsq_f64, m->nn_count,
NULL, NULL, NULL);
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (!(agg_vecs[a] &&
agg_needs_row_gather(agg_vecs[a]->type, aop)))
continue;
ray_t* wsel_blk = NULL; const int64_t* wsel = NULL; int64_t wscan = nrows;
if (g->selection) {
ray_rowsel_t* sm = ray_rowsel_meta(g->selection);
if (sm && sm->nrows == nrows) {
wsel_blk = ray_rowsel_to_indices(g->selection);
wsel = wsel_blk ? (const int64_t*)ray_data(wsel_blk) : NULL;
wscan = sm->total_pass;
}
}
bool hn = ray_vec_may_have_nulls(agg_vecs[a]);
ray_t* atom = agg_wide_reduce(agg_vecs[a], aop, wsel, wscan, hn);
if (wsel_blk) ray_release(wsel_blk);
ray_t* col = agg_vecs[a]->type == RAY_LIST
? ray_list_new(1)
: col_vec_new(agg_vecs[a], 1);
if (col && !RAY_IS_ERR(col)) {
col->len = 1;
if (RAY_ATOM_IS_NULL(atom)) {
ray_vec_set_null(col, 0, true);
} else if (agg_vecs[a]->type == RAY_LIST) {
((ray_t**)ray_data(col))[0] = atom;
atom = NULL;
} else if (agg_vecs[a]->type == RAY_STR) {
ray_t* nv = ray_str_vec_set(col, 0, ray_str_ptr(atom), ray_str_len(atom));
if (nv && !RAY_IS_ERR(nv)) col = nv;
} else {
store_typed_elem(col, 0, atom);
}
ray_table_set_col_idx(result, a, col);
ray_release(col);
}
if (atom) ray_release(atom);
}
{
bool any_rank = false;
for (uint32_t a = 0; a < n_aggs; a++)
if (ext->agg_ops[a] == OP_MEDIAN ||
ext->agg_ops[a] == OP_QUANTILE ||
ext->agg_ops[a] == OP_MODE) { any_rank = true; break; }
if (any_rank) {
ray_t* hsel_blk = NULL; const int64_t* hsel = NULL; int64_t hscan = nrows;
if (g->selection) {
ray_rowsel_t* sm = ray_rowsel_meta(g->selection);
if (sm && sm->nrows == nrows) {
hsel_blk = ray_rowsel_to_indices(g->selection);
hsel = hsel_blk ? (const int64_t*)ray_data(hsel_blk) : NULL;
hscan = sm->total_pass;
}
}
ray_t* ix_hdr = NULL;
int64_t* idxb = (int64_t*)scratch_alloc(&ix_hdr,
(size_t)(hscan > 0 ? hscan : 1) * sizeof(int64_t));
if (idxb) {
if (hsel) { for (int64_t i = 0; i < hscan; i++) idxb[i] = hsel[i]; }
else { for (int64_t i = 0; i < hscan; i++) idxb[i] = i; }
int64_t offs[1] = {0};
int64_t cnts[1] = {hscan};
for (uint32_t a = 0; a < n_aggs; a++) {
if (!agg_vecs[a]) continue;
uint16_t aop = ext->agg_ops[a];
ray_t* hv = NULL;
if (aop == OP_MEDIAN) {
hv = ray_median_per_group_buf(agg_vecs[a], idxb, offs, cnts, 1);
} else if (aop == OP_QUANTILE) {
double q = ext->agg_k
? decode_agg_f64_param(ext->agg_k[a]) : 0.5;
hv = ray_quantile_per_group_buf(agg_vecs[a], idxb, offs, cnts, 1, q);
} else if (aop == OP_MODE) {
hv = ray_mode_per_group_buf(agg_vecs[a], idxb, offs, cnts, 1);
} else {
continue;
}
if (hv && !RAY_IS_ERR(hv)) ray_table_set_col_idx(result, a, hv);
if (hv) ray_release(hv);
}
scratch_free(ix_hdr);
}
if (hsel_blk) ray_release(hsel_blk);
}
}
da_accum_free(&sc_acc[0]); scratch_free(sc_hdr);
scratch_free(sc_vla_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result;
}
da_path:;
{
bool da_eligible = (n_scan > 0 && n_keys > 0 && n_keys <= 8 &&
n_aggs <= 64);
if (da_eligible && rowsel && !match_idx) {
ray_rowsel_t* sm = ray_rowsel_meta(rowsel);
if (sm && sm->total_pass * 4 < nrows)
da_eligible = false;
}
uint64_t da_slot_budget = (uint64_t)(n_scan >> 3);
if (da_slot_budget < 262144) da_slot_budget = 262144;
for (uint32_t a = 0; a < n_aggs && da_eligible; a++) {
uint16_t aop = ext->agg_ops[a];
if (agg_is_binary_agg(aop)) {
ray_t* yv2 = (agg_vecs2 ? agg_vecs2[a] : NULL);
if (!yv2) { da_eligible = false; }
else if (yv2->len < n_scan) {
da_eligible = false;
}
else {
int8_t t2 = yv2->type;
bool y_fp = group_fp_type(t2);
bool y_int = (t2 == RAY_I16 || t2 == RAY_I32 || t2 == RAY_I64 ||
t2 == RAY_DATE || t2 == RAY_TIME || t2 == RAY_TIMESTAMP ||
t2 == RAY_BOOL || t2 == RAY_U8);
ray_t* ysrc = (yv2->attrs & RAY_ATTR_SLICE)
? yv2->slice_parent : yv2;
if (!y_fp && !y_int) da_eligible = false;
else if (!y_fp && ysrc && ray_vec_may_have_nulls(ysrc))
da_eligible = false;
}
}
if (aop == OP_MEDIAN) da_eligible = false;
if (aop == OP_QUANTILE) da_eligible = false;
if (aop == OP_MODE) da_eligible = false;
if (aop == OP_TOP_N) da_eligible = false;
if (aop == OP_BOT_N) da_eligible = false;
if (agg_vecs[a] &&
agg_needs_row_gather(agg_vecs[a]->type, aop))
da_eligible = false;
}
for (uint32_t k = 0; k < n_keys && da_eligible; k++) {
if (!key_data[k]) { da_eligible = false; break; }
int8_t t = key_types[k];
if (t != RAY_I64 && t != RAY_SYM && t != RAY_I32
&& t != RAY_TIMESTAMP && t != RAY_DATE && t != RAY_TIME
&& t != RAY_BOOL && t != RAY_U8 && t != RAY_I16) {
da_eligible = false;
}
if (key_vecs[k]) {
ray_t* src = (key_vecs[k]->attrs & RAY_ATTR_SLICE)
? key_vecs[k]->slice_parent : key_vecs[k];
if (t != RAY_SYM && ray_vec_may_have_nulls(src))
da_eligible = false;
}
}
if (da_eligible && __VM && __VM->grp_card_n > 0) {
uint64_t known_ub = 1;
for (uint32_t k = 0; k < n_keys; k++) {
if (key_scan_sym[k] < 0) continue;
int64_t card = ray_grp_card_lookup(key_scan_sym[k]);
if (card <= 1) continue;
if ((uint64_t)card > da_slot_budget / known_ub) {
da_eligible = false;
break;
}
known_ub *= (uint64_t)card;
if (known_ub > da_slot_budget) { da_eligible = false; break; }
}
}
int64_t da_key_min[8], da_key_range[8], da_key_stride[8];
uint64_t total_slots = 1;
bool da_fits = false;
if (da_eligible) {
da_fits = true;
ray_pool_t* mm_pool = ray_pool_get();
uint32_t mm_n = (mm_pool && nrows >= RAY_PARALLEL_THRESHOLD)
? ray_pool_total_workers(mm_pool) : 1;
int64_t mm_mins[mm_n], mm_maxs[mm_n];
_Atomic(int) mm_abort = 0;
for (uint32_t k = 0; k < n_keys && da_fits; k++) {
int64_t kmin, kmax;
for (uint32_t w = 0; w < mm_n; w++) {
mm_mins[w] = INT64_MAX;
mm_maxs[w] = INT64_MIN;
}
minmax_ctx_t mm_ctx = {
.key_data = key_data[k],
.key_type = key_types[k],
.key_attrs = key_attrs[k],
.per_worker_min = mm_mins,
.per_worker_max = mm_maxs,
.n_workers = mm_n,
.match_idx = match_idx,
.rowsel = rowsel,
.span_budget = (int64_t)da_slot_budget,
.abort_flag = &mm_abort,
};
if (mm_n > 1) {
ray_pool_dispatch(mm_pool, minmax_scan_fn, &mm_ctx, n_scan);
} else {
minmax_scan_fn(&mm_ctx, 0, 0, n_scan);
}
if (atomic_load_explicit(&mm_abort, memory_order_relaxed)) {
da_fits = false;
break;
}
kmin = INT64_MAX; kmax = INT64_MIN;
for (uint32_t w = 0; w < mm_n; w++) {
if (mm_mins[w] < kmin) kmin = mm_mins[w];
if (mm_maxs[w] > kmax) kmax = mm_maxs[w];
}
if (key_types[k] == RAY_SYM) {
if (kmin == INT64_MAX) { kmin = 1; kmax = 0; }
kmin--;
}
da_key_min[k] = kmin;
uint64_t span = (uint64_t)kmax - (uint64_t)kmin + 1;
if (span > (uint64_t)INT64_MAX) { da_fits = false; break; }
da_key_range[k] = (int64_t)span;
if (da_key_range[k] <= 0) { da_fits = false; break; }
uint64_t range = (uint64_t)da_key_range[k];
uint64_t slot_limit = da_slot_budget < INT32_MAX
? da_slot_budget : INT32_MAX;
if (total_slots > slot_limit / range) {
da_fits = false;
break;
}
total_slots *= range;
}
}
if (da_fits) {
uint8_t need_flags = DA_NEED_COUNT;
bool all_sum = true;
bool da_has_first_last = false;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop == OP_SUM || aop == OP_PROD || aop == OP_AVG || aop == OP_ALL || aop == OP_ANY || aop == OP_FIRST || aop == OP_LAST) need_flags |= DA_NEED_SUM;
else if (aop == OP_STDDEV || aop == OP_STDDEV_POP || aop == OP_VAR || aop == OP_VAR_POP)
{ need_flags |= DA_NEED_SUM; need_flags |= DA_NEED_SUMSQ; }
else if (agg_is_binary_agg(aop))
{ need_flags |= DA_NEED_SUM; need_flags |= DA_NEED_SUMSQ; need_flags |= DA_NEED_PAIR; }
else if (aop == OP_MIN) need_flags |= DA_NEED_MIN;
else if (aop == OP_MAX) need_flags |= DA_NEED_MAX;
if (aop != OP_SUM && aop != OP_AVG && aop != OP_COUNT)
all_sum = false;
if (aop == OP_FIRST || aop == OP_LAST) da_has_first_last = true;
}
bool stride_overflow = false;
for (uint32_t k = 0; k < n_keys; k++) {
int64_t s = 1;
for (uint32_t j = k + 1; j < n_keys; j++) {
if (da_key_range[j] != 0 && s > INT64_MAX / da_key_range[j]) {
stride_overflow = true; break;
}
s *= da_key_range[j];
}
if (stride_overflow) break;
da_key_stride[k] = s;
}
if (stride_overflow) da_fits = false;
uint32_t n_slots = (uint32_t)total_slots;
size_t total = (size_t)n_slots * n_aggs;
void* agg_ptrs[vla_aggs];
int8_t agg_types[vla_aggs];
void* agg_ptrs2[vla_aggs];
int8_t agg_types2[vla_aggs];
int64_t da_int_null_sentinel[vla_aggs];
uint64_t agg_f64_mask = 0;
uint64_t da_int_null_mask = 0;
bool da_need_sum128 = false;
bool da_any_nullable = false;
for (uint32_t a = 0; a < n_aggs; a++) {
agg_ptrs2[a] = NULL; agg_types2[a] = 0;
if (agg_is_binary_agg(ext->agg_ops[a]) && agg_vecs2 && agg_vecs2[a]) {
agg_ptrs2[a] = ray_data(agg_vecs2[a]);
agg_types2[a] = agg_vecs2[a]->type;
ray_t* ysrc = (agg_vecs2[a]->attrs & RAY_ATTR_SLICE)
? agg_vecs2[a]->slice_parent : agg_vecs2[a];
if (group_fp_type(agg_vecs2[a]->type) &&
ysrc && ray_vec_may_have_nulls(ysrc))
da_any_nullable = true;
}
if (agg_prod[a].enabled) {
agg_ptrs[a] = NULL;
agg_types[a] = RAY_F64;
agg_f64_mask |= ((uint64_t)1 << a);
da_int_null_sentinel[a] = 0;
continue;
}
if (agg_vecs[a]) {
agg_ptrs[a] = ray_data(agg_vecs[a]);
agg_types[a] = agg_vecs[a]->type;
if (group_fp_type(agg_vecs[a]->type))
agg_f64_mask |= ((uint64_t)1 << a);
da_int_null_sentinel[a] = agg_int_null_sentinel_for(agg_vecs[a]->type);
int8_t t = agg_vecs[a]->type;
bool is_sentinel_typed = (t == RAY_I16 || t == RAY_I32 || t == RAY_I64 ||
t == RAY_DATE || t == RAY_TIME || t == RAY_TIMESTAMP || t == RAY_SYM);
if (is_sentinel_typed && ray_vec_may_have_nulls(agg_vecs[a]))
da_int_null_mask |= ((uint64_t)1 << a);
if (ray_vec_may_have_nulls(agg_vecs[a]) &&
(group_fp_type(agg_vecs[a]->type) || is_sentinel_typed))
da_any_nullable = true;
} else {
agg_ptrs[a] = NULL;
agg_types[a] = 0;
da_int_null_sentinel[a] = 0;
}
if (ext->agg_ops[a] == OP_AVG && !(agg_f64_mask & ((uint64_t)1 << a)) &&
!agg_strlen[a] && group_avg_needs_hi(agg_types[a], nrows))
da_need_sum128 = true;
}
ray_pool_t* da_pool = ray_pool_get();
uint32_t da_n_workers = (da_pool && nrows >= RAY_PARALLEL_THRESHOLD)
? ray_pool_total_workers(da_pool) : 1;
uint32_t arrays_per_agg = 0;
if (need_flags & DA_NEED_SUM) arrays_per_agg += 1;
if (da_need_sum128) arrays_per_agg += 1;
if (need_flags & DA_NEED_MIN) arrays_per_agg += 1;
if (need_flags & DA_NEED_MAX) arrays_per_agg += 1;
if (need_flags & DA_NEED_SUMSQ) arrays_per_agg += 1;
if (need_flags & DA_NEED_PAIR) arrays_per_agg += 3;
if (da_any_nullable) arrays_per_agg += 1;
uint64_t cells_per_worker = (uint64_t)n_slots
* ((uint64_t)arrays_per_agg * n_aggs + 1u);
uint64_t max_workers = cells_per_worker
? (uint64_t)n_scan / cells_per_worker : 1;
if (max_workers < 1) max_workers = 1;
bool da_ranged = false;
uint32_t da_part_tasks = 0;
if ((uint64_t)da_n_workers > max_workers && da_has_first_last) {
da_n_workers = 1;
} else if ((uint64_t)da_n_workers > max_workers) {
if (n_slots >= 2 * da_n_workers && nrows <= INT32_MAX) {
da_part_tasks = 2 * da_n_workers;
if (da_part_tasks > n_slots) da_part_tasks = n_slots;
da_n_workers = 1;
} else {
da_n_workers = (uint32_t)max_workers;
da_ranged = da_n_workers > 1;
}
}
ray_t* accums_hdr;
da_accum_t* accums = (da_accum_t*)scratch_calloc(&accums_hdr,
da_n_workers * sizeof(da_accum_t));
if (!accums) goto ht_path;
bool alloc_ok = true;
for (uint32_t w = 0; w < da_n_workers; w++) {
if (need_flags & DA_NEED_SUM) {
accums[w].sum = (da_val_t*)scratch_calloc(&accums[w]._h_sum,
total * sizeof(da_val_t));
if (!accums[w].sum) { alloc_ok = false; break; }
}
if (da_need_sum128) {
accums[w].sum_hi = (int64_t*)scratch_calloc(&accums[w]._h_sum_hi,
total * sizeof(int64_t));
if (!accums[w].sum_hi) { alloc_ok = false; break; }
}
if (need_flags & DA_NEED_SUMSQ) {
accums[w].sumsq_f64 = (double*)scratch_calloc(&accums[w]._h_sumsq,
total * sizeof(double));
if (!accums[w].sumsq_f64) { alloc_ok = false; break; }
}
if (need_flags & DA_NEED_PAIR) {
accums[w].sum_y = (double*)scratch_calloc(&accums[w]._h_sum_y,
total * sizeof(double));
accums[w].sumsq_y = (double*)scratch_calloc(&accums[w]._h_sumsq_y,
total * sizeof(double));
accums[w].sumxy = (double*)scratch_calloc(&accums[w]._h_sumxy,
total * sizeof(double));
if (!accums[w].sum_y || !accums[w].sumsq_y || !accums[w].sumxy) {
alloc_ok = false; break;
}
}
if (need_flags & DA_NEED_MIN) {
accums[w].min_val = (da_val_t*)scratch_alloc(&accums[w]._h_min,
total * sizeof(da_val_t));
if (!accums[w].min_val) { alloc_ok = false; break; }
for (size_t i = 0; i < total; i++) {
uint8_t a = (uint8_t)(i % n_aggs);
if (group_fp_type(agg_types[a])) accums[w].min_val[i].f = DBL_MAX;
else accums[w].min_val[i].i = INT64_MAX;
}
}
if (need_flags & DA_NEED_MAX) {
accums[w].max_val = (da_val_t*)scratch_alloc(&accums[w]._h_max,
total * sizeof(da_val_t));
if (!accums[w].max_val) { alloc_ok = false; break; }
for (size_t i = 0; i < total; i++) {
uint8_t a = (uint8_t)(i % n_aggs);
if (group_fp_type(agg_types[a])) accums[w].max_val[i].f = -DBL_MAX;
else accums[w].max_val[i].i = INT64_MIN;
}
}
accums[w].count = (int64_t*)scratch_calloc(&accums[w]._h_count,
n_slots * sizeof(int64_t));
if (!accums[w].count) { alloc_ok = false; break; }
if (da_any_nullable) {
accums[w].nn_count = (int64_t*)scratch_calloc(
&accums[w]._h_nn_count, total * sizeof(int64_t));
if (!accums[w].nn_count) { alloc_ok = false; break; }
}
if (da_has_first_last) {
accums[w].first_row = (int64_t*)scratch_alloc(
&accums[w]._h_first_row, n_slots * sizeof(int64_t));
if (!accums[w].first_row) { alloc_ok = false; break; }
for (uint32_t s = 0; s < n_slots; s++)
accums[w].first_row[s] = INT64_MAX;
accums[w].last_row = (int64_t*)scratch_alloc(
&accums[w]._h_last_row, n_slots * sizeof(int64_t));
if (!accums[w].last_row) { alloc_ok = false; break; }
for (uint32_t s = 0; s < n_slots; s++)
accums[w].last_row[s] = INT64_MIN;
}
}
if (!alloc_ok) {
for (uint32_t w = 0; w < da_n_workers; w++)
da_accum_free(&accums[w]);
scratch_free(accums_hdr);
goto ht_path;
}
uint8_t da_key_esz[n_keys];
for (uint32_t k = 0; k < n_keys; k++)
da_key_esz[k] = ray_sym_elem_size(key_types[k], key_attrs[k]);
ray_t** da_sym_strings = NULL;
uint32_t da_sym_count = 0;
for (uint32_t a = 0; a < n_aggs; a++) {
if (agg_strlen[a] && agg_vecs[a] &&
agg_vecs[a]->type == RAY_SYM) {
ray_sym_strings_borrow(&da_sym_strings, &da_sym_count);
break;
}
}
da_ctx_t da_ctx = {
.accums = accums,
.sym_strings = da_sym_strings,
.sym_count = da_sym_count,
.n_accums = da_n_workers,
.key_ptrs = key_data,
.key_types = key_types,
.key_attrs = key_attrs,
.key_esz = da_key_esz,
.key_mins = da_key_min,
.key_strides = da_key_stride,
.n_keys = n_keys,
.agg_ptrs = agg_ptrs,
.agg_ptrs2 = agg_ptrs2,
.agg_types2 = agg_types2,
.agg_types = agg_types,
.agg_cols = agg_vecs,
.agg_strlen = agg_strlen,
.agg_prod = agg_prod,
.agg_ops = ext->agg_ops,
.n_aggs = n_aggs,
.need_flags = need_flags,
.agg_f64_mask = agg_f64_mask,
.agg_int_null_mask = da_int_null_mask,
.agg_int_null_sentinel = da_int_null_sentinel,
.all_sum = all_sum,
.n_slots = n_slots,
.match_idx = match_idx,
.rowsel = rowsel,
.da_n_scan = n_scan,
};
ray_t* agg_raw_hdr = NULL;
if (n_aggs > 0) {
char* rm = (char*)scratch_calloc(&agg_raw_hdr,
(size_t)n_aggs * (sizeof(ray_sym_domain_raw_t) + 1));
if (rm) {
da_ctx.agg_raw = (ray_sym_domain_raw_t*)rm;
da_ctx.agg_raw_ok = (uint8_t*)(rm + (size_t)n_aggs * sizeof(ray_sym_domain_raw_t));
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_vecs[a] && agg_vecs[a]->type == RAY_SYM)
da_ctx.agg_raw_ok[a] = ray_sym_domain_raw_pin(
ray_sym_vec_domain(agg_vecs[a]), &da_ctx.agg_raw[a]) ? 1 : 0;
}
}
if (da_part_tasks > 0) {
if (!da_accum_partitioned(&da_ctx, da_pool, da_part_tasks, n_scan))
da_accum_fn(&da_ctx, 0, 0, n_scan);
} else if (da_ranged)
ray_pool_dispatch_n(da_pool, da_accum_task_fn, &da_ctx, da_n_workers);
else if (da_n_workers > 1)
ray_pool_dispatch(da_pool, da_accum_fn, &da_ctx, n_scan);
else
da_accum_fn(&da_ctx, 0, 0, n_scan);
if (agg_raw_hdr) scratch_free(agg_raw_hdr);
da_accum_t* merged = &accums[0];
bool has_first_last = false;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop == OP_FIRST || aop == OP_LAST) { has_first_last = true; break; }
}
if (has_first_last) {
for (uint32_t w = 1; w < da_n_workers; w++) {
da_accum_t* wa = &accums[w];
if (need_flags & DA_NEED_SUMSQ) {
for (size_t i = 0; i < total; i++)
merged->sumsq_f64[i] += wa->sumsq_f64[i];
}
if (need_flags & DA_NEED_PAIR) {
for (size_t i = 0; i < total; i++) {
merged->sum_y[i] += wa->sum_y[i];
merged->sumsq_y[i] += wa->sumsq_y[i];
merged->sumxy[i] += wa->sumxy[i];
}
}
if (need_flags & DA_NEED_SUM) {
for (uint32_t s = 0; s < n_slots; s++) {
size_t base = (size_t)s * n_aggs;
bool take_first = wa->first_row && merged->first_row &&
wa->first_row[s] < merged->first_row[s];
bool take_last = wa->last_row && merged->last_row &&
wa->last_row[s] > merged->last_row[s];
for (uint32_t a = 0; a < n_aggs; a++) {
size_t idx = base + a;
uint16_t aop = ext->agg_ops[a];
if (aop == OP_SUM || aop == OP_AVG || aop == OP_ALL || aop == OP_ANY || aop == OP_STDDEV || aop == OP_STDDEV_POP || aop == OP_VAR || aop == OP_VAR_POP || agg_is_binary_agg(aop)) {
if (group_fp_type(agg_types[a]) || agg_is_binary_agg(aop)) merged->sum[idx].f += wa->sum[idx].f;
else if (merged->sum_hi) ray_i128_add128(&merged->sum_hi[idx], (uint64_t*)&merged->sum[idx].i, wa->sum_hi[idx], (uint64_t)wa->sum[idx].i);
else merged->sum[idx].i = wrap_add_i64(merged->sum[idx].i, wa->sum[idx].i);
} else if (aop == OP_PROD) {
int64_t mnn = merged->nn_count ? merged->nn_count[idx] : merged->count[s];
int64_t wnn = wa->nn_count ? wa->nn_count[idx] : wa->count[s];
if (wnn > 0) {
if (mnn == 0)
merged->sum[idx] = wa->sum[idx];
else if (group_fp_type(agg_types[a]))
merged->sum[idx].f *= wa->sum[idx].f;
else
merged->sum[idx].i = (int64_t)((uint64_t)merged->sum[idx].i * (uint64_t)wa->sum[idx].i);
}
} else if (aop == OP_FIRST) {
if (take_first) merged->sum[idx] = wa->sum[idx];
} else if (aop == OP_LAST) {
if (take_last) merged->sum[idx] = wa->sum[idx];
}
}
if (take_first) merged->first_row[s] = wa->first_row[s];
if (take_last) merged->last_row[s] = wa->last_row[s];
}
}
if (need_flags & DA_NEED_MIN) {
for (size_t i = 0; i < total; i++) {
uint8_t a = (uint8_t)(i % n_aggs);
if (group_fp_type(agg_types[a])) {
if (wa->min_val[i].f < merged->min_val[i].f)
merged->min_val[i].f = wa->min_val[i].f;
} else if (agg_types[a] == RAY_SYM) {
if (wa->min_val[i].i != INT64_MAX &&
(merged->min_val[i].i == INT64_MAX ||
sym_lex_lt(ray_sym_vec_domain(agg_vecs[a]),
wa->min_val[i].i, merged->min_val[i].i)))
merged->min_val[i].i = wa->min_val[i].i;
} else {
if (wa->min_val[i].i < merged->min_val[i].i)
merged->min_val[i].i = wa->min_val[i].i;
}
}
}
if (need_flags & DA_NEED_MAX) {
for (size_t i = 0; i < total; i++) {
uint8_t a = (uint8_t)(i % n_aggs);
if (group_fp_type(agg_types[a])) {
if (wa->max_val[i].f > merged->max_val[i].f)
merged->max_val[i].f = wa->max_val[i].f;
} else if (agg_types[a] == RAY_SYM) {
if (wa->max_val[i].i != INT64_MIN &&
(merged->max_val[i].i == INT64_MIN ||
sym_lex_gt(ray_sym_vec_domain(agg_vecs[a]),
wa->max_val[i].i, merged->max_val[i].i)))
merged->max_val[i].i = wa->max_val[i].i;
} else {
if (wa->max_val[i].i > merged->max_val[i].i)
merged->max_val[i].i = wa->max_val[i].i;
}
}
}
if (merged->nn_count && wa->nn_count) {
for (size_t i = 0; i < total; i++)
merged->nn_count[i] += wa->nn_count[i];
}
for (uint32_t s = 0; s < n_slots; s++)
merged->count[s] += wa->count[s];
}
} else if (da_n_workers > 1 && n_slots >= 1024 && da_pool) {
da_merge_ctx_t merge_ctx = {
.accums = accums,
.n_src_workers = da_n_workers,
.need_flags = need_flags,
.n_aggs = n_aggs,
.agg_types = agg_types,
.agg_ops = ext->agg_ops,
.agg_cols = agg_vecs,
};
ray_pool_dispatch(da_pool, da_merge_fn, &merge_ctx, (int64_t)n_slots);
} else {
for (uint32_t w = 1; w < da_n_workers; w++) {
da_accum_t* wa = &accums[w];
if (need_flags & DA_NEED_SUMSQ) {
for (size_t i = 0; i < total; i++)
merged->sumsq_f64[i] += wa->sumsq_f64[i];
}
if (need_flags & DA_NEED_PAIR) {
for (size_t i = 0; i < total; i++) {
merged->sum_y[i] += wa->sum_y[i];
merged->sumsq_y[i] += wa->sumsq_y[i];
merged->sumxy[i] += wa->sumxy[i];
}
}
if (need_flags & DA_NEED_SUM) {
for (uint32_t s = 0; s < n_slots; s++) {
size_t base = (size_t)s * n_aggs;
for (uint32_t a = 0; a < n_aggs; a++) {
size_t idx = base + a;
uint16_t aop = ext->agg_ops[a];
int64_t mnn = merged->nn_count ? merged->nn_count[idx] : merged->count[s];
int64_t wnn = wa->nn_count ? wa->nn_count[idx] : wa->count[s];
if (aop == OP_FIRST) {
if (mnn == 0 && wnn > 0)
merged->sum[idx] = wa->sum[idx];
} else if (aop == OP_LAST) {
if (wnn > 0)
merged->sum[idx] = wa->sum[idx];
} else if (aop == OP_PROD) {
if (wnn > 0) {
if (mnn == 0)
merged->sum[idx] = wa->sum[idx];
else if (group_fp_type(agg_types[a]))
merged->sum[idx].f *= wa->sum[idx].f;
else
merged->sum[idx].i = (int64_t)((uint64_t)merged->sum[idx].i * (uint64_t)wa->sum[idx].i);
}
} else if (group_fp_type(agg_types[a]) || agg_is_binary_agg(aop))
merged->sum[idx].f += wa->sum[idx].f;
else if (merged->sum_hi)
ray_i128_add128(&merged->sum_hi[idx], (uint64_t*)&merged->sum[idx].i,
wa->sum_hi[idx], (uint64_t)wa->sum[idx].i);
else
merged->sum[idx].i = wrap_add_i64(merged->sum[idx].i, wa->sum[idx].i);
}
}
}
if (need_flags & DA_NEED_MIN) {
for (size_t i = 0; i < total; i++) {
uint8_t a = (uint8_t)(i % n_aggs);
if (group_fp_type(agg_types[a])) {
if (wa->min_val[i].f < merged->min_val[i].f)
merged->min_val[i].f = wa->min_val[i].f;
} else if (agg_types[a] == RAY_SYM) {
if (wa->min_val[i].i != INT64_MAX &&
(merged->min_val[i].i == INT64_MAX ||
sym_lex_lt(ray_sym_vec_domain(agg_vecs[a]),
wa->min_val[i].i, merged->min_val[i].i)))
merged->min_val[i].i = wa->min_val[i].i;
} else {
if (wa->min_val[i].i < merged->min_val[i].i)
merged->min_val[i].i = wa->min_val[i].i;
}
}
}
if (need_flags & DA_NEED_MAX) {
for (size_t i = 0; i < total; i++) {
uint8_t a = (uint8_t)(i % n_aggs);
if (group_fp_type(agg_types[a])) {
if (wa->max_val[i].f > merged->max_val[i].f)
merged->max_val[i].f = wa->max_val[i].f;
} else if (agg_types[a] == RAY_SYM) {
if (wa->max_val[i].i != INT64_MIN &&
(merged->max_val[i].i == INT64_MIN ||
sym_lex_gt(ray_sym_vec_domain(agg_vecs[a]),
wa->max_val[i].i, merged->max_val[i].i)))
merged->max_val[i].i = wa->max_val[i].i;
} else {
if (wa->max_val[i].i > merged->max_val[i].i)
merged->max_val[i].i = wa->max_val[i].i;
}
}
}
if (merged->nn_count && wa->nn_count) {
for (size_t i = 0; i < total; i++)
merged->nn_count[i] += wa->nn_count[i];
}
for (uint32_t s = 0; s < n_slots; s++)
merged->count[s] += wa->count[s];
}
}
for (uint32_t w = 1; w < da_n_workers; w++)
da_accum_free(&accums[w]);
da_val_t* da_sum = merged->sum;
int64_t* da_sum_hi = merged->sum_hi;
da_val_t* da_min_val = merged->min_val;
da_val_t* da_max_val = merged->max_val;
double* da_sumsq = merged->sumsq_f64;
double* da_sum_y = merged->sum_y;
double* da_sumsq_y = merged->sumsq_y;
double* da_sumxy = merged->sumxy;
int64_t* da_count = merged->count;
int64_t* da_nn_count = merged->nn_count;
uint32_t all_grp_count = 0;
for (uint32_t s = 0; s < n_slots; s++)
if (da_count[s] > 0) all_grp_count++;
int64_t da_keep_min = use_emit_filter
? da_count_emit_keep_min(da_count, n_slots, all_grp_count, emit_filter)
: 1;
uint32_t grp_count = 0;
for (uint32_t s = 0; s < n_slots; s++)
if (da_count[s] >= da_keep_min) grp_count++;
int64_t total_cols = n_keys + n_aggs;
ray_t* result = ray_table_new(total_cols);
if (!result || RAY_IS_ERR(result)) {
da_accum_free(&accums[0]); scratch_free(accums_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result ? result : ray_error("oom", NULL);
}
for (uint32_t k = 0; k < n_keys; k++) {
ray_t* src_col = key_vecs[k];
if (!src_col) continue;
ray_t* key_col = col_vec_new(src_col, (int64_t)grp_count);
out_col_adopt_str_pool(key_col, src_col);
if (!key_col || RAY_IS_ERR(key_col)) continue;
if (key_col->type == RAY_SYM)
ray_sym_vec_adopt_domain(key_col, sym_domain_rep(src_col));
key_col->len = (int64_t)grp_count;
uint32_t gi = 0;
for (uint32_t s = 0; s < n_slots; s++) {
if (da_count[s] < da_keep_min) continue;
int64_t offset = ((int64_t)s / da_key_stride[k]) % da_key_range[k];
int64_t key_val = (src_col->type == RAY_SYM && offset == 0)
? 0 : da_key_min[k] + offset;
write_col_i64(ray_data(key_col), gi, key_val, src_col->type, key_col->attrs);
gi++;
}
ray_op_ext_t* key_ext = find_ext(g, ext->keys[k]);
int64_t name_id = key_ext ? key_ext->sym : (int64_t)k;
result = ray_table_add_col(result, name_id, key_col);
ray_release(key_col);
}
size_t dense_total = (size_t)grp_count * n_aggs;
ray_t *_h_dsum = NULL, *_h_dmin = NULL, *_h_dmax = NULL;
ray_t *_h_dsq = NULL, *_h_dcnt = NULL, *_h_dnn = NULL;
ray_t *_h_dsy = NULL, *_h_dsqy = NULL, *_h_dxy = NULL, *_h_dshi = NULL;
da_val_t* dense_sum = da_sum ? (da_val_t*)scratch_alloc(&_h_dsum, dense_total * sizeof(da_val_t)) : NULL;
int64_t* dense_sum_hi = da_sum_hi ? (int64_t*)scratch_alloc(&_h_dshi, dense_total * sizeof(int64_t)) : NULL;
da_val_t* dense_min_val = da_min_val ? (da_val_t*)scratch_alloc(&_h_dmin, dense_total * sizeof(da_val_t)) : NULL;
da_val_t* dense_max_val = da_max_val ? (da_val_t*)scratch_alloc(&_h_dmax, dense_total * sizeof(da_val_t)) : NULL;
double* dense_sumsq = da_sumsq ? (double*)scratch_alloc(&_h_dsq, dense_total * sizeof(double)) : NULL;
double* dense_sum_y = da_sum_y ? (double*)scratch_alloc(&_h_dsy, dense_total * sizeof(double)) : NULL;
double* dense_sumsq_y = da_sumsq_y ? (double*)scratch_alloc(&_h_dsqy, dense_total * sizeof(double)) : NULL;
double* dense_sumxy = da_sumxy ? (double*)scratch_alloc(&_h_dxy, dense_total * sizeof(double)) : NULL;
int64_t* dense_counts = grp_count
? (int64_t*)scratch_alloc(&_h_dcnt, grp_count * sizeof(int64_t))
: NULL;
int64_t* dense_nn_counts = (da_nn_count && grp_count)
? (int64_t*)scratch_alloc(&_h_dnn, dense_total * sizeof(int64_t))
: NULL;
uint32_t gi = 0;
for (uint32_t s = 0; s < n_slots; s++) {
if (da_count[s] < da_keep_min) continue;
dense_counts[gi] = da_count[s];
for (uint32_t a = 0; a < n_aggs; a++) {
size_t si = (size_t)s * n_aggs + a;
size_t di = (size_t)gi * n_aggs + a;
if (dense_sum) dense_sum[di] = da_sum[si];
if (dense_sum_hi) dense_sum_hi[di] = da_sum_hi[si];
if (dense_min_val) dense_min_val[di] = da_min_val[si];
if (dense_max_val) dense_max_val[di] = da_max_val[si];
if (dense_sumsq) dense_sumsq[di] = da_sumsq[si];
if (dense_sum_y) dense_sum_y[di] = da_sum_y[si];
if (dense_sumsq_y) dense_sumsq_y[di] = da_sumsq_y[si];
if (dense_sumxy) dense_sumxy[di] = da_sumxy[si];
if (dense_nn_counts) dense_nn_counts[di] = da_nn_count[si];
}
gi++;
}
emit_agg_columns(&result, g, ext, agg_vecs, grp_count, n_aggs,
(double*)dense_sum, (int64_t*)dense_sum, dense_sum_hi,
(double*)dense_min_val, (double*)dense_max_val,
(int64_t*)dense_min_val, (int64_t*)dense_max_val,
dense_counts, agg_affine, agg_prod, dense_sumsq,
dense_nn_counts, dense_sum_y, dense_sumsq_y, dense_sumxy);
scratch_free(_h_dsum); scratch_free(_h_dmin);
scratch_free(_h_dmax);
scratch_free(_h_dsq); scratch_free(_h_dcnt);
scratch_free(_h_dnn);
scratch_free(_h_dsy); scratch_free(_h_dsqy); scratch_free(_h_dxy);
scratch_free(_h_dshi);
da_accum_free(&accums[0]); scratch_free(accums_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result;
}
}
{
bool sp_eligible = (nrows > 0 && n_keys == 1 && key_data[0] != NULL &&
n_aggs <= 64);
int8_t kt = sp_eligible ? key_types[0] : 0;
if (sp_eligible && kt != RAY_I64 && kt != RAY_I32 && kt != RAY_I16 &&
kt != RAY_U8 && kt != RAY_BOOL && kt != RAY_DATE &&
kt != RAY_TIME && kt != RAY_TIMESTAMP && kt != RAY_SYM)
sp_eligible = false;
if (sp_eligible && key_vecs[0]) {
ray_t* src = (key_vecs[0]->attrs & RAY_ATTR_SLICE)
? key_vecs[0]->slice_parent : key_vecs[0];
if (kt != RAY_SYM && ray_vec_may_have_nulls(src))
sp_eligible = false;
}
bool sp_need_sum = false;
bool sp_need_sum128 = false;
for (uint32_t a = 0; a < n_aggs && sp_eligible; a++) {
uint16_t op = ext->agg_ops[a];
if (op == OP_COUNT) continue;
if (op != OP_SUM && op != OP_AVG)
sp_eligible = false;
else {
if (agg_vecs[a] && ray_vec_may_have_nulls(agg_vecs[a]))
sp_eligible = false;
else {
sp_need_sum = true;
if (op == OP_AVG && !(agg_vecs[a] && group_fp_type(agg_vecs[a]->type)) &&
!agg_strlen[a] &&
group_avg_needs_hi(agg_vecs[a] ? agg_vecs[a]->type : 0, nrows))
sp_need_sum128 = true;
}
}
}
if (sp_eligible) {
if (!group_materialize_prod_slots(agg_prod, agg_vecs, agg_owned,
n_aggs, nrows))
goto ht_path;
void* agg_ptrs[vla_aggs];
int8_t agg_types[vla_aggs];
uint64_t agg_f64_mask = 0;
for (uint32_t a = 0; a < n_aggs; a++) {
if (agg_vecs[a]) {
agg_ptrs[a] = ray_data(agg_vecs[a]);
agg_types[a] = agg_vecs[a]->type;
if (group_fp_type(agg_vecs[a]->type))
agg_f64_mask |= ((uint64_t)1 << a);
} else {
agg_ptrs[a] = NULL;
agg_types[a] = 0;
}
}
ray_t** strlen_sym_strings = NULL;
uint32_t strlen_sym_count = 0;
for (uint32_t a = 0; a < n_aggs; a++) {
if (agg_strlen[a] && agg_vecs[a] &&
agg_vecs[a]->type == RAY_SYM) {
ray_sym_strings_borrow(&strlen_sym_strings,
&strlen_sym_count);
break;
}
}
uint8_t key_esz = ray_sym_elem_size(key_types[0], key_attrs[0]);
bool sp_dyn_sparse = false;
if (rowsel) {
ray_rowsel_t* sp_sm = ray_rowsel_meta(rowsel);
uint64_t sp_cap = key_esz == 1 ? 256u
: key_esz == 2 ? (1u << 16)
: (1u << 20);
if (sp_sm && (uint64_t)sp_sm->total_pass * 32 < sp_cap)
sp_dyn_sparse = true;
}
if (use_emit_filter && !sp_dyn_sparse &&
(emit_filter.min_count_exclusive > 0 ||
emit_filter.top_count_take > 0) &&
n_scan <= UINT32_MAX) {
sp_dyn_ctx_t sc = {
.g = g, .ext = ext, .key_data = key_data,
.key_types = key_types, .key_vecs = key_vecs,
.key_owned = key_owned, .agg_ptrs = agg_ptrs,
.agg_types = agg_types, .agg_strlen = agg_strlen,
.agg_vecs = agg_vecs, .agg_owned = agg_owned,
.agg_affine = agg_affine, .agg_prod = agg_prod,
.strlen_sym_strings = strlen_sym_strings,
.strlen_sym_count = strlen_sym_count,
.agg_f64_mask = agg_f64_mask, .n_aggs = n_aggs,
.n_keys = n_keys, .n_scan = n_scan, .key_esz = key_esz,
.sp_need_sum = sp_need_sum, .sp_need_sum128 = sp_need_sum128,
.match_idx = match_idx,
.rowsel = rowsel, .match_idx_block = match_idx_block,
.vla_hdr = vla_hdr, .emit_filter = emit_filter,
};
ray_t* sp_r = exec_group_sp_dyn_emit(&sc);
if (sp_r) return sp_r;
}
if (use_emit_filter &&
(emit_filter.min_count_exclusive > 0 ||
emit_filter.top_count_take > 0) &&
key_types[0] != RAY_SYM && n_scan <= UINT32_MAX) {
bool have_key = false;
int64_t min_key = 0, max_key = 0;
for (int64_t i = 0; i < n_scan; i++) {
int64_t r = match_idx ? match_idx[i] : i;
if (!match_idx && rowsel && !group_rowsel_pass(rowsel, r))
continue;
int64_t key = read_by_esz(key_data[0], r, key_esz);
if (!have_key) {
min_key = max_key = key;
have_key = true;
} else {
if (key < min_key) min_key = key;
if (key > max_key) max_key = key;
}
}
uint64_t key_range = have_key
? (uint64_t)((uint64_t)max_key - (uint64_t)min_key + 1u)
: 0u;
if (have_key && key_range > 0 && key_range <= (1u << 26)) {
ray_t *cnt_hdr = NULL, *range_sum_hdr = NULL, *range_hi_hdr = NULL;
ray_t *_h_sum = NULL, *_h_cnt = NULL, *_h_sum_hi = NULL;
uint32_t* range_count = (uint32_t*)scratch_calloc(
&cnt_hdr, (size_t)key_range * sizeof(uint32_t));
if (!range_count)
goto ht_path;
da_val_t* range_sum = NULL;
int64_t* range_sum_hi = NULL;
if (sp_need_sum && key_range <= (1u << 24)) {
range_sum = (da_val_t*)scratch_calloc(
&range_sum_hdr,
(size_t)key_range * n_aggs * sizeof(da_val_t));
if (!range_sum) {
scratch_free(cnt_hdr);
goto ht_path;
}
if (sp_need_sum128) {
range_sum_hi = (int64_t*)scratch_calloc(
&range_hi_hdr,
(size_t)key_range * n_aggs * sizeof(int64_t));
if (!range_sum_hi) {
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr);
scratch_free(cnt_hdr);
goto ht_path;
}
}
}
for (int64_t i = 0; i < n_scan; i++) {
int64_t r = match_idx ? match_idx[i] : i;
if (!match_idx && rowsel && !group_rowsel_pass(rowsel, r))
continue;
int64_t key = read_by_esz(key_data[0], r, key_esz);
uint64_t off = (uint64_t)((uint64_t)key - (uint64_t)min_key);
if (range_count[off] != UINT32_MAX)
range_count[off]++;
if (range_sum) {
da_val_t* sums = &range_sum[(size_t)off * n_aggs];
for (uint32_t a = 0; a < n_aggs; a++) {
if (ext->agg_ops[a] == OP_COUNT || !agg_ptrs[a])
continue;
if (agg_strlen[a])
sums[a].i = wrap_add_i64(
sums[a].i,
group_strlen_at_cached(agg_vecs[a], r,
strlen_sym_strings,
strlen_sym_count));
else if (agg_f64_mask & ((uint64_t)1 << a))
sums[a].f += group_fp_at(agg_ptrs[a], agg_types[a], r);
else if (range_sum_hi)
ray_i128_add(&range_sum_hi[(size_t)off * n_aggs + a],
(uint64_t*)&sums[a].i,
read_col_i64(agg_ptrs[a], r, agg_types[a], 0));
else
sums[a].i = wrap_add_i64(
sums[a].i,
read_col_i64(agg_ptrs[a], r, agg_types[a], 0));
}
}
}
uint32_t total_groups = 0;
for (uint64_t off = 0; off < key_range; off++) {
if (range_count[off] > 0)
total_groups++;
}
int64_t keep_min = da_count_emit_keep_min_u32(
range_count, key_range, total_groups, emit_filter);
uint32_t grp_count = 0;
for (uint64_t off = 0; off < key_range; off++) {
if ((int64_t)range_count[off] >= keep_min)
grp_count++;
}
ray_t* result = ray_table_new((int64_t)n_keys + n_aggs);
if (!result || RAY_IS_ERR(result)) {
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr);
scratch_free(cnt_hdr);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result ? result : ray_error("oom", NULL);
}
ray_t* key_col = col_vec_new(key_vecs[0], (int64_t)grp_count);
out_col_adopt_str_pool(key_col, key_vecs[0]);
if (key_col && !RAY_IS_ERR(key_col) && key_col->type == RAY_SYM)
ray_sym_vec_adopt_domain(key_col, sym_domain_rep(key_vecs[0]));
if (!key_col || RAY_IS_ERR(key_col)) {
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr);
scratch_free(cnt_hdr);
ray_release(result);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return key_col ? key_col : ray_error("oom", NULL);
}
key_col->len = (int64_t)grp_count;
da_val_t* dense_sum = sp_need_sum
? (da_val_t*)scratch_calloc(&_h_sum,
(size_t)grp_count * n_aggs * sizeof(da_val_t))
: NULL;
int64_t* dense_sum_hi = (sp_need_sum && sp_need_sum128)
? (int64_t*)scratch_calloc(&_h_sum_hi,
(size_t)grp_count * n_aggs * sizeof(int64_t))
: NULL;
int64_t* dense_count = (int64_t*)scratch_alloc(
&_h_cnt, (size_t)grp_count * sizeof(int64_t));
if ((sp_need_sum && !dense_sum) || !dense_count ||
(sp_need_sum && sp_need_sum128 && !dense_sum_hi)) {
scratch_free(_h_sum); scratch_free(_h_sum_hi); scratch_free(_h_cnt);
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr);
scratch_free(cnt_hdr);
ray_release(key_col); ray_release(result);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return ray_error("oom", NULL);
}
uint32_t gi = 0;
for (uint64_t off = 0; off < key_range; off++) {
uint32_t cnt = range_count[off];
if ((int64_t)cnt < keep_min) {
range_count[off] = 0;
continue;
}
int64_t key = (int64_t)((uint64_t)min_key + off);
write_col_i64(ray_data(key_col), gi, key,
key_col->type, key_col->attrs);
dense_count[gi] = (int64_t)cnt;
if (range_sum) {
memcpy(&dense_sum[(size_t)gi * n_aggs],
&range_sum[(size_t)off * n_aggs],
(size_t)n_aggs * sizeof(da_val_t));
if (dense_sum_hi)
memcpy(&dense_sum_hi[(size_t)gi * n_aggs],
&range_sum_hi[(size_t)off * n_aggs],
(size_t)n_aggs * sizeof(int64_t));
}
range_count[off] = gi + 1u;
gi++;
}
if (sp_need_sum && !range_sum) {
for (int64_t i = 0; i < n_scan; i++) {
int64_t r = match_idx ? match_idx[i] : i;
if (!match_idx && rowsel && !group_rowsel_pass(rowsel, r))
continue;
int64_t key = read_by_esz(key_data[0], r, key_esz);
uint64_t off = (uint64_t)((uint64_t)key - (uint64_t)min_key);
uint32_t marker = range_count[off];
if (!marker) continue;
da_val_t* sums = &dense_sum[(size_t)(marker - 1u) * n_aggs];
for (uint32_t a = 0; a < n_aggs; a++) {
if (ext->agg_ops[a] == OP_COUNT || !agg_ptrs[a])
continue;
if (agg_strlen[a])
sums[a].i = wrap_add_i64(
sums[a].i,
group_strlen_at_cached(agg_vecs[a], r,
strlen_sym_strings,
strlen_sym_count));
else if (agg_f64_mask & ((uint64_t)1 << a))
sums[a].f += group_fp_at(agg_ptrs[a], agg_types[a], r);
else if (dense_sum_hi)
ray_i128_add(&dense_sum_hi[(size_t)(marker - 1u) * n_aggs + a],
(uint64_t*)&sums[a].i,
read_col_i64(agg_ptrs[a], r, agg_types[a], 0));
else
sums[a].i = wrap_add_i64(
sums[a].i,
read_col_i64(agg_ptrs[a], r, agg_types[a], 0));
}
}
}
scratch_free(range_sum_hdr); scratch_free(range_hi_hdr);
scratch_free(cnt_hdr);
ray_op_ext_t* key_ext = find_ext(g, ext->keys[0]);
int64_t name_id = key_ext ? key_ext->sym : 0;
result = ray_table_add_col(result, name_id, key_col);
ray_release(key_col);
emit_agg_columns(&result, g, ext, agg_vecs, grp_count, n_aggs,
(double*)dense_sum, (int64_t*)dense_sum, dense_sum_hi,
NULL, NULL, NULL, NULL,
dense_count, agg_affine, agg_prod, NULL, NULL,
NULL, NULL, NULL);
scratch_free(_h_sum); scratch_free(_h_sum_hi);
scratch_free(_h_cnt);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result;
}
}
sparse_i64_ht_t sp_ht;
memset(&sp_ht, 0, sizeof(sp_ht));
bool sp_ok = true;
if (use_emit_filter &&
(emit_filter.min_count_exclusive > 0 ||
emit_filter.top_count_take > 0)) {
if (n_scan > (1 << 21)) goto ht_path;
uint64_t expected = (uint64_t)nrows / 64u;
if (expected < 4096) expected = 4096;
if (expected > (1u << 20)) expected = (1u << 20);
if (!sparse_i64_init(&sp_ht, (uint32_t)expected, n_aggs, false, false))
goto ht_path;
for (int64_t i = 0; i < n_scan; i++) {
int64_t r = match_idx ? match_idx[i] : i;
if (!match_idx && rowsel && !group_rowsel_pass(rowsel, r))
continue;
int64_t key = read_by_esz(key_data[0], r, key_esz);
int32_t slot;
if (!sparse_i64_touch(&sp_ht, key, n_aggs, false, &slot)) {
sp_ok = false;
break;
}
sp_ht.counts[slot]++;
}
} else {
uint64_t expected = (uint64_t)nrows / 64u;
if (expected < 4096) expected = 4096;
if (expected > (1u << 20)) expected = (1u << 20);
if (!sparse_i64_init(&sp_ht, (uint32_t)expected, n_aggs, sp_need_sum,
sp_need_sum128))
goto ht_path;
for (int64_t i = 0; i < n_scan; i++) {
int64_t r = match_idx ? match_idx[i] : i;
if (!match_idx && rowsel && !group_rowsel_pass(rowsel, r))
continue;
int64_t key = read_by_esz(key_data[0], r, key_esz);
int32_t slot;
if (!sparse_i64_touch(&sp_ht, key, n_aggs, sp_need_sum, &slot)) {
sp_ok = false;
break;
}
sp_ht.counts[slot]++;
if (!sp_need_sum) continue;
da_val_t* sums = &sp_ht.sums[(size_t)slot * n_aggs];
for (uint32_t a = 0; a < n_aggs; a++) {
if (ext->agg_ops[a] == OP_COUNT || !agg_ptrs[a])
continue;
if (agg_strlen[a])
sums[a].i = wrap_add_i64(
sums[a].i,
group_strlen_at_cached(agg_vecs[a], r,
strlen_sym_strings,
strlen_sym_count));
else if (agg_f64_mask & ((uint64_t)1 << a))
sums[a].f += group_fp_at(agg_ptrs[a], agg_types[a], r);
else if (sp_ht.sums_hi)
ray_i128_add(&sp_ht.sums_hi[(size_t)slot * n_aggs + a],
(uint64_t*)&sums[a].i,
read_col_i64(agg_ptrs[a], r, agg_types[a], 0));
else
sums[a].i = wrap_add_i64(
sums[a].i,
read_col_i64(agg_ptrs[a], r, agg_types[a], 0));
}
}
}
if (!sp_ok) {
sparse_i64_free(&sp_ht);
goto ht_path;
}
uint32_t total_groups = 0;
for (uint32_t s = 0; s < sp_ht.cap; s++) {
if (!sp_ht.used[s]) continue;
total_groups++;
}
int64_t keep_min = use_emit_filter
? da_count_emit_keep_min(sp_ht.counts, sp_ht.cap,
total_groups, emit_filter)
: 1;
uint32_t grp_count = 0;
for (uint32_t s = 0; s < sp_ht.cap; s++) {
if (!sp_ht.used[s]) continue;
if (sp_ht.counts[s] < keep_min) continue;
grp_count++;
}
ray_t* result = ray_table_new((int64_t)n_keys + n_aggs);
if (!result || RAY_IS_ERR(result)) {
sparse_i64_free(&sp_ht);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result ? result : ray_error("oom", NULL);
}
ray_t* key_col = col_vec_new(key_vecs[0], (int64_t)grp_count);
out_col_adopt_str_pool(key_col, key_vecs[0]);
if (key_col && !RAY_IS_ERR(key_col) && key_col->type == RAY_SYM)
ray_sym_vec_adopt_domain(key_col, sym_domain_rep(key_vecs[0]));
if (!key_col || RAY_IS_ERR(key_col)) {
sparse_i64_free(&sp_ht);
ray_release(result);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return key_col ? key_col : ray_error("oom", NULL);
}
key_col->len = (int64_t)grp_count;
ray_t *_h_sum = NULL, *_h_cnt = NULL, *_h_sum_hi = NULL;
da_val_t* dense_sum = sp_need_sum
? (da_val_t*)scratch_alloc(&_h_sum,
(size_t)grp_count * n_aggs * sizeof(da_val_t))
: NULL;
int64_t* dense_sum_hi = (sp_need_sum && sp_need_sum128)
? (int64_t*)scratch_alloc(&_h_sum_hi,
(size_t)grp_count * n_aggs * sizeof(int64_t))
: NULL;
int64_t* dense_count = (int64_t*)scratch_alloc(&_h_cnt,
(size_t)grp_count * sizeof(int64_t));
if ((sp_need_sum && !dense_sum) || !dense_count ||
(sp_need_sum && sp_need_sum128 && !dense_sum_hi)) {
scratch_free(_h_sum); scratch_free(_h_sum_hi); scratch_free(_h_cnt);
ray_release(key_col); ray_release(result);
sparse_i64_free(&sp_ht);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return ray_error("oom", NULL);
}
if (use_emit_filter && sp_need_sum) {
memset(dense_sum, 0, (size_t)grp_count * n_aggs * sizeof(da_val_t));
if (dense_sum_hi)
memset(dense_sum_hi, 0, (size_t)grp_count * n_aggs * sizeof(int64_t));
}
sparse_i64_ht_t heavy_ht;
memset(&heavy_ht, 0, sizeof(heavy_ht));
if (use_emit_filter && grp_count > 0) {
if (!sparse_i64_init(&heavy_ht, grp_count * 2u, n_aggs, false, false)) {
scratch_free(_h_sum); scratch_free(_h_sum_hi); scratch_free(_h_cnt);
ray_release(key_col); ray_release(result);
sparse_i64_free(&sp_ht);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return ray_error("oom", NULL);
}
}
uint32_t gi = 0;
for (uint32_t s = 0; s < sp_ht.cap; s++) {
if (!sp_ht.used[s]) continue;
if (sp_ht.counts[s] < keep_min) continue;
write_col_i64(ray_data(key_col), gi, sp_ht.keys[s],
key_col->type, key_col->attrs);
dense_count[gi] = sp_ht.counts[s];
if (use_emit_filter) {
int32_t hslot;
if (!sparse_i64_touch(&heavy_ht, sp_ht.keys[s], n_aggs, false, &hslot)) {
scratch_free(_h_sum); scratch_free(_h_sum_hi); scratch_free(_h_cnt);
ray_release(key_col); ray_release(result);
sparse_i64_free(&heavy_ht);
sparse_i64_free(&sp_ht);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return ray_error("oom", NULL);
}
heavy_ht.counts[hslot] = gi;
} else if (sp_need_sum) {
memcpy(&dense_sum[(size_t)gi * n_aggs],
&sp_ht.sums[(size_t)s * n_aggs],
(size_t)n_aggs * sizeof(da_val_t));
if (dense_sum_hi)
memcpy(&dense_sum_hi[(size_t)gi * n_aggs],
&sp_ht.sums_hi[(size_t)s * n_aggs],
(size_t)n_aggs * sizeof(int64_t));
}
gi++;
}
sparse_i64_free(&sp_ht);
if (use_emit_filter && sp_need_sum) {
for (int64_t i = 0; i < n_scan; i++) {
int64_t r = match_idx ? match_idx[i] : i;
if (!match_idx && rowsel && !group_rowsel_pass(rowsel, r))
continue;
int64_t key = read_by_esz(key_data[0], r, key_esz);
int32_t hslot = sparse_i64_find_slot(&heavy_ht, key);
if (!heavy_ht.used[hslot] || heavy_ht.keys[hslot] != key)
continue;
uint32_t out_gi = (uint32_t)heavy_ht.counts[hslot];
da_val_t* sums = &dense_sum[(size_t)out_gi * n_aggs];
for (uint32_t a = 0; a < n_aggs; a++) {
if (ext->agg_ops[a] == OP_COUNT || !agg_ptrs[a])
continue;
if (agg_strlen[a])
sums[a].i = wrap_add_i64(
sums[a].i,
group_strlen_at_cached(agg_vecs[a], r,
strlen_sym_strings,
strlen_sym_count));
else if (agg_f64_mask & ((uint64_t)1 << a))
sums[a].f += group_fp_at(agg_ptrs[a], agg_types[a], r);
else if (dense_sum_hi)
ray_i128_add(&dense_sum_hi[(size_t)out_gi * n_aggs + a],
(uint64_t*)&sums[a].i,
read_col_i64(agg_ptrs[a], r, agg_types[a], 0));
else
sums[a].i = wrap_add_i64(
sums[a].i,
read_col_i64(agg_ptrs[a], r, agg_types[a], 0));
}
}
}
sparse_i64_free(&heavy_ht);
ray_op_ext_t* key_ext = find_ext(g, ext->keys[0]);
int64_t name_id = key_ext ? key_ext->sym : 0;
result = ray_table_add_col(result, name_id, key_col);
ray_release(key_col);
emit_agg_columns(&result, g, ext, agg_vecs, grp_count, n_aggs,
(double*)dense_sum, (int64_t*)dense_sum, dense_sum_hi,
NULL, NULL, NULL, NULL,
dense_count, agg_affine, agg_prod, NULL, NULL,
NULL, NULL, NULL);
scratch_free(_h_sum); scratch_free(_h_sum_hi);
scratch_free(_h_cnt);
for (uint32_t a = 0; a < n_aggs; a++)
if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result;
}
}
ht_path:;
if (!group_materialize_prod_slots(agg_prod, agg_vecs, agg_owned,
n_aggs, nrows)) {
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return ray_error("oom", NULL);
}
uint8_t ght_need = 0;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop == OP_SUM || aop == OP_PROD || aop == OP_AVG || aop == OP_ALL || aop == OP_ANY || aop == OP_FIRST || aop == OP_LAST)
ght_need |= GHT_NEED_SUM;
if (aop == OP_AVG && agg_vecs[a] && !group_fp_type(agg_vecs[a]->type) &&
!agg_strlen[a] && group_avg_needs_hi(agg_vecs[a]->type, nrows))
ght_need |= GHT_NEED_SUM128;
if (aop == OP_STDDEV || aop == OP_STDDEV_POP || aop == OP_VAR || aop == OP_VAR_POP)
{ ght_need |= GHT_NEED_SUM; ght_need |= GHT_NEED_SUMSQ; }
if (agg_is_binary_agg(aop))
{ ght_need |= GHT_NEED_SUM; ght_need |= GHT_NEED_SUMSQ;
ght_need |= GHT_NEED_PEARSON; }
if (aop == OP_MIN) ght_need |= GHT_NEED_MIN;
if (aop == OP_MAX) ght_need |= GHT_NEED_MAX;
}
ght_layout_t ght_layout;
if (!ght_compute_layout(&ght_layout, n_keys, n_aggs, agg_vecs, agg_vecs2,
ght_need,
ext->agg_ops, key_types, key_vecs)) {
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]);
if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return ray_error("limit", "group: layout stride/slot budget exceeded (or OOM)");
}
uint32_t ht_cap = 2;
ray_pool_t* pool = ray_pool_get();
uint32_t n_total = pool ? ray_pool_total_workers(pool) : 1;
int64_t radix_rows = n_scan;
if (g->selection) {
ray_rowsel_t* sel_sm = ray_rowsel_meta(g->selection);
if (sel_sm && sel_sm->nrows == nrows && sel_sm->total_pass < radix_rows)
radix_rows = sel_sm->total_pass > 0 ? sel_sm->total_pass : 1;
}
uint32_t radix_n_parts = group_radix_part_count(n_total, radix_rows);
group_ht_t single_ht;
group_ht_t top_ht;
group_ht_t* final_ht = NULL;
bool top_ht_ready = false;
ray_t* result = NULL;
ray_t* radix_bufs_hdr = NULL;
radix_buf_t* radix_bufs = NULL;
ray_t* part_hts_hdr = NULL;
group_ht_t* part_hts = NULL;
ray_t* part_offsets_hdr = NULL;
uint32_t* part_offsets = NULL;
ray_t* group_out_hdr = NULL;
uint32_t* group_out = NULL;
ray_t* topn_fuse_hdr = NULL;
ray_t* topn_fuse_cnt_hdr = NULL;
topn_fuse_t topn_fuse = {0};
uint16_t topn_order_off = 0;
uint8_t topn_desc_dir = 0;
bool topn_order_ok = false;
if (use_topn_filter) {
uint16_t order_op = emit_filter.agg_op ? emit_filter.agg_op
: (uint16_t)OP_COUNT;
uint8_t ai = emit_filter.agg_index;
bool order_is_f64 = (ght_layout.agg_flags[ai] & GHT_AF_F64) != 0;
int8_t agg_slot = ght_layout.agg_val_slot[ai];
topn_order_ok = true;
if (order_op == OP_SUM) {
if (agg_slot < 0 || order_is_f64) topn_order_ok = false;
else topn_order_off = (uint16_t)(ght_layout.off_sum
+ (uint16_t)agg_slot * 8u);
} else if (order_op == OP_MIN) {
if (agg_slot < 0 || order_is_f64
|| (ght_layout.agg_flags[ai] & GHT_AF_SYM)) topn_order_ok = false;
else topn_order_off = (uint16_t)(ght_layout.off_min
+ (uint16_t)agg_slot * 8u);
} else if (order_op == OP_MAX) {
if (agg_slot < 0 || order_is_f64
|| (ght_layout.agg_flags[ai] & GHT_AF_SYM)) topn_order_ok = false;
else topn_order_off = (uint16_t)(ght_layout.off_max
+ (uint16_t)agg_slot * 8u);
}
topn_desc_dir = emit_filter.desc ? 1 : 0;
}
if (use_topn_filter && topn_order_ok && pool && n_total > 1 &&
radix_n_parts > 1 && emit_filter.top_count_take > 0) {
uint64_t kt = (uint64_t)emit_filter.top_count_take;
uint64_t tot = (uint64_t)radix_n_parts * kt;
if (kt <= UINT32_MAX && tot <= (uint64_t)radix_rows / 8 &&
tot <= UINT32_MAX && tot <= SIZE_MAX / sizeof(group_topn_item_t)) {
uint32_t* fcnt = (uint32_t*)scratch_calloc(&topn_fuse_cnt_hdr,
(size_t)radix_n_parts * sizeof(uint32_t));
group_topn_item_t* fit = fcnt
? (group_topn_item_t*)scratch_alloc(&topn_fuse_hdr,
(size_t)tot * sizeof(group_topn_item_t))
: NULL;
if (fit) {
topn_fuse.items = fit;
topn_fuse.item_cnt = fcnt;
topn_fuse.cap = (uint32_t)kt;
topn_fuse.order_off = topn_order_off;
topn_fuse.desc_dir = topn_desc_dir;
} else if (fcnt) {
scratch_free(topn_fuse_cnt_hdr);
topn_fuse_cnt_hdr = NULL;
}
}
}
if (pool && n_total > 1 &&
(nrows >= RAY_PARALLEL_THRESHOLD || use_topn_filter)) {
bool v2_ok = (n_keys >= 1 && n_aggs > 0);
for (uint32_t k = 0; k < n_keys && v2_ok; k++)
if (key_types[k] == RAY_SYM) v2_ok = false;
for (uint32_t a = 0; a < n_aggs && v2_ok; a++) {
uint16_t op = ext->agg_ops[a];
if (op != OP_COUNT && op != OP_SUM && op != OP_AVG) {
v2_ok = false;
break;
}
if (agg_vecs[a]) {
ray_t* src = (agg_vecs[a]->attrs & RAY_ATTR_SLICE)
? agg_vecs[a]->slice_parent : agg_vecs[a];
if (src && ray_vec_may_have_nulls(src))
v2_ok = false;
}
}
if (v2_ok && !(ght_layout.agg_flags_any & (GHT_AF_FIRST | GHT_AF_LAST
| GHT_AF_HOLISTIC | GHT_AF_BINARY))) {
ray_t* wpart_hdr = NULL;
size_t v2_n_w = (size_t)n_total * radix_n_parts;
group_ht_t* wpart_hts = (group_ht_t*)scratch_calloc(
&wpart_hdr, v2_n_w * sizeof(group_ht_t));
ray_t* v2_part_hdr = NULL;
group_ht_t* v2_part_hts = wpart_hts
? (group_ht_t*)scratch_calloc(&v2_part_hdr,
(size_t)radix_n_parts * sizeof(group_ht_t))
: NULL;
if (!wpart_hts || !v2_part_hts) {
if (wpart_hts) scratch_free(wpart_hdr);
if (v2_part_hts) scratch_free(v2_part_hdr);
goto v2_done;
}
uint8_t v2_nullable = 0;
for (uint32_t k = 0; k < n_keys; k++)
if (ray_key_may_be_null(key_vecs[k])) { v2_nullable = 1; break; }
ray_t* sel_match_block = NULL;
const int64_t* sel_match = NULL;
int64_t sel_n = 0;
if (g->selection) {
ray_rowsel_t* sm = ray_rowsel_meta(g->selection);
if (sm && sm->nrows == nrows
&& sm->total_pass < (nrows >> SEL_MATCH_GATE_SHIFT)) {
sel_match_block = ray_rowsel_to_indices(g->selection);
if (sel_match_block) {
sel_match = (const int64_t*)ray_data(sel_match_block);
sel_n = sm->total_pass;
}
}
}
uint64_t v2_rows_disp = (uint64_t)(sel_match ? sel_n : n_scan);
uint64_t v2_exp = v2_rows_disp
/ ((uint64_t)n_total * radix_n_parts) + 1;
uint32_t v2_cap = 2;
while (v2_cap < v2_exp && v2_cap < 256) v2_cap <<= 1;
uint32_t v2_grow = 2;
while (v2_grow < 2 * v2_exp && v2_grow < (1u << 24)) v2_grow <<= 1;
radix_v2_phase1_ctx_t v2p1 = {
.key_data = key_data,
.key_types = key_types,
.key_attrs = key_attrs,
.key_vecs = key_vecs,
.agg_vecs = agg_vecs,
.agg_vecs2 = agg_vecs2,
.agg_strlen = agg_strlen,
.nullable_mask = v2_nullable,
.n_workers = n_total,
.n_parts = radix_n_parts,
.wpart_hts = wpart_hts,
.rowsel = rowsel,
.match_idx = sel_match,
.ht_init_cap = v2_cap,
.ht_init_block = v2_cap > 2 ? v2_cap / 2 : 1,
.ht_grow_cap = v2_grow,
.oom = 0,
};
ght_layout_copy(&v2p1.layout, &ght_layout);
ray_pool_dispatch(pool, radix_v2_phase1_fn, &v2p1,
sel_match ? sel_n : n_scan);
if (sel_match_block) ray_release(sel_match_block);
CHECK_CANCEL_GOTO(pool, cleanup);
if (atomic_load_explicit(&v2p1.oom, memory_order_relaxed)) {
for (size_t i = 0; i < v2_n_w; i++)
group_ht_free(&wpart_hts[i]);
scratch_free(wpart_hdr);
scratch_free(v2_part_hdr);
goto v2_done;
}
radix_v2_phase2_ctx_t v2p2 = {
.wpart_hts = wpart_hts,
.part_hts = v2_part_hts,
.key_types = key_types,
.n_workers = n_total,
.n_parts = radix_n_parts,
.key_data = key_data,
.fuse = topn_fuse,
.oom = 0,
};
ght_layout_copy(&v2p2.layout, &ght_layout);
ray_t* v2p2_kp_hdr = NULL;
if (ght_layout.any_wide_key) {
v2p2.key_pool = (const void**)scratch_alloc(&v2p2_kp_hdr,
(size_t)(n_keys ? n_keys : 1) * sizeof(void*));
if (!v2p2.key_pool) {
for (size_t i = 0; i < v2_n_w; i++)
group_ht_free(&wpart_hts[i]);
scratch_free(wpart_hdr);
scratch_free(v2_part_hdr);
result = ray_error("oom", NULL); goto cleanup;
}
derive_key_pool(&ght_layout, key_vecs, v2p2.key_pool);
}
ray_pool_dispatch_n(pool, radix_v2_phase2_fn, &v2p2,
radix_n_parts);
scratch_free(v2p2_kp_hdr);
CHECK_CANCEL_GOTO(pool, cleanup);
for (size_t i = 0; i < v2_n_w; i++)
group_ht_free(&wpart_hts[i]);
scratch_free(wpart_hdr);
if (atomic_load_explicit(&v2p2.oom, memory_order_relaxed)) {
for (uint32_t p = 0; p < radix_n_parts; p++)
group_ht_free(&v2_part_hts[p]);
scratch_free(v2_part_hdr);
goto v2_done;
}
part_hts = v2_part_hts;
part_hts_hdr = v2_part_hdr;
goto v2_emit;
}
v2_done:;
size_t n_bufs = (size_t)n_total * radix_n_parts;
radix_bufs = (radix_buf_t*)scratch_calloc(&radix_bufs_hdr,
n_bufs * sizeof(radix_buf_t));
if (!radix_bufs) goto sequential_fallback;
uint8_t p1_nullable = 0;
for (uint32_t k = 0; k < n_keys; k++)
if (ray_key_may_be_null(key_vecs[k])) { p1_nullable = 1; break; }
radix_phase1_ctx_t p1ctx = {
.key_data = key_data,
.key_types = key_types,
.key_attrs = key_attrs,
.key_vecs = key_vecs,
.nullable_mask = p1_nullable,
.agg_vecs = agg_vecs,
.agg_vecs2 = agg_vecs2,
.agg_strlen = agg_strlen,
.n_workers = n_total,
.n_parts = radix_n_parts,
.bufs = radix_bufs,
.rowsel = rowsel,
.match_idx = match_idx,
.buf_prime = (uint32_t)((uint64_t)(n_scan > 0 ? n_scan : 1)
/ ((uint64_t)n_total * radix_n_parts) * 5 / 4 + 8),
};
ght_layout_copy(&p1ctx.layout, &ght_layout);
ray_t* p1_kp_hdr = NULL;
if (ght_layout.any_wide_key) {
p1ctx.key_pool = (const void**)scratch_alloc(&p1_kp_hdr,
(size_t)n_keys * sizeof(void*));
if (!p1ctx.key_pool) { result = ray_error("oom", NULL); goto cleanup; }
derive_key_pool(&ght_layout, key_vecs, p1ctx.key_pool);
}
ray_pool_dispatch(pool, radix_phase1_fn, &p1ctx, n_scan);
scratch_free(p1_kp_hdr);
CHECK_CANCEL_GOTO(pool, cleanup);
{
bool phase1_oom = false;
for (size_t i = 0; i < n_bufs; i++) {
if (radix_bufs[i].oom) { phase1_oom = true; break; }
}
if (phase1_oom) {
for (size_t i = 0; i < n_bufs; i++) scratch_free(radix_bufs[i]._hdr);
scratch_free(radix_bufs_hdr);
radix_bufs = NULL;
goto sequential_fallback;
}
}
part_hts = (group_ht_t*)scratch_calloc(&part_hts_hdr,
(size_t)radix_n_parts * sizeof(group_ht_t));
if (!part_hts) {
for (size_t i = 0; i < n_bufs; i++) scratch_free(radix_bufs[i]._hdr);
scratch_free(radix_bufs_hdr);
radix_bufs = NULL;
goto sequential_fallback;
}
radix_phase2_ctx_t p2ctx = {
.key_types = key_types,
.n_keys = n_keys,
.n_workers = n_total,
.n_parts = radix_n_parts,
.bufs = radix_bufs,
.part_hts = part_hts,
.key_data = key_data,
.fuse = topn_fuse,
};
ght_layout_copy(&p2ctx.layout, &ght_layout);
ray_t* p2_kp_hdr = NULL;
if (ght_layout.any_wide_key) {
p2ctx.key_pool = (const void**)scratch_alloc(&p2_kp_hdr,
(size_t)n_keys * sizeof(void*));
if (!p2ctx.key_pool) { result = ray_error("oom", NULL); goto cleanup; }
derive_key_pool(&ght_layout, key_vecs, p2ctx.key_pool);
}
ray_pool_dispatch_n(pool, radix_phase2_fn, &p2ctx, radix_n_parts);
scratch_free(p2_kp_hdr);
CHECK_CANCEL_GOTO(pool, cleanup);
if (radix_bufs) {
size_t n_bufs_free = (size_t)n_total * radix_n_parts;
for (size_t i = 0; i < n_bufs_free; i++)
scratch_free(radix_bufs[i]._hdr);
scratch_free(radix_bufs_hdr);
radix_bufs = NULL;
radix_bufs_hdr = NULL;
}
v2_emit:;
if (use_topn_filter) {
int64_t k_take = emit_filter.top_count_take;
uint64_t total_pre = 0;
for (uint32_t p = 0; p < radix_n_parts; p++)
total_pre += part_hts[p].grp_count;
if (!topn_order_ok) goto topn_compact_skip;
uint16_t order_off = topn_order_off;
uint8_t desc_dir = topn_desc_dir;
if (total_pre > (uint64_t)k_take && k_take > 0 &&
(uint64_t)k_take <= SIZE_MAX / sizeof(group_topn_item_t)) {
ray_t* heap_hdr = NULL;
group_topn_item_t* heap = (group_topn_item_t*)scratch_alloc(
&heap_hdr, (size_t)k_take * sizeof(group_topn_item_t));
if (!heap) goto topn_compact_skip;
int64_t hn = 0;
ray_t* stage_hdr = NULL;
ray_t* item_hdr = NULL;
const uint32_t* item_off = NULL;
const uint32_t* item_cnt = NULL;
const group_topn_item_t* items = NULL;
bool scanned = false;
if (topn_fuse.items) {
items = topn_fuse.items;
item_cnt = topn_fuse.item_cnt;
scanned = true;
} else if (pool && n_total > 1 && radix_n_parts > 1 &&
(size_t)radix_n_parts + 1 <= SIZE_MAX / sizeof(uint32_t) / 2) {
const uint64_t kcap = (uint64_t)k_take;
uint64_t tot = 0;
for (uint32_t p = 0; p < radix_n_parts; p++) {
uint64_t gc = part_hts[p].grp_count;
tot += gc < kcap ? gc : kcap;
}
if (tot > 0 && tot <= total_pre / 8 && tot <= UINT32_MAX &&
tot <= SIZE_MAX / sizeof(group_topn_item_t)) {
uint32_t* off_w = (uint32_t*)scratch_calloc(&stage_hdr,
((size_t)radix_n_parts * 2 + 1) * sizeof(uint32_t));
group_topn_item_t* items_w = off_w
? (group_topn_item_t*)scratch_alloc(&item_hdr,
(size_t)tot * sizeof(group_topn_item_t))
: NULL;
if (items_w) {
uint32_t* cnt_w = off_w + radix_n_parts + 1;
uint32_t run = 0;
for (uint32_t p = 0; p < radix_n_parts; p++) {
uint64_t gc = part_hts[p].grp_count;
off_w[p] = run;
run += (uint32_t)(gc < kcap ? gc : kcap);
}
off_w[radix_n_parts] = run;
items = items_w;
item_off = off_w;
item_cnt = cnt_w;
topn_scan_ctx_t sctx = {
.part_hts = part_hts, .items = items_w,
.item_off = off_w, .item_cnt = cnt_w,
.order_off = order_off, .desc_dir = desc_dir,
};
ray_pool_dispatch_n(pool, topn_scan_fn, &sctx,
radix_n_parts);
if (pool_cancelled(pool)) {
scratch_free(item_hdr);
scratch_free(stage_hdr);
scratch_free(heap_hdr);
result = ray_error("cancel", NULL);
goto cleanup;
}
scanned = true;
}
}
}
if (scanned)
for (uint32_t p = 0; p < radix_n_parts; p++) {
const group_topn_item_t* h = items
+ (item_off ? (size_t)item_off[p]
: (size_t)p * topn_fuse.cap);
uint32_t n = item_cnt[p];
for (uint32_t i = 0; i < n; i++)
hn = topn_heap_push(heap, hn, k_take, desc_dir,
h[i].value, h[i].part,
h[i].gid);
}
scratch_free(item_hdr);
scratch_free(stage_hdr);
if (!scanned)
for (uint32_t p = 0; p < radix_n_parts; p++) {
group_ht_t* ph = &part_hts[p];
uint16_t rs = ph->layout.row_stride;
uint32_t gc = ph->grp_count;
for (uint32_t gi = 0; gi < gc; gi++) {
int64_t v = *(const int64_t*)(const void*)
(ph->rows + (size_t)gi * rs + order_off);
hn = topn_heap_push(heap, hn, k_take, desc_dir,
v, p, gi);
}
}
if (hn > 0) {
qsort(heap, (size_t)hn, sizeof(*heap), topn_part_gid_cmp);
ray_t* dirty_hdr = NULL;
uint8_t* dirty = (uint8_t*)scratch_calloc(&dirty_hdr,
(size_t)radix_n_parts);
if (dirty) {
bool any_dirty = false;
int64_t i = 0;
for (uint32_t p = 0; p < radix_n_parts; p++) {
group_ht_t* ph = &part_hts[p];
uint16_t rs = ph->layout.row_stride;
uint32_t kn = 0;
for (; i < hn && heap[i].part == p; i++) {
uint32_t gi = heap[i].gid;
if (gi != kn)
memmove(ph->rows + (size_t)kn * rs,
ph->rows + (size_t)gi * rs, rs);
kn++;
}
if (kn == ph->grp_count) continue;
ph->grp_count = kn;
uint64_t need = (uint64_t)kn * 2;
uint32_t want = 2;
while ((uint64_t)want < need && want < ph->ht_cap)
want <<= 1;
if (want < ph->ht_cap) ph->ht_cap = want;
dirty[p] = 1;
any_dirty = true;
}
if (any_dirty) {
if (pool && n_total > 1 && radix_n_parts > 1) {
topn_rebuild_ctx_t rctx = {
.part_hts = part_hts,
.key_types = key_types, .dirty = dirty,
};
ray_pool_dispatch_n(pool, topn_rebuild_fn,
&rctx, radix_n_parts);
if (pool_cancelled(pool)) {
scratch_free(dirty_hdr);
scratch_free(heap_hdr);
result = ray_error("cancel", NULL);
goto cleanup;
}
} else {
for (uint32_t p = 0; p < radix_n_parts; p++)
if (dirty[p])
group_ht_rebuild_slots(&part_hts[p],
key_types);
}
}
}
scratch_free(dirty_hdr);
}
scratch_free(heap_hdr);
}
topn_compact_skip:;
}
part_offsets = (uint32_t*)scratch_alloc(&part_offsets_hdr,
((size_t)radix_n_parts + 1) * sizeof(uint32_t));
if (!part_offsets) { result = ray_error("oom", NULL); goto cleanup; }
part_offsets[0] = 0;
for (uint32_t p = 0; p < radix_n_parts; p++)
part_offsets[p + 1] = part_offsets[p] + part_hts[p].grp_count;
uint32_t total_grps = part_offsets[radix_n_parts];
if (total_grps > 0) {
group_out = (uint32_t*)scratch_alloc(&group_out_hdr,
(size_t)total_grps * sizeof(uint32_t));
if (!group_out) {
result = ray_error("oom", NULL);
goto cleanup;
}
uint64_t row_count = nrows > 0 ? (uint64_t)nrows : 0;
if (row_count > SIZE_MAX / sizeof(uint32_t)) {
result = ray_error("limit", "group row index is too large");
goto cleanup;
}
if ((uint64_t)total_grps <= row_count / 16) {
ray_t* pair_hdr = NULL;
group_order_pair_t* pairs = (group_order_pair_t*)scratch_alloc(
&pair_hdr, (size_t)total_grps * sizeof(group_order_pair_t));
if (!pairs) {
result = ray_error("oom", NULL);
goto cleanup;
}
bool order_ok = true;
uint32_t np = 0;
for (uint32_t p = 0; p < radix_n_parts && order_ok; p++) {
group_ht_t* ph = &part_hts[p];
for (uint32_t gi = 0; gi < ph->grp_count; gi++) {
uint32_t flat = part_offsets[p] + gi;
const char* group_row = ph->rows
+ (size_t)gi * ph->layout.row_stride;
int64_t first;
memcpy(&first,
group_row + ph->layout.off_group_first, 8);
if (first < 0 || first >= nrows) {
order_ok = false;
break;
}
pairs[np].first = first;
pairs[np].flat = flat;
np++;
}
}
if (order_ok) {
qsort(pairs, np, sizeof(group_order_pair_t),
group_order_pair_cmp);
for (uint32_t i = 0; i < np; i++) {
if (i > 0 && pairs[i].first == pairs[i - 1].first) {
order_ok = false;
break;
}
group_out[pairs[i].flat] = i;
}
}
scratch_free(pair_hdr);
if (!order_ok || np != total_grps) {
result = ray_error("group", "failed to order all groups");
goto cleanup;
}
} else {
ray_t* first_hdr = NULL;
uint32_t* first_to_flat = (uint32_t*)scratch_alloc(
&first_hdr, (size_t)row_count * sizeof(uint32_t));
if (!first_to_flat) {
result = ray_error("oom", NULL);
goto cleanup;
}
memset(first_to_flat, 0xff,
(size_t)row_count * sizeof(uint32_t));
bool order_ok = true;
for (uint32_t p = 0; p < radix_n_parts && order_ok; p++) {
group_ht_t* ph = &part_hts[p];
for (uint32_t gi = 0; gi < ph->grp_count; gi++) {
uint32_t flat = part_offsets[p] + gi;
const char* group_row = ph->rows
+ (size_t)gi * ph->layout.row_stride;
int64_t first;
memcpy(&first,
group_row + ph->layout.off_group_first, 8);
if (first < 0 || first >= nrows ||
first_to_flat[first] != UINT32_MAX) {
order_ok = false;
break;
}
first_to_flat[first] = flat;
}
}
uint32_t out = 0;
if (order_ok)
for (int64_t row = 0; row < nrows; row++) {
uint32_t flat = first_to_flat[row];
if (flat != UINT32_MAX) group_out[flat] = out++;
}
scratch_free(first_hdr);
if (!order_ok || out != total_grps) {
result = ray_error("group", "failed to order all groups");
goto cleanup;
}
}
}
int64_t total_cols = n_keys + n_aggs;
result = ray_table_new(total_cols);
if (!result || RAY_IS_ERR(result)) goto cleanup;
ray_t* key_cols[n_keys];
char* key_dsts[n_keys];
int8_t key_out_types[n_keys];
uint8_t key_esizes[n_keys];
for (uint32_t k = 0; k < n_keys; k++) {
ray_t* src_col = key_vecs[k];
key_cols[k] = NULL;
key_dsts[k] = NULL;
key_out_types[k] = 0;
key_esizes[k] = 0;
if (!src_col) continue;
uint8_t esz = ray_sym_elem_size(src_col->type, src_col->attrs);
ray_t* new_col;
if (src_col->type == RAY_SYM) {
new_col = ray_sym_vec_new(src_col->attrs & RAY_SYM_W_MASK, (int64_t)total_grps);
if (new_col && !RAY_IS_ERR(new_col))
ray_sym_vec_adopt_domain(new_col, sym_domain_rep(src_col));
} else if (src_col->type == RAY_STR) {
new_col = ray_vec_new(RAY_STR, (int64_t)total_grps);
if (new_col && !RAY_IS_ERR(new_col)) {
ray_t* owner = (src_col->attrs & RAY_ATTR_SLICE)
? src_col->slice_parent : src_col;
if (owner && owner->str_pool && !RAY_IS_ERR(owner->str_pool)) {
ray_retain(owner->str_pool);
new_col->str_pool = owner->str_pool;
}
}
} else
new_col = ray_vec_new(src_col->type, (int64_t)total_grps);
if (!new_col || RAY_IS_ERR(new_col)) continue;
new_col->len = (int64_t)total_grps;
const ray_t* null_src = (src_col->attrs & RAY_ATTR_SLICE)
? src_col->slice_parent : src_col;
if (null_src && (null_src->attrs & RAY_ATTR_HAS_NULLS))
new_col->attrs |= RAY_ATTR_HAS_NULLS;
key_cols[k] = new_col;
key_dsts[k] = (char*)ray_data(new_col);
key_out_types[k] = src_col->type;
key_esizes[k] = esz;
}
agg_out_t agg_outs[vla_aggs];
ray_t* agg_cols[vla_aggs];
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t agg_op = ext->agg_ops[a];
ray_t* agg_col = agg_vecs[a];
bool is_f64 = (ght_layout.agg_flags[a] & GHT_AF_F64) != 0;
int8_t out_type;
switch (agg_op) {
case OP_AVG:
case OP_STDDEV: case OP_STDDEV_POP:
case OP_VAR: case OP_VAR_POP:
case OP_PEARSON_CORR:
case OP_COV: case OP_SCOV:
case OP_WSUM: case OP_WAVG:
case OP_MEDIAN:
case OP_QUANTILE:
out_type = RAY_F64; break;
case OP_ALL:
case OP_ANY:
out_type = RAY_BOOL; break;
case OP_COUNT: out_type = RAY_I64; break;
case OP_SUM:
out_type = is_f64 ? RAY_F64
: (agg_col && agg_col->type == RAY_TIME ? RAY_TIME : RAY_I64);
break;
case OP_PROD:
out_type = is_f64 ? RAY_F64 : RAY_I64; break;
default:
out_type = is_f64 ? RAY_F64
: (agg_col ? agg_col->type : RAY_I64); break;
}
ray_t* new_col = out_type == RAY_LIST
? ray_list_new((int64_t)total_grps)
: ray_vec_new(out_type, (int64_t)total_grps);
if (!new_col || RAY_IS_ERR(new_col)) {
agg_cols[a] = NULL;
memset(&agg_outs[a], 0, sizeof(agg_outs[a]));
continue;
}
if (out_type == RAY_SYM)
ray_sym_vec_adopt_domain(new_col, sym_domain_rep(agg_col));
new_col->len = (int64_t)total_grps;
agg_cols[a] = new_col;
agg_outs[a] = (agg_out_t){
.out_type = out_type, .src_f64 = is_f64,
.agg_op = agg_op,
.affine = agg_affine[a].enabled,
.bias_f64 = agg_affine[a].bias_f64,
.bias_i64 = agg_affine[a].bias_i64,
.dst = ray_data(new_col),
.vec = new_col,
};
}
{
radix_phase3_ctx_t p3ctx = {
.part_hts = part_hts,
.part_offsets = part_offsets,
.group_out = group_out,
.key_dsts = key_dsts,
.key_types = key_out_types,
.key_attrs = key_attrs,
.key_esizes = key_esizes,
.n_keys = n_keys,
.agg_outs = agg_outs,
.n_aggs = n_aggs,
.key_src_data = key_data,
};
ray_pool_dispatch_n(pool, radix_phase3_fn, &p3ctx,
radix_n_parts);
}
if (ght_layout.agg_flags_any & GHT_AF_HOLISTIC) {
int64_t n_groups = (int64_t)total_grps;
ray_t* rg_hdr = NULL;
int64_t* row_gid = (int64_t*)scratch_alloc(&rg_hdr,
(size_t)nrows * sizeof(int64_t));
if (!row_gid) { result = ray_error("oom", NULL); goto cleanup; }
memset(row_gid, 0xff, (size_t)nrows * sizeof(int64_t));
uint8_t reprobe_nullable = 0;
for (uint32_t k = 0; k < n_keys; k++)
if (ray_key_may_be_null(key_vecs[k])) { reprobe_nullable = 1; break; }
reprobe_ctx_t rp = {
.key_data = key_data,
.key_types = key_types,
.key_attrs = key_attrs,
.key_vecs = key_vecs,
.n_keys = n_keys,
.nullable_mask = reprobe_nullable,
.layout = &ght_layout,
.part_hts = part_hts,
.part_offsets = part_offsets,
.group_out = group_out,
.n_parts = radix_n_parts,
.row_gid = row_gid,
.match_idx = match_idx,
.rowsel = rowsel,
};
ray_t* rp_kp_hdr = NULL;
if (ght_layout.any_wide_key) {
rp.key_pool = (const void**)scratch_alloc(&rp_kp_hdr,
(size_t)n_keys * sizeof(void*));
if (!rp.key_pool) { scratch_free(rg_hdr); result = ray_error("oom", NULL); goto cleanup; }
derive_key_pool(&ght_layout, key_vecs, rp.key_pool);
}
ray_pool_dispatch(pool, reprobe_rows_fn, &rp, n_scan);
scratch_free(rp_kp_hdr);
int64_t n_workers = (int64_t)ray_pool_total_workers(pool);
int64_t med_ntasks = n_workers > 1 ? n_workers : 1;
int64_t min_grain = 8192;
if (med_ntasks * min_grain > nrows)
med_ntasks = (nrows + min_grain - 1) / min_grain;
if (med_ntasks < 1) med_ntasks = 1;
int64_t med_grain = (nrows + med_ntasks - 1) / med_ntasks;
if (med_grain < 1) med_grain = 1;
med_ntasks = (nrows + med_grain - 1) / med_grain;
ray_t* hist_hdr = NULL;
ray_t* cur_hdr = NULL;
ray_t* cnt_hdr = NULL;
ray_t* off_hdr = NULL;
int64_t* hist = (int64_t*)scratch_calloc(&hist_hdr,
(size_t)med_ntasks * (size_t)n_groups * sizeof(int64_t));
int64_t* cur = (int64_t*)scratch_alloc(&cur_hdr,
(size_t)med_ntasks * (size_t)n_groups * sizeof(int64_t));
int64_t* grp_cnt = (int64_t*)scratch_alloc(&cnt_hdr,
(size_t)n_groups * sizeof(int64_t));
int64_t* offsets = (int64_t*)scratch_alloc(&off_hdr,
(size_t)n_groups * sizeof(int64_t));
ray_t* idx_hdr = NULL;
int64_t* idx_buf = NULL;
if (hist && cur && grp_cnt && offsets) {
med_idx_ctx_t mctx = {
.row_gid = row_gid,
.hist = hist,
.cursor = cur,
.idx_buf = NULL,
.n_groups = n_groups,
.grain = med_grain,
.nrows = nrows,
};
ray_pool_dispatch_n(pool, med_idx_hist_fn, &mctx,
(uint32_t)med_ntasks);
int64_t total = 0;
for (int64_t gi = 0; gi < n_groups; gi++) {
int64_t cum = total;
for (int64_t t = 0; t < med_ntasks; t++) {
int64_t cn = hist[t * n_groups + gi];
cur[t * n_groups + gi] = cum;
cum += cn;
}
grp_cnt[gi] = cum - total;
offsets[gi] = total;
total = cum;
}
idx_buf = (int64_t*)scratch_alloc(&idx_hdr,
(size_t)(total > 0 ? total : 1) * sizeof(int64_t));
if (idx_buf) {
mctx.idx_buf = idx_buf;
ray_pool_dispatch_n(pool, med_idx_scat_fn, &mctx,
(uint32_t)med_ntasks);
}
}
if (idx_buf) {
for (uint32_t a = 0; a < n_aggs; a++) {
if (!(ght_layout.agg_flags[a] & GHT_AF_HOLISTIC)) continue;
if (!agg_vecs[a] || !agg_cols[a]) continue;
uint16_t aop = ext->agg_ops[a];
ray_t* hol_vec = NULL;
const char* err_tag = "median: type";
if (aop == OP_MEDIAN) {
hol_vec = ray_median_per_group_buf(
agg_vecs[a], idx_buf, offsets, grp_cnt, n_groups);
} else if (aop == OP_QUANTILE) {
double q = ext->agg_k
? decode_agg_f64_param(ext->agg_k[a]) : 0.5;
hol_vec = ray_quantile_per_group_buf(
agg_vecs[a], idx_buf, offsets, grp_cnt, n_groups, q);
err_tag = "quantile: type";
} else if (aop == OP_MODE) {
hol_vec = ray_mode_per_group_buf(
agg_vecs[a], idx_buf, offsets, grp_cnt, n_groups);
err_tag = "mode: type";
} else if (aop == OP_TOP_N || aop == OP_BOT_N) {
int64_t k_val = (ext->agg_k && ext->agg_k[a] > 0)
? ext->agg_k[a] : 1;
hol_vec = ray_topk_per_group_buf(
agg_vecs[a], k_val,
aop == OP_TOP_N ? 1 : 0,
idx_buf, offsets, grp_cnt, n_groups);
err_tag = "top/bot: type";
} else if (aop == OP_MIN || aop == OP_MAX ||
aop == OP_FIRST || aop == OP_LAST) {
hol_vec = ray_wide_minmax_per_group_buf(
agg_vecs[a], aop, idx_buf, offsets, grp_cnt, n_groups);
err_tag = "minmax: type";
}
if (!hol_vec) {
if (hist_hdr) scratch_free(hist_hdr);
if (cur_hdr) scratch_free(cur_hdr);
if (cnt_hdr) scratch_free(cnt_hdr);
if (off_hdr) scratch_free(off_hdr);
if (idx_hdr) scratch_free(idx_hdr);
scratch_free(rg_hdr);
result = ray_error("nyi", err_tag);
goto cleanup;
}
if (RAY_IS_ERR(hol_vec)) {
if (hist_hdr) scratch_free(hist_hdr);
if (cur_hdr) scratch_free(cur_hdr);
if (cnt_hdr) scratch_free(cnt_hdr);
if (off_hdr) scratch_free(off_hdr);
if (idx_hdr) scratch_free(idx_hdr);
scratch_free(rg_hdr);
result = hol_vec;
goto cleanup;
}
ray_release(agg_cols[a]);
agg_cols[a] = hol_vec;
agg_outs[a].vec = hol_vec;
}
} else {
if (hist_hdr) scratch_free(hist_hdr);
if (cur_hdr) scratch_free(cur_hdr);
if (cnt_hdr) scratch_free(cnt_hdr);
if (off_hdr) scratch_free(off_hdr);
if (idx_hdr) scratch_free(idx_hdr);
scratch_free(rg_hdr);
result = ray_error("oom", NULL);
goto cleanup;
}
if (hist_hdr) scratch_free(hist_hdr);
if (cur_hdr) scratch_free(cur_hdr);
if (cnt_hdr) scratch_free(cnt_hdr);
if (off_hdr) scratch_free(off_hdr);
if (idx_hdr) scratch_free(idx_hdr);
scratch_free(rg_hdr);
}
for (uint32_t a = 0; a < n_aggs; a++) {
if (!agg_cols[a]) continue;
if (agg_outs[a].vec && agg_outs[a].vec->type == RAY_LIST) continue;
grp_finalize_nulls(agg_outs[a].vec);
}
for (uint32_t k = 0; k < n_keys; k++) {
if (!key_cols[k]) continue;
grp_finalize_nulls(key_cols[k]);
}
for (uint32_t k = 0; k < n_keys; k++) {
if (!key_cols[k]) continue;
ray_op_ext_t* key_ext = find_ext(g, ext->keys[k]);
int64_t name_id = key_ext ? key_ext->sym : k;
result = ray_table_add_col(result, name_id, key_cols[k]);
ray_release(key_cols[k]);
}
for (uint32_t a = 0; a < n_aggs; a++) {
if (!agg_cols[a]) continue;
uint16_t agg_op = ext->agg_ops[a];
ray_op_ext_t* agg_ext = find_ext(g, ext->agg_ins[a]);
int64_t name_id;
if (agg_ext && agg_ext->base.opcode == OP_SCAN) {
ray_t* name_atom = ray_sym_str(agg_ext->sym);
const char* base = name_atom ? ray_str_ptr(name_atom) : NULL;
size_t blen = base ? ray_str_len(name_atom) : 0;
const char* sfx = "";
size_t slen = 0;
switch (agg_op) {
case OP_SUM: sfx = "_sum"; slen = 4; break;
case OP_PROD: sfx = "_prod"; slen = 5; break;
case OP_ALL: sfx = "_all"; slen = 4; break;
case OP_ANY: sfx = "_any"; slen = 4; break;
case OP_COUNT: sfx = "_count"; slen = 6; break;
case OP_AVG: sfx = "_mean"; slen = 5; break;
case OP_MIN: sfx = "_min"; slen = 4; break;
case OP_MAX: sfx = "_max"; slen = 4; break;
case OP_FIRST: sfx = "_first"; slen = 6; break;
case OP_LAST: sfx = "_last"; slen = 5; break;
case OP_STDDEV: sfx = "_stddev"; slen = 7; break;
case OP_STDDEV_POP: sfx = "_stddev_pop"; slen = 11; break;
case OP_VAR: sfx = "_var"; slen = 4; break;
case OP_VAR_POP: sfx = "_var_pop"; slen = 8; break;
case OP_MEDIAN: sfx = "_median"; slen = 7; break;
case OP_QUANTILE: sfx = "_quantile"; slen = 9; break;
case OP_MODE: sfx = "_mode"; slen = 5; break;
case OP_COV: sfx = "_cov"; slen = 4; break;
case OP_SCOV: sfx = "_scov"; slen = 5; break;
case OP_WSUM: sfx = "_wsum"; slen = 5; break;
case OP_WAVG: sfx = "_wavg"; slen = 5; break;
case OP_TOP_N: sfx = "_top"; slen = 4; break;
case OP_BOT_N: sfx = "_bot"; slen = 4; break;
}
char buf[256];
ray_t* name_dyn_hdr = NULL;
char* nbp = buf;
size_t nbc = sizeof(buf);
if (base && blen + slen >= sizeof(buf)) {
nbp = (char*)scratch_alloc(&name_dyn_hdr, blen + slen + 1);
if (nbp) nbc = blen + slen + 1;
else { nbp = buf; nbc = sizeof(buf); }
}
if (base && blen + slen < nbc) {
memcpy(nbp, base, blen);
memcpy(nbp + blen, sfx, slen);
name_id = ray_sym_intern(nbp, blen + slen);
} else {
name_id = agg_ext->sym;
}
scratch_free(name_dyn_hdr);
} else {
name_id = (int64_t)(n_keys + a);
}
result = ray_table_add_col(result, name_id, agg_cols[a]);
ray_release(agg_cols[a]);
}
goto cleanup;
}
sequential_fallback:;
if (!group_ht_init(&single_ht, ht_cap, &ght_layout)) {
result = ray_error("oom", NULL);
goto cleanup;
}
group_rows_range(&single_ht, key_data, key_types, key_attrs, key_vecs, agg_vecs,
agg_vecs2, agg_strlen, rowsel,
0, n_scan, match_idx);
final_ht = &single_ht;
if (ray_interrupted()) { result = ray_error("cancel", "interrupted"); goto cleanup; }
if (single_ht.oom) { result = ray_error("oom", NULL); goto cleanup; }
{
uint32_t grp_count = final_ht->grp_count;
const ght_layout_t* ly = &final_ht->layout;
int64_t total_cols = n_keys + n_aggs;
result = ray_table_new(total_cols);
if (!result || RAY_IS_ERR(result)) goto cleanup;
for (uint32_t k = 0; k < n_keys; k++) {
ray_t* src_col = key_vecs[k];
if (!src_col) continue;
uint8_t esz = col_esz(src_col);
int8_t kt = src_col->type;
ray_t* new_col = col_vec_new(src_col, (int64_t)grp_count);
if (!new_col || RAY_IS_ERR(new_col)) continue;
out_col_adopt_str_pool(new_col, src_col);
if (new_col->type == RAY_SYM)
ray_sym_vec_adopt_domain(new_col, sym_domain_rep(src_col));
new_col->len = (int64_t)grp_count;
bool is_wide = (ly->key_flags[k] & GHT_KEYF_WIDE) != 0;
const char* src_base = is_wide ? (const char*)key_data[k] : NULL;
bool inline_str_k = (ly->key_flags[k] & GHT_KEYF_INLINE_STR) != 0;
size_t null_woff = ly->null_words
? (size_t)ly->key_off[n_keys] + (size_t)(k >> 6) * 8 : 0;
int64_t null_kbit = (int64_t)((uint64_t)1 << (k & 63));
for (uint32_t gi = 0; gi < grp_count; gi++) {
const char* row = final_ht->rows + (size_t)gi * ly->row_stride;
const char* rk = row + 8;
int64_t null_word = ly->null_words
? *(const int64_t*)(rk + null_woff) : 0;
if (null_word & null_kbit) {
ray_vec_set_null(new_col, (int64_t)gi, true);
switch (kt) {
case RAY_F64:
((double*)ray_data(new_col))[gi] = NULL_F64; break;
case RAY_I64: case RAY_TIMESTAMP:
((int64_t*)ray_data(new_col))[gi] = NULL_I64; break;
case RAY_I32: case RAY_DATE: case RAY_TIME:
((int32_t*)ray_data(new_col))[gi] = NULL_I32; break;
case RAY_I16:
((int16_t*)ray_data(new_col))[gi] = NULL_I16; break;
case RAY_STR:
memset((char*)ray_data(new_col) + (size_t)gi * esz, 0,
sizeof(ray_str_t)); break;
default: break;
}
continue;
}
if (inline_str_k) {
memcpy((char*)ray_data(new_col) + (size_t)gi * esz,
rk + ly->key_off[k], sizeof(ray_str_t));
continue;
}
int64_t kv = *(const int64_t*)(rk + ly->key_off[k]);
if (is_wide) {
char* dst = (char*)ray_data(new_col) + (size_t)gi * esz;
memcpy(dst, src_base + (size_t)kv * esz, esz);
} else if (kt == RAY_F64) {
char* dst = (char*)ray_data(new_col) + (size_t)gi * esz;
memcpy(dst, &kv, 8);
} else {
write_col_i64(ray_data(new_col), gi, kv, kt, new_col->attrs);
}
}
ray_op_ext_t* key_ext = find_ext(g, ext->keys[k]);
int64_t name_id = key_ext ? key_ext->sym : k;
result = ray_table_add_col(result, name_id, new_col);
ray_release(new_col);
}
ray_t** med_out = NULL;
ray_t* med_hdr = NULL;
if (ly->agg_flags_any & GHT_AF_HOLISTIC) {
med_out = (ray_t**)scratch_calloc(&med_hdr,
(size_t)n_aggs * sizeof(ray_t*));
if (med_out) {
ray_t* rg_hdr = NULL;
int64_t* row_gid = (int64_t*)scratch_alloc(&rg_hdr,
(size_t)nrows * sizeof(int64_t));
ray_t* cnt_hdr_s = NULL;
int64_t* grp_cnt_s = (int64_t*)scratch_calloc(&cnt_hdr_s,
(size_t)grp_count * sizeof(int64_t));
ray_t* off_hdr_s = NULL;
int64_t* offsets_s = (int64_t*)scratch_alloc(&off_hdr_s,
(size_t)grp_count * sizeof(int64_t));
ray_t* pos_hdr_s = NULL;
int64_t* pos_s = (int64_t*)scratch_alloc(&pos_hdr_s,
(size_t)grp_count * sizeof(int64_t));
if (row_gid && grp_cnt_s && offsets_s && pos_s) {
uint8_t reprobe_nullable_s = 0;
for (uint32_t k = 0; k < n_keys; k++)
if (ray_key_may_be_null(key_vecs[k])) { reprobe_nullable_s = 1; break; }
const void* reprobe_pool_stk[8];
int64_t ek_buf_stk[9];
char keybuf_stk[136];
const void** reprobe_pool = reprobe_pool_stk;
int64_t* ek_buf = ek_buf_stk;
char* keybuf = keybuf_stk;
ray_t* hol_stage_hdr = NULL;
if ((size_t)ly->key_region > sizeof(ek_buf_stk) || n_keys > 8) {
size_t kb = ly->key_region;
char* blk = (char*)scratch_alloc(&hol_stage_hdr,
kb + kb + (size_t)n_keys * sizeof(void*));
if (!blk) {
scratch_free(rg_hdr); scratch_free(cnt_hdr_s);
scratch_free(off_hdr_s); scratch_free(pos_hdr_s);
result = ray_error("oom", NULL); goto cleanup;
}
keybuf = blk;
ek_buf = (int64_t*)(blk + kb);
reprobe_pool = (const void**)(blk + kb + kb);
}
derive_key_pool(ly, key_vecs, reprobe_pool);
for (int64_t i = 0; i < n_scan; i++) {
int64_t row = match_idx ? match_idx[i] : i;
if (!match_idx && !group_rowsel_pass(rowsel, row)) continue;
uint64_t h = 0;
const int64_t* lookup_keys;
if (ly->any_inline_str) {
h = inline_build_keys(ly, key_types, key_data, key_attrs,
reprobe_pool, key_vecs, reprobe_nullable_s,
row, keybuf);
lookup_keys = (const int64_t*)keybuf;
} else {
int64_t* nullw = ek_buf + n_keys;
uint32_t null_words = ly->null_words;
for (uint32_t w = 0; w < null_words; w++) nullw[w] = 0;
if (ly->packed_key) {
for (uint32_t k = 0; k < n_keys; k++) {
if (__builtin_expect(reprobe_nullable_s && ray_key_may_be_null(key_vecs[k])
&& ray_vec_is_null(key_vecs[k], row), 0)) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
ek_buf[k] = 0;
} else {
ek_buf[k] = read_col_i64(key_data[k], row, key_types[k], key_attrs[k]);
}
}
h = ght_hash_lanes(ek_buf, (uint32_t)n_keys + null_words);
} else {
for (uint32_t k = 0; k < n_keys; k++) {
int8_t t = key_types[k];
uint64_t kh;
bool is_null = reprobe_nullable_s && ray_key_may_be_null(key_vecs[k])
&& ray_vec_is_null(key_vecs[k], row);
if (is_null) {
nullw[k >> 6] |= (int64_t)((uint64_t)1 << (k & 63));
ek_buf[k] = 0;
kh = ray_hash_i64(0);
} else if (ly->key_flags[k] & GHT_KEYF_WIDE) {
ek_buf[k] = row;
kh = wide_key_hash_at(ly, k, key_data, reprobe_pool, row);
} else if (t == RAY_F64) {
int64_t kv = group_key_f64_bits(key_data[k], row);
double dv; memcpy(&dv, &kv, 8);
ek_buf[k] = kv;
kh = ray_hash_f64(dv);
} else {
int64_t kv = read_col_i64(key_data[k], row, t, key_attrs[k]);
ek_buf[k] = kv;
kh = ray_hash_i64(kv);
}
h = (k == 0) ? kh : ray_hash_combine(h, kh);
}
h = ght_hash_null_words(h, nullw, null_words);
}
lookup_keys = ek_buf;
}
uint32_t gid = group_ht_lookup_gid(final_ht, h, lookup_keys, key_types);
row_gid[row] = (gid == UINT32_MAX) ? -1 : (int64_t)gid;
if (gid != UINT32_MAX) grp_cnt_s[gid]++;
}
int64_t total_s = 0;
for (uint32_t gi = 0; gi < grp_count; gi++) {
offsets_s[gi] = total_s;
pos_s[gi] = total_s;
total_s += grp_cnt_s[gi];
}
ray_t* ix_hdr_s = NULL;
int64_t* idx_buf_s = (int64_t*)scratch_alloc(&ix_hdr_s,
(size_t)(total_s > 0 ? total_s : 1) * sizeof(int64_t));
if (idx_buf_s) {
for (int64_t i = 0; i < n_scan; i++) {
int64_t row = match_idx ? match_idx[i] : i;
if (!match_idx && !group_rowsel_pass(rowsel, row)) continue;
int64_t gi = row_gid[row];
if (gi >= 0) idx_buf_s[pos_s[gi]++] = row;
}
for (uint32_t a = 0; a < n_aggs; a++) {
if (!(ly->agg_flags[a] & GHT_AF_HOLISTIC)) continue;
if (!agg_vecs[a]) continue;
uint16_t aop = ext->agg_ops[a];
ray_t* hol_vec = NULL;
if (aop == OP_MEDIAN) {
hol_vec = ray_median_per_group_buf(
agg_vecs[a], idx_buf_s, offsets_s, grp_cnt_s,
(int64_t)grp_count);
} else if (aop == OP_QUANTILE) {
double q = ext->agg_k
? decode_agg_f64_param(ext->agg_k[a]) : 0.5;
hol_vec = ray_quantile_per_group_buf(
agg_vecs[a], idx_buf_s, offsets_s, grp_cnt_s,
(int64_t)grp_count, q);
} else if (aop == OP_MODE) {
hol_vec = ray_mode_per_group_buf(
agg_vecs[a], idx_buf_s, offsets_s, grp_cnt_s,
(int64_t)grp_count);
} else if (aop == OP_TOP_N || aop == OP_BOT_N) {
int64_t k_val = (ext->agg_k && ext->agg_k[a] > 0)
? ext->agg_k[a] : 1;
hol_vec = ray_topk_per_group_buf(
agg_vecs[a], k_val,
aop == OP_TOP_N ? 1 : 0,
idx_buf_s, offsets_s, grp_cnt_s,
(int64_t)grp_count);
} else if (aop == OP_MIN || aop == OP_MAX ||
aop == OP_FIRST || aop == OP_LAST) {
hol_vec = ray_wide_minmax_per_group_buf(
agg_vecs[a], aop, idx_buf_s, offsets_s,
grp_cnt_s, (int64_t)grp_count);
}
med_out[a] = hol_vec;
}
scratch_free(ix_hdr_s);
}
scratch_free(hol_stage_hdr);
}
scratch_free(rg_hdr);
scratch_free(cnt_hdr_s);
scratch_free(off_hdr_s);
scratch_free(pos_hdr_s);
}
}
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t agg_op = ext->agg_ops[a];
ray_t* agg_col = agg_vecs[a];
bool is_f64 = (ly->agg_flags[a] & GHT_AF_F64) != 0;
int8_t out_type;
switch (agg_op) {
case OP_AVG:
case OP_STDDEV: case OP_STDDEV_POP:
case OP_VAR: case OP_VAR_POP:
case OP_PEARSON_CORR:
case OP_COV: case OP_SCOV:
case OP_WSUM: case OP_WAVG:
case OP_MEDIAN:
case OP_QUANTILE:
out_type = RAY_F64; break;
case OP_ALL:
case OP_ANY:
out_type = RAY_BOOL; break;
case OP_COUNT: out_type = RAY_I64; break;
case OP_SUM: {
int8_t src_t = agg_col ? agg_col->type
: (op_node(g, ext->agg_ins[a]) ? op_node(g, ext->agg_ins[a])->out_type : 0);
out_type = is_f64 ? RAY_F64 : (src_t == RAY_TIME ? RAY_TIME : RAY_I64);
break;
}
case OP_PROD:
out_type = is_f64 ? RAY_F64 : RAY_I64; break;
default:
out_type = is_f64 ? RAY_F64
: (agg_col ? agg_col->type : RAY_I64); break;
}
ray_t* new_col;
bool is_holistic = (ly->agg_flags[a] & GHT_AF_HOLISTIC) != 0;
if (is_holistic && med_out && med_out[a]
&& !RAY_IS_ERR(med_out[a])) {
new_col = med_out[a];
med_out[a] = NULL;
} else if (is_holistic) {
continue;
} else {
new_col = ray_vec_new(out_type, (int64_t)grp_count);
if (!new_col || RAY_IS_ERR(new_col)) continue;
if (out_type == RAY_SYM)
ray_sym_vec_adopt_domain(new_col, sym_domain_rep(agg_col));
new_col->len = (int64_t)grp_count;
}
int8_t s = ly->agg_val_slot[a];
if (is_holistic) goto med_attach;
for (uint32_t gi = 0; gi < grp_count; gi++) {
const char* row = final_ht->rows + (size_t)gi * ly->row_stride;
int64_t cnt = *(const int64_t*)(const void*)row;
int64_t nn = (ly->off_nn && s >= 0)
? ((const int64_t*)(const void*)(row + ly->off_nn))[s] : cnt;
if (out_type == RAY_F64) {
double v;
switch (agg_op) {
case OP_SUM:
v = is_f64 ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
if (agg_affine[a].enabled) v += agg_affine[a].bias_f64 * cnt;
break;
case OP_PROD:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
break;
case OP_AVG:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? ROW_RD_F64(row, ly->off_sum, s) / nn
: (ly->need_flags & GHT_NEED_SUM128)
? ray_i128_to_f64(ROW_RD_I64(row, ly->off_sum_hi, s),
(uint64_t)ROW_RD_I64(row, ly->off_sum, s)) / nn
: (double)ROW_RD_I64(row, ly->off_sum, s) / nn;
if (agg_affine[a].enabled) v += agg_affine[a].bias_f64;
break;
case OP_MIN:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? ROW_RD_F64(row, ly->off_min, s)
: (double)ROW_RD_I64(row, ly->off_min, s);
break;
case OP_MAX:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? ROW_RD_F64(row, ly->off_max, s)
: (double)ROW_RD_I64(row, ly->off_max, s);
break;
case OP_FIRST: case OP_LAST:
if (nn == 0) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
v = is_f64 ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
break;
case OP_VAR: case OP_VAR_POP:
case OP_STDDEV: case OP_STDDEV_POP: {
bool insuf = (agg_op == OP_VAR || agg_op == OP_STDDEV) ? nn <= 1 : nn <= 0;
if (insuf) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
double sum_val = is_f64 ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
double sq_val = ly->off_sumsq ? ROW_RD_F64(row, ly->off_sumsq, s) : 0.0;
double mean = sum_val / nn;
double var_pop = sq_val / nn - mean * mean;
if (var_pop < 0) var_pop = 0;
if (agg_op == OP_VAR_POP) v = var_pop;
else if (agg_op == OP_VAR) v = var_pop * nn / (nn - 1);
else if (agg_op == OP_STDDEV_POP) v = sqrt(var_pop);
else v = sqrt(var_pop * nn / (nn - 1));
break;
}
case OP_PEARSON_CORR:
case OP_COV:
case OP_SCOV:
case OP_WSUM:
case OP_WAVG: {
if ((agg_op == OP_PEARSON_CORR || agg_op == OP_SCOV) && nn < 2) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
if ((agg_op == OP_COV || agg_op == OP_WAVG) && nn < 1) { v = NULL_F64; ray_vec_set_null(new_col, gi, true); break; }
double sx = is_f64 ? ROW_RD_F64(row, ly->off_sum, s)
: (double)ROW_RD_I64(row, ly->off_sum, s);
double sxx = ly->off_sumsq ? ROW_RD_F64(row, ly->off_sumsq, s) : 0.0;
double sy = ly->off_sum_y ? ROW_RD_F64(row, ly->off_sum_y, s) : 0.0;
double syy = ly->off_sumsq_y ? ROW_RD_F64(row, ly->off_sumsq_y, s) : 0.0;
double sxy = ly->off_sumxy ? ROW_RD_F64(row, ly->off_sumxy, s) : 0.0;
double dn = (double)nn;
double num = dn * sxy - sx * sy;
double dx = dn * sxx - sx * sx;
double dy = dn * syy - sy * sy;
if (agg_op == OP_WSUM) { v = sxy; break; }
if (agg_op == OP_WAVG) { if (sx == 0.0) { v = NULL_F64; break; } v = sxy / sx; break; }
if (agg_op == OP_COV) { v = (sxy - sx * sy / dn) / dn; break; }
if (agg_op == OP_SCOV) { v = (sxy - sx * sy / dn) / (dn - 1.0); break; }
if (dx <= 0.0 || dy <= 0.0) { v = NULL_F64; break; }
v = num / sqrt(dx * dy);
break;
}
default: v = 0.0; break;
}
((double*)ray_data(new_col))[gi] = ray_f64_fin(v);
} else {
int64_t v;
int64_t int_null = agg_int_null_sentinel_for(out_type);
switch (agg_op) {
case OP_SUM:
v = ROW_RD_I64(row, ly->off_sum, s);
if (agg_affine[a].enabled)
v = wrap_add_i64(v, wrap_mul_i64(agg_affine[a].bias_i64, cnt));
break;
case OP_PROD:
if (nn == 0) { v = int_null; ray_vec_set_null(new_col, gi, true); break; }
v = ROW_RD_I64(row, ly->off_sum, s); break;
case OP_COUNT: v = cnt; break;
case OP_ALL: v = (ROW_RD_I64(row, ly->off_sum, s) == nn) ? 1 : 0; break;
case OP_ANY: v = (ROW_RD_I64(row, ly->off_sum, s) > 0) ? 1 : 0; break;
case OP_MIN:
if (nn == 0) { v = int_null; ray_vec_set_null(new_col, gi, true); break; }
v = ROW_RD_I64(row, ly->off_min, s); break;
case OP_MAX:
if (nn == 0) { v = int_null; ray_vec_set_null(new_col, gi, true); break; }
v = ROW_RD_I64(row, ly->off_max, s); break;
case OP_FIRST: case OP_LAST:
if (nn == 0) { v = int_null; ray_vec_set_null(new_col, gi, true); break; }
v = ROW_RD_I64(row, ly->off_sum, s); break;
default: v = 0; break;
}
topk_write_i64(ray_data(new_col), gi, v,
ray_sym_elem_size(out_type, new_col->attrs));
}
}
med_attach:;
ray_op_ext_t* agg_ext = find_ext(g, ext->agg_ins[a]);
int64_t name_id;
if (agg_ext && agg_ext->base.opcode == OP_SCAN) {
name_id = agg_result_col_name(agg_ext->sym, agg_op);
} else {
char nbuf[32];
int np = 0;
nbuf[np++] = '_'; nbuf[np++] = 'e';
{ uint8_t v = a; char dig[3]; int nd = 0;
do { dig[nd++] = (char)('0' + v % 10); v /= 10; } while (v);
while (nd--) nbuf[np++] = dig[nd]; }
const char* nsfx = "";
size_t nslen = 0;
switch (agg_op) {
case OP_SUM: nsfx = "_sum"; nslen = 4; break;
case OP_PROD: nsfx = "_prod"; nslen = 5; break;
case OP_ALL: nsfx = "_all"; nslen = 4; break;
case OP_ANY: nsfx = "_any"; nslen = 4; break;
case OP_COUNT: nsfx = "_count"; nslen = 6; break;
case OP_AVG: nsfx = "_mean"; nslen = 5; break;
case OP_MIN: nsfx = "_min"; nslen = 4; break;
case OP_MAX: nsfx = "_max"; nslen = 4; break;
case OP_FIRST: nsfx = "_first"; nslen = 6; break;
case OP_LAST: nsfx = "_last"; nslen = 5; break;
case OP_STDDEV: nsfx = "_stddev"; nslen = 7; break;
case OP_STDDEV_POP: nsfx = "_stddev_pop"; nslen = 11; break;
case OP_VAR: nsfx = "_var"; nslen = 4; break;
case OP_VAR_POP: nsfx = "_var_pop"; nslen = 8; break;
case OP_MEDIAN: nsfx = "_median"; nslen = 7; break;
case OP_QUANTILE: nsfx = "_quantile"; nslen = 9; break;
case OP_MODE: nsfx = "_mode"; nslen = 5; break;
case OP_COV: nsfx = "_cov"; nslen = 4; break;
case OP_SCOV: nsfx = "_scov"; nslen = 5; break;
case OP_WSUM: nsfx = "_wsum"; nslen = 5; break;
case OP_WAVG: nsfx = "_wavg"; nslen = 5; break;
case OP_TOP_N: nsfx = "_top"; nslen = 4; break;
case OP_BOT_N: nsfx = "_bot"; nslen = 4; break;
}
memcpy(nbuf + np, nsfx, nslen);
name_id = ray_sym_intern(nbuf, (size_t)np + nslen);
}
if (new_col->type == RAY_F64) grp_finalize_nulls(new_col);
result = ray_table_add_col(result, name_id, new_col);
ray_release(new_col);
}
if (med_out) {
for (uint32_t a = 0; a < n_aggs; a++)
if (med_out[a] && !RAY_IS_ERR(med_out[a])) ray_release(med_out[a]);
scratch_free(med_hdr);
}
}
cleanup:
if (final_ht == &single_ht) {
group_ht_free(&single_ht);
}
if (top_ht_ready) {
group_ht_free(&top_ht);
}
if (radix_bufs) {
size_t n_bufs = (size_t)n_total * radix_n_parts;
for (size_t i = 0; i < n_bufs; i++) scratch_free(radix_bufs[i]._hdr);
scratch_free(radix_bufs_hdr);
}
if (part_hts) {
for (uint32_t p = 0; p < radix_n_parts; p++) {
if (part_hts[p].rows) group_ht_free(&part_hts[p]);
}
scratch_free(part_hts_hdr);
}
scratch_free(part_offsets_hdr);
scratch_free(group_out_hdr);
scratch_free(topn_fuse_hdr);
scratch_free(topn_fuse_cnt_hdr);
ght_layout_free(&ght_layout);
for (uint32_t a = 0; a < n_aggs; a++)
{ if (agg_owned[a] && agg_vecs[a]) ray_release(agg_vecs[a]); if (agg_owned2[a] && agg_vecs2[a]) ray_release(agg_vecs2[a]); }
for (uint32_t k = 0; k < n_keys; k++)
if (key_owned[k] && key_vecs[k]) ray_release(key_vecs[k]);
if (match_idx_block) { ray_release(match_idx_block); } scratch_free(vla_hdr);
return result;
}
static inline int64_t grp_sorted_lower_bound(const int64_t* a, int64_t n,
int64_t key) {
int64_t lo = 0, hi = n;
while (lo < hi) {
int64_t mid = lo + ((hi - lo) >> 1);
if (a[mid] < key) lo = mid + 1;
else hi = mid;
}
return lo;
}
static ray_t* __attribute__((noinline))
exec_group_per_partition(ray_graph_t* g, ray_t* parted_tbl, ray_op_ext_t* ext,
int32_t n_parts, const int64_t* key_syms,
const int64_t* agg_syms, const uint8_t* key_is_expr,
const uint8_t* agg_is_expr, const int64_t* borrow_syms,
int n_borrow, int has_avg, int has_stddev,
int64_t group_limit) {
atomic_fetch_add_explicit(&ray_group_perpart_runs_ctr, 1,
memory_order_relaxed);
uint32_t n_keys = ext->n_keys;
uint32_t n_aggs = ext->n_aggs;
uint32_t n_avg = 0, n_std = 0;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop == OP_AVG) n_avg++;
else if (aop == OP_STDDEV || aop == OP_STDDEV_POP ||
aop == OP_VAR || aop == OP_VAR_POP) n_std++;
}
uint32_t part_n_aggs_max = n_aggs + n_avg + 2 * n_std;
size_t key_n = n_keys ? n_keys : 1;
size_t agg_n = n_aggs ? n_aggs : 1;
size_t pna_n = part_n_aggs_max ? part_n_aggs_max : 1;
ray_t* ret = NULL;
ray_t* saved_sel = g->selection;
ray_t* sel_ids_block = NULL;
const int64_t* sel_ids = NULL;
int64_t sel_total = 0;
ray_t* part_off_hdr = NULL;
int64_t* part_off = NULL;
ray_t* local_ids_hdr = NULL;
int64_t* local_ids = NULL;
ray_t *mc_sym_ids_hdr = NULL, *pk_syms_hdr = NULL,
*pkeys_hdr = NULL, *mkeys_hdr = NULL;
int64_t* mc_sym_ids = scratch_alloc(&mc_sym_ids_hdr, key_n * sizeof(int64_t));
int64_t* pk_syms = scratch_alloc(&pk_syms_hdr, key_n * sizeof(int64_t));
ray_op_t** pkeys = scratch_alloc(&pkeys_hdr, key_n * sizeof(ray_op_t*));
ray_op_t** mkeys = scratch_alloc(&mkeys_hdr, key_n * sizeof(ray_op_t*));
ray_t *avg_idx_hdr = NULL, *std_idx_hdr = NULL, *std_orig_op_hdr = NULL,
*std_sq_slot_hdr = NULL, *std_cnt_slot_hdr = NULL;
uint8_t* avg_idx = scratch_alloc(&avg_idx_hdr, agg_n * sizeof(uint8_t));
uint8_t* std_idx = scratch_alloc(&std_idx_hdr, agg_n * sizeof(uint8_t));
uint16_t* std_orig_op = scratch_alloc(&std_orig_op_hdr, agg_n * sizeof(uint16_t));
uint8_t* std_sq_slot = scratch_alloc(&std_sq_slot_hdr, agg_n * sizeof(uint8_t));
uint8_t* std_cnt_slot = scratch_alloc(&std_cnt_slot_hdr, agg_n * sizeof(uint8_t));
ray_t *part_ops_hdr = NULL, *merge_ops_hdr = NULL, *part_agg_syms_hdr = NULL,
*part_needs_sq_hdr = NULL, *agg_name_ids_hdr = NULL, *pagg_ins_hdr = NULL,
*magg_ins_hdr = NULL, *unique_agg_hdr = NULL;
uint16_t* part_ops = scratch_alloc(&part_ops_hdr, pna_n * sizeof(uint16_t));
uint16_t* merge_ops = scratch_alloc(&merge_ops_hdr, pna_n * sizeof(uint16_t));
int64_t* part_agg_syms = scratch_alloc(&part_agg_syms_hdr, pna_n * sizeof(int64_t));
int* part_needs_sq = scratch_alloc(&part_needs_sq_hdr, pna_n * sizeof(int));
int64_t* agg_name_ids = scratch_alloc(&agg_name_ids_hdr, pna_n * sizeof(int64_t));
ray_op_t** pagg_ins = scratch_alloc(&pagg_ins_hdr, pna_n * sizeof(ray_op_t*));
ray_op_t** magg_ins = scratch_alloc(&magg_ins_hdr, pna_n * sizeof(ray_op_t*));
int64_t* unique_agg = scratch_alloc(&unique_agg_hdr, pna_n * sizeof(int64_t));
if (!mc_sym_ids || !pk_syms || !pkeys || !mkeys || !avg_idx || !std_idx ||
!std_orig_op || !std_sq_slot || !std_cnt_slot || !part_ops ||
!merge_ops || !part_agg_syms || !part_needs_sq || !agg_name_ids ||
!pagg_ins || !magg_ins || !unique_agg)
goto pp_cleanup;
uint32_t n_mc_keys = 0;
uint32_t n_part_keys = 0;
for (uint32_t k = 0; k < n_keys; k++) {
ray_t* pcol = ray_table_get_col(parted_tbl, key_syms[k]);
if (pcol && pcol->type == RAY_MAPCOMMON) {
mc_sym_ids[n_mc_keys++] = key_syms[k];
} else {
pk_syms[n_part_keys++] = key_syms[k];
}
}
if (group_limit > 0 && n_part_keys == 0 && group_limit < n_parts && !saved_sel)
n_parts = (int32_t)group_limit;
if (saved_sel) {
ray_rowsel_t* sm = ray_rowsel_meta(saved_sel);
sel_total = sm->total_pass;
sel_ids_block = ray_rowsel_to_indices(saved_sel);
if (!sel_ids_block) goto pp_cleanup;
sel_ids = (const int64_t*)ray_data(sel_ids_block);
g->selection = NULL;
part_off = scratch_alloc(&part_off_hdr,
(size_t)(n_parts + 1) * sizeof(int64_t));
if (!part_off) goto pp_cleanup;
int64_t pncols = ray_table_ncols(parted_tbl);
ray_t* ref = NULL;
for (int64_t c = 0; c < pncols; c++) {
ray_t* col = ray_table_get_col_idx(parted_tbl, c);
if (col && (RAY_IS_PARTED(col->type) ||
col->type == RAY_MAPCOMMON)) { ref = col; break; }
}
if (!ref) goto pp_cleanup;
const int64_t* mc_counts = (ref->type == RAY_MAPCOMMON)
? (const int64_t*)ray_data(((ray_t**)ray_data(ref))[1]) : NULL;
ray_t** ref_segs = (ref->type == RAY_MAPCOMMON)
? NULL : (ray_t**)ray_data(ref);
int64_t cum = 0, max_len = 0;
for (int32_t p = 0; p < n_parts; p++) {
part_off[p] = cum;
int64_t len_p = mc_counts ? mc_counts[p]
: (ref_segs[p] ? ref_segs[p]->len : 0);
cum += len_p;
if (len_p > max_len) max_len = len_p;
}
part_off[n_parts] = cum;
local_ids = scratch_alloc(&local_ids_hdr,
(size_t)(max_len > 0 ? max_len : 1) * sizeof(int64_t));
if (!local_ids) goto pp_cleanup;
if (ray_interrupted()) { ret = ray_error("cancel", NULL); goto pp_cleanup; }
}
n_avg = 0;
n_std = 0;
uint32_t part_n_aggs = n_aggs;
for (uint32_t a = 0; a < n_aggs; a++) {
uint16_t aop = ext->agg_ops[a];
if (aop == OP_AVG) {
part_ops[a] = OP_SUM;
avg_idx[n_avg++] = a;
} else if (aop == OP_STDDEV || aop == OP_STDDEV_POP ||
aop == OP_VAR || aop == OP_VAR_POP) {
part_ops[a] = OP_SUM;
std_orig_op[n_std] = aop;
std_idx[n_std++] = a;
} else {
part_ops[a] = aop;
}
}
for (uint32_t i = 0; i < n_std; i++) {
std_sq_slot[i] = part_n_aggs;
part_ops[part_n_aggs++] = OP_SUM;
}
for (uint32_t i = 0; i < n_avg; i++)
part_ops[part_n_aggs++] = OP_COUNT;
for (uint32_t i = 0; i < n_std; i++) {
std_cnt_slot[i] = part_n_aggs;
part_ops[part_n_aggs++] = OP_COUNT;
}
for (uint32_t a = 0; a < part_n_aggs; a++) {
merge_ops[a] = part_ops[a];
if (merge_ops[a] == OP_COUNT) merge_ops[a] = OP_SUM;
}
memset(part_needs_sq, 0, pna_n * sizeof(int));
for (uint32_t a = 0; a < n_aggs; a++)
part_agg_syms[a] = agg_syms[a];
for (uint32_t i = 0; i < n_std; i++) {
part_agg_syms[std_sq_slot[i]] = agg_syms[std_idx[i]];
part_needs_sq[std_sq_slot[i]] = 1;
}
for (uint32_t i = 0; i < n_avg; i++)
part_agg_syms[n_aggs + n_std + i] = agg_syms[avg_idx[i]];
for (uint32_t i = 0; i < n_std; i++)
part_agg_syms[std_cnt_slot[i]] = agg_syms[std_idx[i]];
#define MERGE_BATCH 8
int agg_names_captured = 0;
ray_t* running = NULL;
ray_t* merge_tbl = NULL;
for (int32_t batch_start = 0; batch_start < n_parts;
batch_start += MERGE_BATCH) {
int32_t batch_end = batch_start + MERGE_BATCH;
if (batch_end > n_parts) batch_end = n_parts;
int32_t batch_n = batch_end - batch_start;
ray_t* bp[MERGE_BATCH];
memset(bp, 0, sizeof(bp));
int32_t ref_bi = -1;
for (int32_t bi = 0; bi < batch_n; bi++) {
int32_t p = batch_start + bi;
if (ray_interrupted()) {
ret = ray_error("cancel", NULL);
goto batch_fail;
}
ray_progress_update("group", "per-partition aggregate",
(uint64_t)p, (uint64_t)n_parts);
int64_t sel_off = 0, sel_len = 0, sel_lo = 0, sel_cnt = -1;
if (saved_sel) {
sel_off = part_off[p];
sel_len = part_off[p + 1] - sel_off;
sel_lo = grp_sorted_lower_bound(sel_ids, sel_total, sel_off);
int64_t sel_hi = grp_sorted_lower_bound(sel_ids, sel_total,
sel_off + sel_len);
sel_cnt = sel_hi - sel_lo;
if (sel_cnt == 0 && (running || ref_bi >= 0)) {
bp[bi] = NULL;
continue;
}
}
(void)unique_agg;
ray_t* sub = ray_table_new((int64_t)(n_borrow + n_keys + n_aggs));
if (!sub || RAY_IS_ERR(sub)) goto batch_fail;
int64_t part_nrows = -1;
for (int j = 0; j < n_borrow; j++) {
ray_t* pcol = ray_table_get_col(parted_tbl, borrow_syms[j]);
if (!pcol || !RAY_IS_PARTED(pcol->type)) {
ray_release(sub); goto batch_fail;
}
ray_t* seg = ((ray_t**)ray_data(pcol))[p];
if (!seg) { ray_release(sub); goto batch_fail; }
if (part_nrows < 0) part_nrows = seg->len;
ray_retain(seg);
sub = ray_table_add_col(sub, borrow_syms[j], seg);
ray_release(seg);
}
for (uint32_t k = 0; k < n_keys; k++) {
if (!key_is_expr[k]) continue;
ray_t* saved = g->table;
g->table = sub;
ray_t* v = exec_node(g, op_node(g, ext->keys[k]));
g->table = saved;
if (!v || RAY_IS_ERR(v) || !ray_is_vec(v) ||
(part_nrows >= 0 && v->len != part_nrows)) {
if (v && !RAY_IS_ERR(v)) ray_release(v);
ray_release(sub); goto batch_fail;
}
sub = ray_table_add_col(sub, key_syms[k], v);
ray_release(v);
}
for (uint32_t a = 0; a < n_aggs; a++) {
if (!agg_is_expr[a]) continue;
ray_t* saved = g->table;
g->table = sub;
ray_t* v = exec_node(g, op_node(g, ext->agg_ins[a]));
g->table = saved;
if (!v || RAY_IS_ERR(v) || !ray_is_vec(v) ||
(part_nrows >= 0 && v->len != part_nrows)) {
if (v && !RAY_IS_ERR(v)) ray_release(v);
ray_release(sub); goto batch_fail;
}
sub = ray_table_add_col(sub, agg_syms[a], v);
ray_release(v);
}
ray_graph_t* pg = ray_graph_new(sub);
if (!pg) { ray_release(sub); goto batch_fail; }
for (uint32_t k = 0; k < n_part_keys; k++) {
ray_t* sym_atom = ray_sym_str(pk_syms[k]);
pkeys[k] = ray_scan(pg, ray_str_ptr(sym_atom));
}
for (uint32_t a = 0; a < part_n_aggs; a++) {
ray_t* sym_atom = ray_sym_str(part_agg_syms[a]);
pagg_ins[a] = ray_scan(pg, ray_str_ptr(sym_atom));
}
for (uint32_t j = 0; j < n_std; j++) {
uint8_t sq = std_sq_slot[j];
ray_op_t* x = pagg_ins[sq];
ray_op_t* xf = (x->out_type == RAY_F64) ? x : ray_cast(pg, x, RAY_F64);
pagg_ins[sq] = ray_mul(pg, xf, xf);
}
ray_op_t* proot = ray_group(pg, pkeys, n_part_keys,
part_ops, pagg_ins, part_n_aggs);
proot = ray_optimize(pg, proot);
if (saved_sel) {
if (part_nrows >= 0 && part_nrows != sel_len) {
ray_graph_free(pg); ray_release(sub); goto batch_fail;
}
for (int64_t j = 0; j < sel_cnt; j++)
local_ids[j] = sel_ids[sel_lo + j] - sel_off;
ray_t* lsel = ray_index_rowsel_from_ids(sel_len, local_ids,
sel_cnt);
if (!lsel) { ray_graph_free(pg); ray_release(sub); goto batch_fail; }
pg->selection = lsel;
}
bp[bi] = ray_execute(pg, proot);
ray_graph_free(pg);
ray_release(sub);
if (!bp[bi] || RAY_IS_ERR(bp[bi])) goto batch_fail;
if (ref_bi < 0) ref_bi = bi;
if (!agg_names_captured) {
for (uint32_t a = 0; a < part_n_aggs; a++)
agg_name_ids[a] = ray_table_col_name(
bp[bi], (int64_t)n_part_keys + a);
agg_names_captured = 1;
}
}
int64_t mrows = running ? ray_table_nrows(running) : 0;
for (int32_t i = 0; i < batch_n; i++)
mrows += ray_table_nrows(bp[i]);
if (merge_tbl) { ray_release(merge_tbl); merge_tbl = NULL; }
merge_tbl = ray_table_new((int64_t)(n_keys + part_n_aggs));
if (!merge_tbl || RAY_IS_ERR(merge_tbl)) {
merge_tbl = NULL; goto batch_fail;
}
for (uint32_t k = 0; k < n_keys; k++) {
if (ray_interrupted()) {
ret = ray_error("cancel", NULL);
goto batch_fail;
}
int is_mc = 0;
for (uint32_t m = 0; m < n_mc_keys; m++)
if (mc_sym_ids[m] == key_syms[k]) { is_mc = 1; break; }
ray_t* tref = NULL;
if (running) {
tref = ray_table_get_col(running, key_syms[k]);
} else if (is_mc) {
ray_t* mc_col = ray_table_get_col(parted_tbl, key_syms[k]);
tref = ((ray_t**)ray_data(mc_col))[0];
} else {
tref = ray_table_get_col(bp[ref_bi >= 0 ? ref_bi : 0], key_syms[k]);
}
if (!tref) goto batch_fail;
uint8_t out_attrs = tref->attrs;
if (tref->type == RAY_SYM) {
uint8_t w = tref->attrs & RAY_SYM_W_MASK;
if (running) {
ray_t* rc = ray_table_get_col(running, key_syms[k]);
if (rc && rc->type == RAY_SYM &&
(rc->attrs & RAY_SYM_W_MASK) > w)
w = rc->attrs & RAY_SYM_W_MASK;
}
for (int32_t i = 0; i < batch_n; i++) {
ray_t* pc = is_mc ? NULL
: ray_table_get_col(bp[i], key_syms[k]);
if (pc && pc->type == RAY_SYM &&
(pc->attrs & RAY_SYM_W_MASK) > w)
w = pc->attrs & RAY_SYM_W_MASK;
}
out_attrs = w;
}
ray_t* flat = typed_vec_new(tref->type, out_attrs, mrows);
if (!flat || RAY_IS_ERR(flat)) goto batch_fail;
if (flat->type == RAY_SYM)
ray_sym_vec_adopt_domain(flat, sym_domain_rep(tref));
flat->len = mrows;
char* out = (char*)ray_data(flat);
int64_t off = 0;
uint8_t src_nulls = 0;
if (running) {
ray_t* rc = ray_table_get_col(running, key_syms[k]);
if (rc && rc->len > 0) {
parted_copy_cells(out, flat->type, out_attrs, 0,
rc, 0, rc->len);
src_nulls |= rc->attrs & RAY_ATTR_HAS_NULLS;
off = rc->len;
}
}
for (int32_t i = 0; i < batch_n; i++) {
int64_t pnrows = ray_table_nrows(bp[i]);
if (is_mc) {
int32_t p = batch_start + i;
ray_t* mc_col = ray_table_get_col(parted_tbl, key_syms[k]);
ray_t* mc_kv = ((ray_t**)ray_data(mc_col))[0];
for (int64_t r = 0; r < pnrows; r++) {
if ((r & 0xFFFF) == 0 && ray_interrupted()) {
ray_release(flat);
ret = ray_error("cancel", NULL);
goto batch_fail;
}
parted_copy_cells(out, flat->type, out_attrs,
off + r, mc_kv, p, 1);
}
src_nulls |= mc_kv->attrs & RAY_ATTR_HAS_NULLS;
off += pnrows;
} else {
ray_t* pc = ray_table_get_col(bp[i], key_syms[k]);
if (pc && pc->len > 0) {
parted_copy_cells(out, flat->type, out_attrs, off,
pc, 0, pc->len);
src_nulls |= pc->attrs & RAY_ATTR_HAS_NULLS;
off += pc->len;
}
}
}
flat->attrs |= src_nulls;
merge_tbl = ray_table_add_col(merge_tbl, key_syms[k], flat);
ray_release(flat);
}
for (uint32_t a = 0; a < part_n_aggs; a++) {
if (ray_interrupted()) {
ret = ray_error("cancel", NULL);
goto batch_fail;
}
ray_t* tref = running
? ray_table_get_col_idx(running, (int64_t)n_keys + a)
: ray_table_get_col_idx(bp[ref_bi >= 0 ? ref_bi : 0],
(int64_t)n_part_keys + a);
if (!tref) goto batch_fail;
uint8_t out_attrs = tref->attrs;
if (tref->type == RAY_SYM) {
uint8_t w = tref->attrs & RAY_SYM_W_MASK;
for (int32_t i = 0; i < batch_n; i++) {
ray_t* pc = ray_table_get_col_idx(bp[i],
(int64_t)n_part_keys + a);
if (pc && pc->type == RAY_SYM &&
(pc->attrs & RAY_SYM_W_MASK) > w)
w = pc->attrs & RAY_SYM_W_MASK;
}
out_attrs = w;
}
ray_t* flat = typed_vec_new(tref->type, out_attrs, mrows);
if (!flat || RAY_IS_ERR(flat)) goto batch_fail;
if (flat->type == RAY_SYM)
ray_sym_vec_adopt_domain(flat, sym_domain_rep(tref));
flat->len = mrows;
char* out = (char*)ray_data(flat);
int64_t off = 0;
uint8_t src_nulls = 0;
if (running) {
ray_t* rc = ray_table_get_col_idx(running, (int64_t)n_keys + a);
if (rc && rc->len > 0) {
parted_copy_cells(out, flat->type, out_attrs, 0,
rc, 0, rc->len);
src_nulls |= rc->attrs & RAY_ATTR_HAS_NULLS;
off = rc->len;
}
}
for (int32_t i = 0; i < batch_n; i++) {
ray_t* pc = ray_table_get_col_idx(bp[i],
(int64_t)n_part_keys + a);
if (pc && pc->len > 0) {
parted_copy_cells(out, flat->type, out_attrs, off,
pc, 0, pc->len);
src_nulls |= pc->attrs & RAY_ATTR_HAS_NULLS;
off += pc->len;
}
}
flat->attrs |= src_nulls;
merge_tbl = ray_table_add_col(merge_tbl, agg_name_ids[a], flat);
ray_release(flat);
}
for (int32_t i = 0; i < batch_n; i++) {
ray_release(bp[i]);
bp[i] = NULL;
}
ray_graph_t* mg = ray_graph_new(merge_tbl);
if (!mg) goto batch_fail;
for (uint32_t k = 0; k < n_keys; k++) {
ray_t* sym_atom = ray_sym_str(key_syms[k]);
mkeys[k] = ray_scan(mg, ray_str_ptr(sym_atom));
}
for (uint32_t a = 0; a < part_n_aggs; a++) {
ray_t* agg_name = ray_sym_str(agg_name_ids[a]);
magg_ins[a] = ray_scan(mg, ray_str_ptr(agg_name));
}
ray_op_t* mroot = ray_group(mg, mkeys, n_keys,
merge_ops, magg_ins, part_n_aggs);
mroot = ray_optimize(mg, mroot);
ray_t* new_running = ray_execute(mg, mroot);
ray_graph_free(mg);
if (running) ray_release(running);
running = new_running;
if (!running || RAY_IS_ERR(running)) {
ray_release(merge_tbl);
goto pp_cleanup;
}
for (uint32_t a = 0; a < part_n_aggs; a++)
ray_table_set_col_name(running, (int64_t)n_keys + a, agg_name_ids[a]);
continue;
batch_fail:
for (int32_t i = 0; i < batch_n; i++)
if (bp[i]) ray_release(bp[i]);
if (running) ray_release(running);
if (merge_tbl) ray_release(merge_tbl);
goto pp_cleanup;
}
ray_progress_update("group", "per-partition aggregate",
(uint64_t)n_parts, (uint64_t)n_parts);
ray_t* result = running;
if (!result || RAY_IS_ERR(result)) {
if (merge_tbl) ray_release(merge_tbl);
goto pp_cleanup;
}
int64_t rncols = ray_table_ncols(result);
if (has_avg || has_stddev) {
ray_t* trimmed = ray_table_new((int64_t)(n_keys + n_aggs));
if (!trimmed || RAY_IS_ERR(trimmed)) {
ray_release(result);
if (merge_tbl) ray_release(merge_tbl);
goto pp_cleanup;
}
for (int64_t c = 0; c < (int64_t)(n_keys + n_aggs) && c < rncols; c++) {
int64_t nm = ray_table_col_name(result, c);
int is_avg_slot = 0, is_std_slot = 0;
uint8_t avg_i = 0, std_i = 0;
if (c >= n_keys) {
uint8_t a = (uint8_t)(c - n_keys);
for (uint32_t j = 0; j < n_avg; j++) {
if (avg_idx[j] == a) { is_avg_slot = 1; avg_i = j; break; }
}
for (uint32_t j = 0; j < n_std; j++) {
if (std_idx[j] == a) { is_std_slot = 1; std_i = j; break; }
}
}
if (is_avg_slot) {
int64_t sum_ci = c;
int64_t cnt_ci = (int64_t)n_keys + n_aggs + n_std + avg_i;
ray_t* sum_col = ray_table_get_col_idx(result, sum_ci);
ray_t* cnt_col = (cnt_ci < rncols) ? ray_table_get_col_idx(result, cnt_ci) : NULL;
if (!sum_col || !cnt_col) {
if (sum_col) {
ray_retain(sum_col);
trimmed = ray_table_add_col(trimmed, nm, sum_col);
ray_release(sum_col);
}
continue;
}
int64_t nrows = sum_col->len;
ray_t* avg_col = ray_vec_new(RAY_F64, nrows);
if (!avg_col || RAY_IS_ERR(avg_col)) {
ray_release(trimmed); ray_release(result);
if (merge_tbl) ray_release(merge_tbl);
goto pp_cleanup;
}
avg_col->len = nrows;
double* out = (double*)ray_data(avg_col);
if (sum_col->type == RAY_F64) {
const double* sv = (const double*)ray_data(sum_col);
const int64_t* cv = (const int64_t*)ray_data(cnt_col);
for (int64_t r = 0; r < nrows; r++)
out[r] = cv[r] > 0 ? ray_f64_fin(sv[r] / (double)cv[r]) : 0.0;
} else {
const int64_t* sv = (const int64_t*)ray_data(sum_col);
const int64_t* cv = (const int64_t*)ray_data(cnt_col);
for (int64_t r = 0; r < nrows; r++)
out[r] = cv[r] > 0 ? ray_f64_fin((double)sv[r] / (double)cv[r]) : 0.0;
}
grp_finalize_nulls(avg_col);
trimmed = ray_table_add_col(trimmed, nm, avg_col);
ray_release(avg_col);
} else if (is_std_slot) {
int64_t sum_ci = c;
int64_t sq_ci = (int64_t)n_keys + std_sq_slot[std_i];
int64_t cnt_ci = (int64_t)n_keys + std_cnt_slot[std_i];
ray_t* sum_col = ray_table_get_col_idx(result, sum_ci);
ray_t* sq_col = (sq_ci < rncols) ? ray_table_get_col_idx(result, sq_ci) : NULL;
ray_t* cnt_col = (cnt_ci < rncols) ? ray_table_get_col_idx(result, cnt_ci) : NULL;
if (!sum_col || !sq_col || !cnt_col) {
if (sum_col) {
ray_retain(sum_col);
trimmed = ray_table_add_col(trimmed, nm, sum_col);
ray_release(sum_col);
}
continue;
}
int64_t nrows = sum_col->len;
ray_t* out_col = ray_vec_new(RAY_F64, nrows);
if (!out_col || RAY_IS_ERR(out_col)) {
ray_release(trimmed); ray_release(result);
if (merge_tbl) ray_release(merge_tbl);
goto pp_cleanup;
}
out_col->len = nrows;
double* out = (double*)ray_data(out_col);
uint16_t orig_op = std_orig_op[std_i];
const double* sq = (const double*)ray_data(sq_col);
const int64_t* cv = (const int64_t*)ray_data(cnt_col);
if (sum_col->type == RAY_F64) {
const double* sv = (const double*)ray_data(sum_col);
for (int64_t r = 0; r < nrows; r++) {
double n = (double)cv[r];
if (n <= 0) { out[r] = NULL_F64; ray_vec_set_null(out_col, r, true); continue; }
double mean = sv[r] / n;
double var_pop = sq[r] / n - mean * mean;
if (var_pop < 0) var_pop = 0;
bool insuf = (orig_op == OP_VAR || orig_op == OP_STDDEV) && n <= 1;
if (insuf) { out[r] = NULL_F64; ray_vec_set_null(out_col, r, true); continue; }
if (orig_op == OP_VAR_POP) out[r] = ray_f64_fin(var_pop);
else if (orig_op == OP_VAR) out[r] = ray_f64_fin(var_pop * n / (n - 1));
else if (orig_op == OP_STDDEV_POP) out[r] = ray_f64_fin(sqrt(var_pop));
else out[r] = ray_f64_fin(sqrt(var_pop * n / (n - 1)));
}
} else {
const int64_t* sv = (const int64_t*)ray_data(sum_col);
for (int64_t r = 0; r < nrows; r++) {
double n = (double)cv[r];
if (n <= 0) { out[r] = NULL_F64; ray_vec_set_null(out_col, r, true); continue; }
double mean = (double)sv[r] / n;
double var_pop = sq[r] / n - mean * mean;
if (var_pop < 0) var_pop = 0;
bool insuf = (orig_op == OP_VAR || orig_op == OP_STDDEV) && n <= 1;
if (insuf) { out[r] = NULL_F64; ray_vec_set_null(out_col, r, true); continue; }
if (orig_op == OP_VAR_POP) out[r] = ray_f64_fin(var_pop);
else if (orig_op == OP_VAR) out[r] = ray_f64_fin(var_pop * n / (n - 1));
else if (orig_op == OP_STDDEV_POP) out[r] = ray_f64_fin(sqrt(var_pop));
else out[r] = ray_f64_fin(sqrt(var_pop * n / (n - 1)));
}
}
grp_finalize_nulls(out_col);
trimmed = ray_table_add_col(trimmed, nm, out_col);
ray_release(out_col);
} else {
ray_t* col = ray_table_get_col_idx(result, c);
if (col) {
ray_retain(col);
trimmed = ray_table_add_col(trimmed, nm, col);
ray_release(col);
}
}
}
ray_release(result);
result = trimmed;
rncols = ray_table_ncols(result);
}
for (uint32_t k = 0; k < n_keys && (int64_t)k < rncols; k++) {
ray_op_ext_t* ke = find_ext(g, ext->keys[k]);
ray_table_set_col_name(result, (int64_t)k, ke ? ke->sym : (int64_t)k);
}
for (uint32_t a = 0; a < n_aggs && (int64_t)(n_keys + a) < rncols; a++)
ray_table_set_col_name(result, (int64_t)n_keys + a, agg_name_ids[a]);
if (merge_tbl) ray_release(merge_tbl);
ret = result;
pp_cleanup:
scratch_free(mc_sym_ids_hdr);
scratch_free(pk_syms_hdr);
scratch_free(pkeys_hdr);
scratch_free(mkeys_hdr);
scratch_free(avg_idx_hdr);
scratch_free(std_idx_hdr);
scratch_free(std_orig_op_hdr);
scratch_free(std_sq_slot_hdr);
scratch_free(std_cnt_slot_hdr);
scratch_free(part_ops_hdr);
scratch_free(merge_ops_hdr);
scratch_free(part_agg_syms_hdr);
scratch_free(part_needs_sq_hdr);
scratch_free(agg_name_ids_hdr);
scratch_free(pagg_ins_hdr);
scratch_free(magg_ins_hdr);
scratch_free(unique_agg_hdr);
if (sel_ids_block) ray_release(sel_ids_block);
scratch_free(part_off_hdr);
scratch_free(local_ids_hdr);
g->selection = saved_sel;
return ret;
}
static void pivot_ingest_sequential(pivot_ingest_t* out, const ght_layout_t* ly,
void** key_data, int8_t* key_types,
uint8_t* key_attrs, ray_t** key_vecs,
ray_t** agg_vecs, int64_t n_scan,
group_ht_t* scratch_ht) {
(void)key_data;
out->part_hts = scratch_ht;
out->n_parts = 1;
out->row_stride = ly->row_stride;
group_rows_range(scratch_ht, key_data, key_types, key_attrs, key_vecs,
agg_vecs, NULL, NULL, NULL, 0, n_scan, NULL);
out->total_grps = scratch_ht->grp_count;
out->part_offsets[0] = 0;
out->part_offsets[1] = scratch_ht->grp_count;
out->part_hts = scratch_ht;
}
bool pivot_ingest_run(pivot_ingest_t* out,
const ght_layout_t* ly,
void** key_data, int8_t* key_types, uint8_t* key_attrs,
ray_t** key_vecs, ray_t** agg_vecs,
int64_t n_scan) {
memset(out, 0, sizeof(*out));
out->row_stride = ly->row_stride;
ray_pool_t* pool = ray_pool_get();
uint32_t n_total = pool ? ray_pool_total_workers(pool) : 1;
bool parallel_ok = (pool && n_scan >= RAY_PARALLEL_THRESHOLD && n_total > 1);
uint32_t n_parts = parallel_ok
? group_radix_part_count(n_total, n_scan) : 1;
out->part_offsets = (uint32_t*)scratch_alloc(&out->_offsets_hdr,
((size_t)n_parts + 1) * sizeof(uint32_t));
if (!out->part_offsets) return false;
uint32_t n_keys = ly->n_keys;
if (!parallel_ok) {
group_ht_t* seq = (group_ht_t*)scratch_calloc(&out->_part_hts_hdr,
sizeof(group_ht_t));
if (!seq) return false;
out->part_hts = seq;
out->n_parts = 1;
if (!group_ht_init(seq, 2, ly)) return false;
pivot_ingest_sequential(out, ly, key_data, key_types, key_attrs,
key_vecs, agg_vecs, n_scan, seq);
if (seq->oom) return false;
return true;
}
size_t n_bufs = (size_t)n_total * n_parts;
out->_n_bufs = n_bufs;
radix_buf_t* radix_bufs = (radix_buf_t*)scratch_calloc(&out->_radix_bufs_hdr,
n_bufs * sizeof(radix_buf_t));
if (!radix_bufs) return false;
out->_radix_bufs = radix_bufs;
uint8_t p1_nullable = 0;
for (uint32_t k = 0; k < n_keys; k++)
if (ray_key_may_be_null(key_vecs[k])) { p1_nullable = 1; break; }
radix_phase1_ctx_t p1ctx = {
.key_data = key_data,
.key_types = key_types,
.key_attrs = key_attrs,
.key_vecs = key_vecs,
.nullable_mask = p1_nullable,
.agg_vecs = agg_vecs,
.agg_vecs2 = NULL,
.n_workers = n_total,
.n_parts = n_parts,
.bufs = radix_bufs,
.match_idx = NULL,
.buf_prime = (uint32_t)((uint64_t)(n_scan > 0 ? n_scan : 1)
/ ((uint64_t)n_total * n_parts) * 5 / 4 + 8),
};
ght_layout_copy(&p1ctx.layout, ly);
ray_t* p1_kp_hdr = NULL;
if (ly->any_wide_key) {
p1ctx.key_pool = (const void**)scratch_alloc(&p1_kp_hdr,
(size_t)n_keys * sizeof(void*));
if (!p1ctx.key_pool) return false;
derive_key_pool(ly, key_vecs, p1ctx.key_pool);
}
ray_pool_dispatch(pool, radix_phase1_fn, &p1ctx, n_scan);
scratch_free(p1_kp_hdr);
if (ray_interrupted()) return true;
ray_progress_update(NULL, "hash-partition", (uint64_t)n_scan, (uint64_t)n_scan);
for (size_t i = 0; i < n_bufs; i++)
if (radix_bufs[i].oom) return false;
group_ht_t* part_hts = (group_ht_t*)scratch_calloc(&out->_part_hts_hdr,
(size_t)n_parts * sizeof(group_ht_t));
if (!part_hts) return false;
out->part_hts = part_hts;
out->n_parts = n_parts;
radix_phase2_ctx_t p2ctx = {
.key_types = key_types,
.n_keys = n_keys,
.n_workers = n_total,
.n_parts = n_parts,
.bufs = radix_bufs,
.part_hts = part_hts,
.key_data = key_data,
};
ght_layout_copy(&p2ctx.layout, ly);
ray_t* p2_kp_hdr = NULL;
if (ly->any_wide_key) {
p2ctx.key_pool = (const void**)scratch_alloc(&p2_kp_hdr,
(size_t)n_keys * sizeof(void*));
if (!p2ctx.key_pool) return false;
derive_key_pool(ly, key_vecs, p2ctx.key_pool);
}
ray_pool_dispatch_n(pool, radix_phase2_fn, &p2ctx, n_parts);
scratch_free(p2_kp_hdr);
if (ray_interrupted()) return true;
ray_progress_update(NULL, "per-partition aggregate", n_parts, n_parts);
for (uint32_t p = 0; p < n_parts; p++) {
if (part_hts[p].oom) return false;
if (part_hts[p].rows) continue;
uint32_t pcount = 0;
for (uint32_t w = 0; w < n_total; w++)
pcount += radix_bufs[(size_t)w * n_parts + p].count;
if (pcount) return false;
}
out->part_offsets[0] = 0;
for (uint32_t p = 0; p < n_parts; p++)
out->part_offsets[p + 1] = out->part_offsets[p] + part_hts[p].grp_count;
out->total_grps = out->part_offsets[n_parts];
return true;
}
void pivot_ingest_free(pivot_ingest_t* out) {
if (!out) return;
if (out->part_hts) {
for (uint32_t p = 0; p < out->n_parts; p++) {
if (out->part_hts[p].rows || out->part_hts[p].slots)
group_ht_free(&out->part_hts[p]);
}
scratch_free(out->_part_hts_hdr);
}
if (out->_radix_bufs) {
radix_buf_t* bufs = (radix_buf_t*)out->_radix_bufs;
for (size_t i = 0; i < out->_n_bufs; i++) scratch_free(bufs[i]._hdr);
scratch_free(out->_radix_bufs_hdr);
}
scratch_free(out->_offsets_hdr);
memset(out, 0, sizeof(*out));
}