rayforce-sys 1.2.0

/*
 *   Copyright (c) 2025-2026 Anton Kundenko <singaraiona@gmail.com>
 *   All rights reserved.

 *   Permission is hereby granted, free of charge, to any person obtaining a copy
 *   of this software and associated documentation files (the "Software"), to deal
 *   in the Software without restriction, including without limitation the rights
 *   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 *   copies of the Software, and to permit persons to whom the Software is
 *   furnished to do so, subject to the following conditions:

 *   The above copyright notice and this permission notice shall be included in all
 *   copies or substantial portions of the Software.

 *   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 *   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 *   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 *   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 *   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 *   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 *   SOFTWARE.
 */

/*
 * Fused filter + top-K.  One pass over the source rows: predicate eval
 * per morsel into stack bits[], then for each passing row a bounded-heap
 * insert that keeps only the K best by the sort key.  No intermediate
 * filtered-table materialization (the FILTER → SORT_TAKE pipeline writes
 * a 2.5M-row temp table; this skips it).
 *
 * Targets the ORDER BY + LIMIT cluster: previously the engine
 * materialised a full filtered table and sorted it before applying
 * TAKE; the heap-bounded merge here keeps only K rows in flight.
 *
 * Single sort key for now.  SYM keys go through the comparator path
 * (dict-id → string → memcmp).  Numeric keys do direct value compare.
 *
 * Public surface: ray_fused_topk_supported() (gate probe) and
 * ray_fused_topk_select() (executor).  The planner in query.c calls
 * these before building the DAG; on a shape miss the caller falls back
 * to the unfused FILTER + apply_sort_take path.
 */

#include "ops/fused_topk.h"
#include "ops/fused_pred.h"
#include "ops/fused_group.h"   /* ray_fused_group_supported */
#include "ops/internal.h"
#include "lang/internal.h"
#include "core/pool.h"
#include "ops/idxop.h"      /* chunk-zone extrema: skip chunks that cannot enter the top-K */

#include <string.h>
#include <stdatomic.h>

/* Use the same predicate-shape detector as fused_group.  Single comparison
 * or AND of comparisons against literals on flat int/temporal/SYM columns. */
int ray_fused_topk_supported(ray_t* where_expr, ray_t* tbl) {
    return ray_fused_group_supported(where_expr, tbl);
}

/* ───── Heap entry ─────────────────────────────────────────────────────
 * Two parallel arrays per worker: heap_idx[] (source row indices) and
 * heap_key[] (encoded uint64 sort key for numeric, dict-id for SYM).
 * For numeric the encoded key is order-preserving so a single uint64
 * compare gives the right order.  For SYM we fall back to comparator
 * (sym_str / memcmp).  Stored as max-heap for ASC top-K (the WORST is
 * at the root and gets evicted first), or min-heap for DESC.
 * ──────────────────────────────────────────────────────────────────── */

#define FPK_MAX_KEYS 16

typedef struct {
    int8_t      type;
    uint8_t     attrs;
    uint8_t     esz;
    uint8_t     desc;        /* 0 = asc, 1 = desc */
    /* When the column carries nulls, fpk_cmp consults them before
     * reading the raw payload and places them at the end this flag
     * names.  It is filled from sort_nulls_first, the one place that
     * says where a null goes, so this path and the ordinary sort agree:
     * they must, since a parted top-k runs this path per partition and
     * is checked against the flat sort.  has_nulls is the compile-time
     * flag that gates the per-row probe. */
    uint8_t     has_nulls;
    uint8_t     nulls_first;
    /* SYM keys only: 1 when the column's domain is the runtime
     * singleton — cell ids then index the borrowed global snapshot
     * (ctx.sym_strings).  0 routes the compare through the column's
     * own domain (ray_sym_domain_str). */
    uint8_t     dom_runtime;
    /* SYM keys on a FILE domain: the mapped prefix pinned for the
     * dispatch — compares read raw bytes, no atom per symbol. */
    uint8_t     raw_ok;
    ray_sym_domain_raw_t raw;
    int64_t     sym;
    const void* base;
    ray_t*      col;         /* for ray_vec_is_null when has_nulls */
} fpk_keyspec_t;
static void fpk_unpin_keys(fpk_keyspec_t* keys, uint32_t n) {
    for (uint32_t i = 0; i < n; i++)
        if (keys[i].raw_ok) { ray_sym_domain_raw_unpin(ray_sym_vec_domain(keys[i].col)); keys[i].raw_ok = 0; }
}

typedef struct {
    fp_pred_t      pred;
    fpk_keyspec_t  keys[FPK_MAX_KEYS];
    /* n_keys mirrors n_sort_keys (now uint32_t at the API boundary), but
     * its VALUE stays permanently <= FPK_MAX_KEYS==16 — a routing/arity
     * admission gate below (n_sort_keys > FPK_MAX_KEYS -> NULL), not a
     * cut-3 cap slated to lift.  Widened anyway so no narrower carrier
     * stores a value that arrived through the wide parameter. */
    uint32_t       n_keys;
    int64_t        k;
    /* Per-worker bounded heap of row indices.  Comparator is invoked
     * via fpk_cmp; multi-key comparisons short-circuit on the first
     * key that breaks the tie.  No payload — keys are re-read from the
     * source columns at compare time. */
    int64_t*       heap_idx;       /* [nw * k] */
    int32_t*       heap_n;         /* [nw] */
    ray_t*         tbl;
    /* Lock-free snapshot of the GLOBAL SYM dict — populated once per
     * query for any runtime-domain RAY_SYM sort key.  ray_sym_str()
     * takes a global lock per call, which on the SYM heap-compare hot
     * path turns 28-way parallel into 12× *slower* due to lock
     * contention.  Borrowing the strings array once lets workers index
     * into it without the lock.  Non-runtime-domain keys bypass the
     * snapshot and resolve per compare via ray_sym_domain_str (locked;
     * a per-domain snapshot is a flip-time optimization if needed). */
    ray_t**        sym_strings;
    uint32_t       sym_count;
    _Atomic(uint32_t) oom;
    /* Chunk pruning on the first sort key (integer / temporal column with a
     * chunk-zone index): zmin/zmax per 1<<zlog2 rows, znull chunk null
     * bits.  `bound` is the best worst-key any full worker heap holds —
     * the final K-th key is at least as good — so a chunk whose extremum
     * cannot beat it is skipped without evaluating the predicate. */
    const int64_t* zmin;
    const int64_t* zmax;
    const uint8_t* znull;
    uint32_t       zn;
    uint8_t        zlog2;
    uint8_t        zdesc;
    uint8_t        znulls_better;   /* nulls sort ahead of values */
    /* Best-chunk-first visiting order: zorder[v] is the chunk to scan in
     * virtual slot v (chunks sorted by the extremum that matters for the
     * direction, null chunks first when nulls sort ahead).  Workers are
     * dispatched over the virtual row space zn << zlog2 and map each slot
     * back to its chunk, so the K-th key tightens on the most promising
     * chunks first and the rest are pruned — whatever the physical order
     * of the data.  NULL when there is no zone index. */
    const uint32_t* zorder;
    int64_t        nrows;
    /* written by every worker with a full heap: own cache line, away from
     * the read-only zone pointers above */
    _Alignas(64) _Atomic(int64_t) bound;
    _Atomic(uint8_t) bound_set;
} fpk_par_ctx_t;

static inline int64_t fpk_key_i64(const fpk_keyspec_t* ks, int64_t row) {
    switch (ks->esz) {
    case 1: return (int64_t)((const uint8_t*)ks->base)[row];
    case 2: return (int64_t)((const int16_t*)ks->base)[row];
    case 4: return (int64_t)((const int32_t*)ks->base)[row];
    default: return ((const int64_t*)ks->base)[row];
    }
}

/* A full heap's worst row publishes its first key as a pruning bound. */
static void fpk_publish_bound(fpk_par_ctx_t* c, int64_t worst_row) {
    if (!c->zmin) return;
    const fpk_keyspec_t* ks = &c->keys[0];
    if (ks->has_nulls && ray_vec_is_null(ks->col, worst_row)) return;
    int64_t v = fpk_key_i64(ks, worst_row);
    int64_t cur = atomic_load_explicit(&c->bound, memory_order_relaxed);
    bool set = atomic_load_explicit(&c->bound_set, memory_order_relaxed);
    for (;;) {
        bool better = !set || (c->zdesc ? v > cur : v < cur);
        if (!better) return;
        if (atomic_compare_exchange_weak_explicit(&c->bound, &cur, v,
                memory_order_relaxed, memory_order_relaxed)) {
            atomic_store_explicit(&c->bound_set, 1, memory_order_release);
            return;
        }
        set = true;
    }
}

/* Chunk `g` cannot hold a row that beats the bound. */
static inline bool fpk_chunk_pruned(const fpk_par_ctx_t* c, int64_t g) {
    if (!atomic_load_explicit(&c->bound_set, memory_order_acquire)) return false;
    if (g < 0 || g >= (int64_t)c->zn) return false;
    if (c->znulls_better && c->znull && (c->znull[g >> 3] & (1u << (g & 7)))) return false;
    int64_t b = atomic_load_explicit(&c->bound, memory_order_relaxed);
    return c->zdesc ? c->zmax[g] < b : c->zmin[g] > b;
}

/* Chunk `a` is more promising than chunk `b` for the first key's
 * direction: null chunks lead when nulls sort ahead, then the smaller
 * minimum (asc) or the larger maximum (desc); ties keep chunk order. */
static inline bool fpk_chunk_better(const fpk_par_ctx_t* c, uint32_t a, uint32_t b) {
    if (c->znulls_better && c->znull) {
        bool na = (c->znull[a >> 3] >> (a & 7)) & 1;
        bool nb = (c->znull[b >> 3] >> (b & 7)) & 1;
        if (na != nb) return na;
    }
    int64_t ka = c->zdesc ? c->zmax[a] : c->zmin[a];
    int64_t kb = c->zdesc ? c->zmax[b] : c->zmin[b];
    if (ka != kb) return c->zdesc ? ka > kb : ka < kb;
    return a < b;
}

/* Heap-sort the chunk ids in `ord` (initially 0..n-1) best first. */
static void fpk_order_chunks(const fpk_par_ctx_t* c, uint32_t* ord, uint32_t n) {
    /* max-heap on "worse": the root is the least promising chunk, so
     * popping it to the tail leaves the best chunk at ord[0]. */
    for (uint32_t i = n; i-- > 0;) {
        /* sift ord[i] down */
        uint32_t k = i;
        for (;;) {
            uint32_t l = 2 * k + 1, r = l + 1, w = k;
            if (l < n && fpk_chunk_better(c, ord[w], ord[l])) w = l;
            if (r < n && fpk_chunk_better(c, ord[w], ord[r])) w = r;
            if (w == k) break;
            uint32_t t = ord[k]; ord[k] = ord[w]; ord[w] = t;
            k = w;
        }
        if (i == 0) break;
    }
    for (uint32_t end = n; end > 1;) {
        end--;
        uint32_t t = ord[0]; ord[0] = ord[end]; ord[end] = t;
        uint32_t k = 0;
        for (;;) {
            uint32_t l = 2 * k + 1, r = l + 1, w = k;
            if (l < end && fpk_chunk_better(c, ord[w], ord[l])) w = l;
            if (r < end && fpk_chunk_better(c, ord[w], ord[r])) w = r;
            if (w == k) break;
            uint32_t u = ord[k]; ord[k] = ord[w]; ord[w] = u;
            k = w;
        }
    }
}

/* Compare two source rows by the multi-key sort spec.  Returns
 * "a is worse than b" sense: positive means evict-a-first in the
 * max-heap of K-best entries.  Short-circuits on first non-equal key.
 *
 * Numeric/temporal narrow widths use sign-aware reads so a stored -1
 * compares as -1 and not 65535 — same fix as the fused-group agg
 * read.  Signed/unsigned matches the column class: BOOL/U8/SYM-id
 * are unsigned; I16/I32/I64 and the temporals (DATE/TIME/TIMESTAMP)
 * are signed.
 *
 * Final tie-break is by source row index ascending — produces a
 * deterministic, source-order-preserving result that matches a
 * stable sort of the surviving rows.  Without the tie-break, two
 * runs of the same query against the same data could return
 * different rows for ties because morsel scheduling between
 * workers is non-deterministic. */
static inline int fpk_cmp(const fpk_par_ctx_t* c, int64_t row_a, int64_t row_b) {
    if (RAY_UNLIKELY(row_a == row_b)) return 0;
    for (uint32_t k = 0; k < c->n_keys; k++) {
        const fpk_keyspec_t* ks = &c->keys[k];
        int cmp = 0;
        /* Null-aware leg.  In the max-heap-of-K-best comparator, "worse"
         * → evicted first, so a null asked for first is better than any
         * non-null and a null asked for last is worse.  Both legs
         * short-circuit before the raw payload read. */
        if (ks->has_nulls) {
            bool a_null = ray_vec_is_null(ks->col, row_a);
            bool b_null = ray_vec_is_null(ks->col, row_b);
            if (a_null && b_null) continue;       /* tie on this key */
            if (a_null) return ks->nulls_first ? -1 : 1;
            if (b_null) return ks->nulls_first ?  1 : -1;
        }
        if (ks->type == RAY_SYM) {
            /* cell-data: ids resolve over the COLUMN's domain.  Runtime-
             * domain columns use the lock-free global snapshot; other
             * domains resolve through ray_sym_domain_str. */
            uint32_t ia = (uint32_t)read_by_esz(ks->base, row_a, ks->esz);
            uint32_t ib = (uint32_t)read_by_esz(ks->base, row_b, ks->esz);
            if (ia == ib) continue;
            if (ks->raw_ok && (int64_t)ia < ks->raw.count && (int64_t)ib < ks->raw.count) {
                /* ray_str_cmp order: bytes of the common prefix, then length. */
                size_t la, lb;
                const char* pa = ray_sym_domain_raw_str(&ks->raw, (int64_t)ia, &la);
                const char* pb = ray_sym_domain_raw_str(&ks->raw, (int64_t)ib, &lb);
                size_t ml = la < lb ? la : lb;
                cmp = ml ? memcmp(pa, pb, ml) : 0;
                if (cmp == 0) cmp = la < lb ? -1 : (la > lb ? 1 : 0);
            } else {
                ray_t* sa;
                ray_t* sb;
                if (ks->dom_runtime) {
                    if (ia >= c->sym_count || ib >= c->sym_count) continue;
                    sa = c->sym_strings[ia];
                    sb = c->sym_strings[ib];
                } else {
                    struct ray_sym_domain_s* dom = ray_sym_vec_domain(ks->col);
                    sa = ray_sym_domain_str(dom, (int64_t)ia);
                    sb = ray_sym_domain_str(dom, (int64_t)ib);
                }
                if (!sa || !sb) continue;
                cmp = ray_str_cmp(sa, sb);
            }
        } else if (ks->type == RAY_STR) {
            /* Variable-length STR key: compare the actual strings (dict
             * codes, if any, are first-occurrence order — not sorted).
             * No null in the STR model; empty "" is a normal value. */
            size_t la = 0, lb = 0;
            const char* sa = ray_str_vec_get(ks->col, row_a, &la);
            const char* sb = ray_str_vec_get(ks->col, row_b, &lb);
            if (!sa) sa = "";
            if (!sb) sb = "";
            size_t ml = la < lb ? la : lb;
            cmp = ml ? memcmp(sa, sb, ml) : 0;
            if (cmp == 0) {
                if (la < lb) cmp = -1;
                else if (la > lb) cmp = 1;
            }
            if (cmp == 0) continue;
        } else if (ks->type == RAY_GUID) {
            cmp = memcmp((const char*)ks->base + (size_t)row_a * 16,
                         (const char*)ks->base + (size_t)row_b * 16, 16);
        } else if (ks->type == RAY_F32 || ks->type == RAY_F64) {
            double va = ks->type == RAY_F32 ? ((const float*)ks->base)[row_a] : ((const double*)ks->base)[row_a];
            double vb = ks->type == RAY_F32 ? ((const float*)ks->base)[row_b] : ((const double*)ks->base)[row_b];
            cmp = (va > vb) - (va < vb);
        } else if (ks->esz == 8) {
            int64_t va = ((const int64_t*)ks->base)[row_a];
            int64_t vb = ((const int64_t*)ks->base)[row_b];
            if (va < vb) cmp = -1; else if (va > vb) cmp = 1;
        } else if (ks->esz == 4) {
            int32_t va = ((const int32_t*)ks->base)[row_a];
            int32_t vb = ((const int32_t*)ks->base)[row_b];
            if (va < vb) cmp = -1; else if (va > vb) cmp = 1;
        } else if (ks->esz == 2) {
            int16_t va = ((const int16_t*)ks->base)[row_a];
            int16_t vb = ((const int16_t*)ks->base)[row_b];
            if (va < vb) cmp = -1; else if (va > vb) cmp = 1;
        } else { /* esz == 1 */
            uint8_t va = ((const uint8_t*)ks->base)[row_a];
            uint8_t vb = ((const uint8_t*)ks->base)[row_b];
            if (va < vb) cmp = -1; else if (va > vb) cmp = 1;
        }
        if (cmp != 0) return ks->desc ? -cmp : cmp;
    }
    /* All sort keys tie: break by source row index.  A stable sort
     * preserves source order on ties — the prefix of K rows that
     * survives is the K rows with the smallest original indices.
     * In the max-heap of K best entries, the root holds the worst
     * survivor (the one most likely to be evicted), so a row with
     * a higher source index ranks worse than a row with a lower
     * one.  Direction-independent: for both ASC and DESC top-K we
     * want stable source-order semantics on ties.
     *
     * Future: caller-specified NULLS FIRST / NULLS LAST.  The
     * has_nulls leg above already reads a per-key flag, but the flag
     * is filled from the direction alone; an explicit clause in the
     * query DSL would thread the caller's choice into it instead.
     * Tracked separately. */
    if (row_a > row_b) return  1;   /* a is worse — evict first */
    if (row_a < row_b) return -1;
    return 0;
}

/* Sift root down (root = current "worst" entry that gets evicted next). */
static inline void fpk_sift_down(const fpk_par_ctx_t* c, int64_t* heap,
                                 int32_t n, int32_t i) {
    for (;;) {
        int32_t worst = i;
        int32_t l = 2 * i + 1, r = 2 * i + 2;
        if (l < n && fpk_cmp(c, heap[l], heap[worst]) > 0) worst = l;
        if (r < n && fpk_cmp(c, heap[r], heap[worst]) > 0) worst = r;
        if (worst == i) return;
        int64_t ti = heap[i]; heap[i] = heap[worst]; heap[worst] = ti;
        i = worst;
    }
}

static inline void fpk_heapify(const fpk_par_ctx_t* c, int64_t* heap, int32_t n) {
    for (int32_t i = n / 2 - 1; i >= 0; i--)
        fpk_sift_down(c, heap, n, i);
}

/* Scan physical rows [row, end): eval predicate per morsel, heap-insert
 * the passing rows.  `*hnp` is the worker's heap fill on entry and exit.
 * `chunk` >= 0 is the zone chunk the rows belong to: the bound may tighten
 * while the chunk is being scanned (another worker's heap filled, or this
 * one's), so it is re-tested before every morsel — one relaxed load and a
 * compare — and the rest of the chunk is abandoned once it cannot beat it. */
static inline void fpk_scan_rows(fpk_par_ctx_t* c, int64_t* hidx, int32_t* hnp,
                                 int32_t k, int64_t row, int64_t end, int64_t chunk) {
    int32_t hn = *hnp;
    while (row < end) {
        if (chunk >= 0 && fpk_chunk_pruned(c, chunk)) break;
        int64_t mend = row + RAY_MORSEL_ELEMS;
        if (mend > end) mend = end;
        int64_t mlen = mend - row;
        uint8_t bits[RAY_MORSEL_ELEMS];
        fp_eval_pred(&c->pred, row, mend, bits);

        for (int64_t r = 0; r < mlen; r++) {
            if (!bits[r]) continue;
            int64_t src_row = row + r;
            if (hn < k) {
                hidx[hn++] = src_row;
                if (hn == k) fpk_heapify(c, hidx, k);
            } else {
                /* Reject fast: skip if new ≥ current worst (heap root). */
                if (fpk_cmp(c, src_row, hidx[0]) >= 0) continue;
                hidx[0] = src_row;
                fpk_sift_down(c, hidx, k, 0);
            }
        }
        /* Once per morsel: publishing on every heap replacement contends
         * on the bound when the rows arrive in the order being sought. */
        if (hn == k) fpk_publish_bound(c, hidx[0]);
        row = mend;
    }
    *hnp = hn;
}

/* Worker fn: [start, end) is a range of physical rows, or — with a zone
 * order — of virtual rows whose chunk slots map to physical chunks. */
static void fpk_par_fn(void* raw, uint32_t worker_id, int64_t start, int64_t end) {
    fpk_par_ctx_t* c = (fpk_par_ctx_t*)raw;
    if (atomic_load_explicit(&c->oom, memory_order_relaxed)) return;
    int32_t  k    = (int32_t)c->k;
    int64_t* hidx = &c->heap_idx[(size_t)worker_id * (size_t)k];
    int32_t  hn   = c->heap_n[worker_id];

    if (!c->zmin) {
        fpk_scan_rows(c, hidx, &hn, k, start, end, -1);
    } else {
        int64_t span = (int64_t)1 << c->zlog2;
        int64_t v = start;
        while (v < end) {
            int64_t slot = v >> c->zlog2;
            int64_t sbeg = slot << c->zlog2;
            int64_t send = sbeg + span;
            if (send > end) send = end;
            int64_t g = c->zorder ? (int64_t)c->zorder[slot] : slot;
            if (!fpk_chunk_pruned(c, g)) {
                int64_t p0 = (g << c->zlog2) + (v - sbeg);
                int64_t p1 = (g << c->zlog2) + (send - sbeg);
                if (p1 > c->nrows) p1 = c->nrows;
                if (p0 < p1) fpk_scan_rows(c, hidx, &hn, k, p0, p1, g);
            }
            v = send;
        }
    }
    c->heap_n[worker_id] = hn;
}

/* Insertion sort over K entries so output rows match full-sort + take order. */
static void fpk_sort_final(const fpk_par_ctx_t* c, int64_t* idx, int32_t n) {
    for (int32_t i = 1; i < n; i++) {
        int64_t v = idx[i];
        int32_t j = i - 1;
        while (j >= 0 && fpk_cmp(c, idx[j], v) > 0) {
            idx[j + 1] = idx[j];
            j--;
        }
        idx[j + 1] = v;
    }
}

ray_t* ray_fused_topk_select(ray_t* tbl,
                             ray_t* where_expr,
                             const int64_t* sort_key_syms,
                             const uint8_t* sort_descs,
                             uint32_t n_sort_keys,
                             int64_t k,
                             const int64_t* out_col_syms,
                             const int64_t* out_alias_syms,
                             uint32_t n_out)
{
    if (!tbl || tbl->type != RAY_TABLE || k <= 0 || n_out == 0) return NULL;
    if (k > FPK_MAX_K) return NULL;
    if (n_sort_keys == 0 || n_sort_keys > FPK_MAX_KEYS) return NULL;
    int64_t nrows = ray_table_nrows(tbl);
    if (nrows <= 0 || k >= nrows) return NULL;

    for (uint32_t c = 0; c < n_out; c++) {
        ray_t* col = ray_table_get_col(tbl, out_col_syms[c]);
        if (!col) return NULL;
        int8_t ot = col->type;
        if (RAY_IS_PARTED(ot) || ot == RAY_MAPCOMMON) return NULL;
        if (!ray_is_vec(col))
            return NULL;
    }

    /* Resolve sort-key columns + decide if any need the SYM dict snapshot. */
    fpk_par_ctx_t ctx;
    memset(&ctx, 0, sizeof(ctx));
    int sym_needed = 0;
    for (uint32_t i = 0; i < n_sort_keys; i++) {
        ray_t* col = ray_table_get_col(tbl, sort_key_syms[i]);
        if (!col) { fpk_unpin_keys(ctx.keys, i); return NULL; }
        int8_t kt = col->type;
        if (RAY_IS_PARTED(kt) || kt == RAY_MAPCOMMON) { fpk_unpin_keys(ctx.keys, i); return NULL; }
        if (kt != RAY_SYM && kt != RAY_STR && kt != RAY_BOOL && kt != RAY_U8
            && kt != RAY_I16 && kt != RAY_I32 && kt != RAY_I64
            && kt != RAY_DATE && kt != RAY_TIME && kt != RAY_TIMESTAMP
            && kt != RAY_F32 && kt != RAY_F64 && kt != RAY_GUID)
            { fpk_unpin_keys(ctx.keys, i); return NULL; }
        ctx.keys[i].type      = kt;
        ctx.keys[i].attrs     = col->attrs;
        ctx.keys[i].esz       = ray_sym_elem_size(kt, col->attrs);
        ctx.keys[i].desc      = sort_descs[i];
        ctx.keys[i].has_nulls = ray_vec_may_have_nulls(col) ? 1 : 0;
        ctx.keys[i].nulls_first = sort_nulls_first(sort_descs[i]);
        ctx.keys[i].dom_runtime =
            (kt == RAY_SYM &&
             ray_sym_vec_domain(col) == ray_sym_runtime_domain()) ? 1 : 0;
        ctx.keys[i].sym       = sort_key_syms[i];
        ctx.keys[i].base      = ray_data(col);
        ctx.keys[i].col       = col;
        ctx.keys[i].raw_ok    = 0;
        if (ctx.keys[i].dom_runtime) sym_needed = 1;
        else if (kt == RAY_SYM)
            ctx.keys[i].raw_ok = ray_sym_domain_raw_pin(ray_sym_vec_domain(col), &ctx.keys[i].raw) ? 1 : 0;
    }
    ctx.n_keys = n_sort_keys;
    ctx.k      = k;
    ctx.tbl    = tbl;
    ctx.nrows  = nrows;
    {
        ray_t* kc = ctx.keys[0].col;
        int8_t kt = ctx.keys[0].type;
        if ((kt == RAY_I16 || kt == RAY_I32 || kt == RAY_I64 || kt == RAY_DATE ||
             kt == RAY_TIME || kt == RAY_TIMESTAMP) &&
            ray_index_kind(kc) == RAY_IDX_CHUNK_ZONE) {
            ray_index_t* zx = ray_index_payload(kc->index);
            if (zx->built_for_len == kc->len && !zx->u.chunk_zone.is_f64 &&
                zx->u.chunk_zone.mins && zx->u.chunk_zone.maxs && zx->u.chunk_zone.null_bits) {
                ctx.zmin  = (const int64_t*)ray_data(zx->u.chunk_zone.mins);
                ctx.zmax  = (const int64_t*)ray_data(zx->u.chunk_zone.maxs);
                ctx.znull = (const uint8_t*)ray_data(zx->u.chunk_zone.null_bits);
                ctx.zn    = zx->u.chunk_zone.n_chunks;
                ctx.zlog2 = zx->u.chunk_zone.chunk_log2;
                ctx.zdesc = ctx.keys[0].desc;
                ctx.znulls_better = ctx.keys[0].nulls_first;
                if (((int64_t)ctx.zn << ctx.zlog2) < nrows)
                    ctx.zmin = NULL; /* index shorter than the column: no pruning */
            }
        }
    }

    /* Compile the predicate via a temp graph just for the WHERE clause.  No
     * where: is a predicate with no children — every row passes — and the
     * heap does the whole sort-and-take in one pass. */
    ray_graph_t* g = ray_graph_new(tbl);
    if (!g) { fpk_unpin_keys(ctx.keys, n_sort_keys); return NULL; }
    if (where_expr) {
        ray_op_t* pred_dag = compile_expr_dag(g, where_expr);
        if (!pred_dag) { ray_graph_free(g); fpk_unpin_keys(ctx.keys, n_sort_keys); return NULL; }
        if (fp_compile_pred(g, pred_dag, tbl, &ctx.pred) != 0) {
            fp_pred_cleanup(&ctx.pred);
            ray_graph_free(g);
            fpk_unpin_keys(ctx.keys, n_sort_keys);
            return NULL;
        }
    }

    if (sym_needed) {
        ray_sym_strings_borrow(&ctx.sym_strings, &ctx.sym_count);
        if (!ctx.sym_strings) { ray_graph_free(g); fpk_unpin_keys(ctx.keys, n_sort_keys); return NULL; }
    }
    atomic_store_explicit(&ctx.oom, 0, memory_order_relaxed);

    /* Allocate per-worker heaps. */
    ray_pool_t* pool = ray_pool_get();
    uint32_t nw = pool ? ray_pool_total_workers(pool) : 1;
    ray_t* idx_hdr = NULL;
    ray_t* hn_hdr  = NULL;
    ctx.heap_idx = (int64_t*)scratch_alloc(&idx_hdr,
                                           (size_t)nw * (size_t)k * sizeof(int64_t));
    ctx.heap_n   = (int32_t*)scratch_calloc(&hn_hdr,
                                            (size_t)nw * sizeof(int32_t));
    if (!ctx.heap_idx || !ctx.heap_n) {
        if (idx_hdr) scratch_free(idx_hdr);
        if (hn_hdr)  scratch_free(hn_hdr);
        fp_pred_cleanup(&ctx.pred);
        ray_graph_free(g);
        fpk_unpin_keys(ctx.keys, n_sort_keys);
        return NULL;
    }

    /* With a zone index the workers walk the chunks best-first over the
     * virtual row space (the order is optional: without it the slots map
     * to the physical chunks and pruning still applies). */
    int64_t span_rows = nrows;
    ray_t*  zord_hdr  = NULL;
    if (ctx.zmin) {
        span_rows = (int64_t)ctx.zn << ctx.zlog2;
        uint32_t* ord = (uint32_t*)scratch_alloc(&zord_hdr,
                                                 (size_t)ctx.zn * sizeof(uint32_t));
        if (ord) {
            for (uint32_t g = 0; g < ctx.zn; g++) ord[g] = g;
            fpk_order_chunks(&ctx, ord, ctx.zn);
            ctx.zorder = ord;
        }
    }
    if (pool) ray_pool_dispatch(pool, fpk_par_fn, &ctx, span_rows);
    else      fpk_par_fn(&ctx, 0, 0, span_rows);
    if (zord_hdr) { scratch_free(zord_hdr); ctx.zorder = NULL; }

    if (atomic_load_explicit(&ctx.oom, memory_order_relaxed)) {
        scratch_free(idx_hdr); scratch_free(hn_hdr);
        fp_pred_cleanup(&ctx.pred);
        ray_graph_free(g);
        fpk_unpin_keys(ctx.keys, n_sort_keys);
        return NULL;
    }

    /* Combine per-worker heaps into one global K-heap.  Stack-resident
     * to avoid an alloc on the hot combine path.  k <= FPK_MAX_K is
     * enforced at function entry. */
    int64_t global_idx[FPK_MAX_K];
    int32_t global_n = 0;
    for (uint32_t w = 0; w < nw; w++) {
        int32_t  hn   = ctx.heap_n[w];
        int64_t* hidx = &ctx.heap_idx[(size_t)w * (size_t)k];
        for (int32_t i = 0; i < hn; i++) {
            if (global_n < (int32_t)k) {
                // cppcheck-suppress objectIndex // heap_idx is an nw*k scratch buffer and hn <= k; cppcheck misreads ray_data's inline payload as a scalar
                global_idx[global_n++] = hidx[i];
                if (global_n == (int32_t)k)
                    fpk_heapify(&ctx, global_idx, global_n);
            } else {
                // cppcheck-suppress objectIndex // same nw*k buffer as above
                if (fpk_cmp(&ctx, hidx[i], global_idx[0]) >= 0) continue;
                // cppcheck-suppress objectIndex // same nw*k buffer as above
                global_idx[0] = hidx[i];
                fpk_sift_down(&ctx, global_idx, global_n, 0);
            }
        }
    }

    scratch_free(idx_hdr);
    scratch_free(hn_hdr);

    fpk_sort_final(&ctx, global_idx, global_n);
    fpk_unpin_keys(ctx.keys, n_sort_keys);   /* no more compares past here */

    /* Materialize n_out output columns by gathering rows[global_idx]. */
    ray_t* result = ray_table_new(n_out);
    if (!result || RAY_IS_ERR(result)) {
        fp_pred_cleanup(&ctx.pred);
        ray_graph_free(g);
        return result ? result : ray_error("oom", NULL);
    }
    int build_ok = 1;
    for (uint32_t c = 0; c < n_out; c++) {
        int64_t cs    = out_col_syms[c];
        int64_t alias = out_alias_syms ? out_alias_syms[c] : cs;
        ray_t* src = ray_table_get_col(tbl, cs);
        if (!src) { build_ok = 0; break; }
        ray_t* col = gather_by_idx(src, global_idx, global_n);
        if (!col || RAY_IS_ERR(col)) { build_ok = 0; break; }
        result = ray_table_add_col(result, alias, col);
        ray_release(col);
    }
    ray_graph_free(g);
    fp_pred_cleanup(&ctx.pred);
    if (!build_ok) {
        ray_release(result);
        return ray_error("schema", NULL);
    }
    return result;
}

/* ───── Fused filter + positional take ────────────────────────────────
 * Chunks of FTK_CHUNK_ROWS rows, numbered from the end the answer comes
 * from (row 0 for the first K, the last row for the last |K|), one pool
 * task per chunk.  Each worker keeps the |K| passing rows nearest that
 * end it has seen, in scan order; once it holds |K| they bound the
 * answer, so it publishes its |K|-th row as the cutoff: no row beyond it
 * can be among the first |K| passing rows of the table, and every chunk
 * past it returns at once.  The pool makes no promise about the order
 * tasks run in (see pool.h), so a worker whose list is full may still be
 * handed a chunk that lies BEFORE the rows it holds — it scans it and
 * lets nearer rows displace the farthest; only the cutoff, which is
 * sound whatever the order, skips work.  The lists are merged by row id
 * at the end and the |K| nearest the scanned end are gathered.
 * ──────────────────────────────────────────────────────────────────── */

#define FTK_CHUNK_ROWS (64 * 1024)

typedef struct {
    fp_pred_t         pred;
    int64_t           nrows;
    int64_t           k;         /* |K| */
    bool              from_end;
    int64_t*          rows;      /* [nw * k] per-worker row ids, in scan order */
    int32_t*          rows_n;    /* [nw] */
    _Atomic(int64_t)  cutoff;    /* forward: rows >= cutoff are out; backward: rows <= cutoff */
} ftk_ctx_t;

static inline bool ftk_beyond(const ftk_ctx_t* c, int64_t row) {
    int64_t cut = atomic_load_explicit(&c->cutoff, memory_order_relaxed);
    return c->from_end ? row <= cut : row >= cut;
}

static void ftk_publish(ftk_ctx_t* c, int64_t row) {
    int64_t cur = atomic_load_explicit(&c->cutoff, memory_order_relaxed);
    for (;;) {
        bool tighter = c->from_end ? row > cur : row < cur;
        if (!tighter) return;
        if (atomic_compare_exchange_weak_explicit(&c->cutoff, &cur, row,
                memory_order_relaxed, memory_order_relaxed))
            return;
    }
}

/* Admit `row` to the worker's list of the |K| nearest passing rows, kept in
 * scan order.  Rows arrive in scan order within a chunk, but a later chunk
 * may lie before an earlier one, so a full list still takes a nearer row:
 * it is inserted in place and the farthest drops off.  False when the list
 * is full and `row` is not nearer than its last entry — then nothing after
 * it in this chunk is either. */
static inline bool ftk_push(const ftk_ctx_t* c, int64_t* my, int32_t* pn, int64_t row) {
    int32_t n = *pn, k = (int32_t)c->k;
    if (n == k) {
        if (c->from_end ? row <= my[k - 1] : row >= my[k - 1]) return false;
        n--;                                  /* the farthest drops off */
    }
    int32_t i = n;
    while (i > 0 && (c->from_end ? my[i - 1] < row : my[i - 1] > row)) {
        my[i] = my[i - 1];
        i--;
    }
    my[i] = row;
    *pn = n + 1;
    return true;
}

static void ftk_task_fn(void* raw, uint32_t worker_id, int64_t start, int64_t end) {
    (void)end;
    ftk_ctx_t* c = (ftk_ctx_t*)raw;
    int32_t  k  = (int32_t)c->k;
    int64_t* my = &c->rows[(size_t)worker_id * (size_t)k];
    int32_t  n  = c->rows_n[worker_id];

    /* chunk `start` counted from the scanned end */
    int64_t lo, hi;
    if (!c->from_end) {
        lo = start * FTK_CHUNK_ROWS;
        hi = lo + FTK_CHUNK_ROWS;
        if (hi > c->nrows) hi = c->nrows;
    } else {
        hi = c->nrows - start * FTK_CHUNK_ROWS;
        lo = hi - FTK_CHUNK_ROWS;
        if (lo < 0) lo = 0;
    }
    if (ftk_beyond(c, c->from_end ? hi - 1 : lo)) return;

    uint8_t bits[RAY_MORSEL_ELEMS];
    bool more = true;                         /* false: list full, rest of chunk is farther */
    if (!c->from_end) {
        for (int64_t row = lo; row < hi && more; ) {
            if (ftk_beyond(c, row)) break;
            int64_t mend = row + RAY_MORSEL_ELEMS;
            if (mend > hi) mend = hi;
            fp_eval_pred(&c->pred, row, mend, bits);
            for (int64_t r = 0; r < mend - row && more; r++)
                if (bits[r]) more = ftk_push(c, my, &n, row + r);
            row = mend;
        }
    } else {
        for (int64_t mend = hi; mend > lo && more; ) {
            if (ftk_beyond(c, mend - 1)) break;
            int64_t row = mend - RAY_MORSEL_ELEMS;
            if (row < lo) row = lo;
            fp_eval_pred(&c->pred, row, mend, bits);
            for (int64_t r = mend - row - 1; r >= 0 && more; r--)
                if (bits[r]) more = ftk_push(c, my, &n, row + r);
            mend = row;
        }
    }
    c->rows_n[worker_id] = n;
    if (n >= k) ftk_publish(c, my[k - 1]);
}

ray_t* ray_fused_take_select(ray_t* tbl,
                             ray_t* where_expr,
                             int64_t k,
                             const int64_t* out_col_syms,
                             const int64_t* out_alias_syms,
                             uint32_t n_out)
{
    if (!tbl || tbl->type != RAY_TABLE || !where_expr || k == 0 || n_out == 0) return NULL;
    if (k == INT64_MIN) return NULL;
    bool from_end = k < 0;
    if (from_end) k = -k;
    if (k > FPK_MAX_K) return NULL;
    int64_t nrows = ray_table_nrows(tbl);
    if (nrows <= 0 || k >= nrows) return NULL;

    for (uint32_t c = 0; c < n_out; c++) {
        ray_t* col = ray_table_get_col(tbl, out_col_syms[c]);
        if (!col) return NULL;
        int8_t ot = col->type;
        if (RAY_IS_PARTED(ot) || ot == RAY_MAPCOMMON) return NULL;
        if (!ray_is_vec(col)) return NULL;
    }

    ftk_ctx_t ctx;
    memset(&ctx, 0, sizeof(ctx));
    ctx.nrows = nrows;
    ctx.k = k;
    ctx.from_end = from_end;
    atomic_store_explicit(&ctx.cutoff, from_end ? -1 : INT64_MAX, memory_order_relaxed);

    ray_graph_t* g = ray_graph_new(tbl);
    if (!g) return NULL;
    ray_op_t* pred_dag = compile_expr_dag(g, where_expr);
    if (!pred_dag) { ray_graph_free(g); return NULL; }
    if (fp_compile_pred(g, pred_dag, tbl, &ctx.pred) != 0) {
        fp_pred_cleanup(&ctx.pred);
        ray_graph_free(g);
        return NULL;
    }

    ray_pool_t* pool = ray_pool_get();
    uint32_t nw = pool ? ray_pool_total_workers(pool) : 1;
    ray_t* rows_hdr = NULL;
    ray_t* n_hdr    = NULL;
    ctx.rows   = (int64_t*)scratch_alloc(&rows_hdr, (size_t)nw * (size_t)k * sizeof(int64_t));
    ctx.rows_n = (int32_t*)scratch_calloc(&n_hdr, (size_t)nw * sizeof(int32_t));
    if (!ctx.rows || !ctx.rows_n) {
        if (rows_hdr) scratch_free(rows_hdr);
        if (n_hdr) scratch_free(n_hdr);
        fp_pred_cleanup(&ctx.pred);
        ray_graph_free(g);
        return NULL;
    }

    int64_t n_chunks = (nrows + FTK_CHUNK_ROWS - 1) / FTK_CHUNK_ROWS;
    if (ray_pool_par_dispatch_ok(pool, n_chunks, 2))
        ray_pool_dispatch_n(pool, ftk_task_fn, &ctx, (uint32_t)n_chunks);
    else
        for (int64_t t = 0; t < n_chunks; t++) ftk_task_fn(&ctx, 0, t, t + 1);

    /* Merge: each list is in scan order; pick the row nearest the scanned
     * end across lists k times, then present in table order. */
    int64_t out[FPK_MAX_K];
    int32_t out_n = 0;
    ray_t* pos_hdr = NULL;
    int32_t* pos = (int32_t*)scratch_calloc(&pos_hdr, (size_t)nw * sizeof(int32_t));
    if (!pos) {
        scratch_free(rows_hdr); scratch_free(n_hdr);
        fp_pred_cleanup(&ctx.pred);
        ray_graph_free(g);
        return NULL;
    }
    while (out_n < (int32_t)k) {
        int64_t best = -1;
        uint32_t bw = 0;
        for (uint32_t w = 0; w < nw; w++) {
            if (pos[w] >= ctx.rows_n[w]) continue;
            int64_t r = ctx.rows[(size_t)w * (size_t)k + (size_t)pos[w]];
            if (best < 0 || (from_end ? r > best : r < best)) { best = r; bw = w; }
        }
        if (best < 0) break;
        pos[bw]++;
        out[out_n++] = best;
    }
    scratch_free(pos_hdr);
    scratch_free(rows_hdr);
    scratch_free(n_hdr);
    if (from_end) {
        for (int32_t i = 0, j = out_n - 1; i < j; i++, j--) {
            int64_t t = out[i]; out[i] = out[j]; out[j] = t;
        }
    }

    ray_t* result = ray_table_new(n_out);
    if (!result || RAY_IS_ERR(result)) {
        fp_pred_cleanup(&ctx.pred);
        ray_graph_free(g);
        return result ? result : ray_error("oom", NULL);
    }
    int build_ok = 1;
    for (uint32_t c = 0; c < n_out; c++) {
        int64_t cs    = out_col_syms[c];
        int64_t alias = out_alias_syms ? out_alias_syms[c] : cs;
        ray_t* src = ray_table_get_col(tbl, cs);
        if (!src) { build_ok = 0; break; }
        ray_t* col = gather_by_idx(src, out, out_n);
        if (!col || RAY_IS_ERR(col)) { build_ok = 0; break; }
        result = ray_table_add_col(result, alias, col);
        ray_release(col);
    }
    ray_graph_free(g);
    fp_pred_cleanup(&ctx.pred);
    if (!build_ok) {
        ray_release(result);
        return ray_error("schema", NULL);
    }
    return result;
}