tranfi 0.1.2 → 0.2.0

package/csrc/op_sort.c

@@ -1,94 +1,164 @@
 /*
- * op_sort.c — Sort all rows by column(s). Requires buffering all data.
+ * op_sort.c -- Sort rows by column(s).
  *
- * Config: {"columns": [{"name": "age", "desc": false}, {"name": "name", "desc": true}]}
- * Buffers all incoming batches, sorts on flush, emits as a single batch.
+ * Default mode preserves the historical in-memory/blocking sort. When the
+ * args contain spill_dir, the operator writes memory-sized sorted runs to
+ * temporary files and merges them at flush in bounded output batches.
  */
 
 #include "internal.h"
+#include "spill.h"
 #include "cJSON.h"
+
+#include <errno.h>
+#include <stdint.h>
+#include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#include <stdio.h>
+#include <unistd.h>
+
+#define SORT_DEFAULT_RUN_ROWS 8192
+#define SORT_DEFAULT_OUTPUT_ROWS 1024
+#define SORT_MIN_RUN_ROWS 16
 
+/* Sort spec. */
 typedef struct {
     char *name;
-    int   desc;   /* 1 = descending */
+    int   desc;
 } sort_col;
 
+typedef tf_owned_cell_value sort_cell;
+
+typedef struct {
+    uint8_t *nulls;
+    sort_cell *cells;
+} sort_spill_row;
+
+typedef struct {
+    FILE *file;
+    sort_spill_row row;
+    int has_row;
+    int done;
+} sort_run_reader;
+
 typedef struct {
-    /* Accumulated data */
-    tf_batch *buf;       /* growing buffer of all rows */
+    tf_batch *buf;
     int       has_schema;
+    char    **schema_names;
+    tf_type  *schema_types;
+    size_t    n_schema_cols;
 
-    /* Sort spec */
     sort_col *cols;
     size_t    n_cols;
+    int      *col_indices;
+    int      *col_desc;
+
+    int       use_spill;
+    char     *spill_dir;
+    tf_spill_session *spill;
+    size_t    spill_memory_bytes;
+    size_t    configured_run_rows;
+    size_t    run_rows;
+    size_t    output_batch_rows;
+
+    char    **run_paths;
+    size_t    n_runs;
+    size_t    cap_runs;
+    size_t    run_seq;
+    size_t    spilled_bytes;
+    size_t    spill_runs_created;
+    size_t    spill_output_batches;
+    size_t    spill_output_rows;
+
+    sort_run_reader *readers;
+    size_t    n_readers;
+    int       merge_started;
+    int       merge_done;
 } sort_state;
 
-/* Copy a row from src batch to dst batch */
-static int copy_row(tf_batch *dst, size_t dst_row,
-                    const tf_batch *src, size_t src_row) {
-    if (tf_batch_ensure_capacity(dst, dst_row + 1) != TF_OK) return TF_ERROR;
-    for (size_t c = 0; c < src->n_cols; c++) {
-        if (tf_batch_is_null(src, src_row, c)) {
-            tf_batch_set_null(dst, dst_row, c);
-            continue;
-        }
-        switch (src->col_types[c]) {
-            case TF_TYPE_BOOL:
-                tf_batch_set_bool(dst, dst_row, c, tf_batch_get_bool(src, src_row, c));
-                break;
-            case TF_TYPE_INT64:
-                tf_batch_set_int64(dst, dst_row, c, tf_batch_get_int64(src, src_row, c));
-                break;
-            case TF_TYPE_FLOAT64:
-                tf_batch_set_float64(dst, dst_row, c, tf_batch_get_float64(src, src_row, c));
-                break;
-            case TF_TYPE_STRING:
-                tf_batch_set_string(dst, dst_row, c, tf_batch_get_string(src, src_row, c));
-                break;
-            case TF_TYPE_DATE:
-                tf_batch_set_date(dst, dst_row, c, tf_batch_get_date(src, src_row, c));
-                break;
-            case TF_TYPE_TIMESTAMP:
-                tf_batch_set_timestamp(dst, dst_row, c, tf_batch_get_timestamp(src, src_row, c));
-                break;
-            default:
-                tf_batch_set_null(dst, dst_row, c);
-                break;
+static tf_batch *create_buffer_from_schema(const sort_state *st, size_t capacity) {
+    tf_batch *b = tf_batch_create(st->n_schema_cols, capacity ? capacity : 16);
+    if (!b) return NULL;
+    for (size_t c = 0; c < st->n_schema_cols; c++) {
+        if (tf_batch_set_schema(b, c, st->schema_names[c], st->schema_types[c]) != TF_OK) {
+            tf_batch_free(b);
+            return NULL;
         }
     }
-    return TF_OK;
+    return b;
 }
 
-static int sort_process(tf_step *self, tf_batch *in, tf_batch **out,
-                        tf_side_channels *side) {
-    (void)side;
-    sort_state *st = self->state;
-    *out = NULL;
+static size_t sort_estimated_row_bytes(const sort_state *st) {
+    size_t bytes = 32;
+    for (size_t c = 0; c < st->n_schema_cols; c++) {
+        bytes += 1;
+        switch (st->schema_types[c]) {
+            case TF_TYPE_BOOL: bytes += 1; break;
+            case TF_TYPE_INT64: bytes += sizeof(int64_t); break;
+            case TF_TYPE_FLOAT64: bytes += sizeof(double); break;
+            case TF_TYPE_STRING: bytes += sizeof(char *) + 64; break;
+            case TF_TYPE_DATE: bytes += sizeof(int32_t); break;
+            case TF_TYPE_TIMESTAMP: bytes += sizeof(int64_t); break;
+            default: break;
+        }
+    }
+    return bytes < 64 ? 64 : bytes;
+}
 
-    /* Initialize buffer on first batch */
-    if (!st->has_schema) {
-        st->buf = tf_batch_create(in->n_cols, in->n_rows > 0 ? in->n_rows : 16);
-        if (!st->buf) return TF_ERROR;
-        for (size_t c = 0; c < in->n_cols; c++) {
-            tf_batch_set_schema(st->buf, c, in->col_names[c], in->col_types[c]);
+static int resolve_sort_columns(sort_state *st) {
+    free(st->col_indices);
+    free(st->col_desc);
+    st->col_indices = tf_callocarray_checked(st->n_cols ? st->n_cols : 1, sizeof(int));
+    st->col_desc = tf_callocarray_checked(st->n_cols ? st->n_cols : 1, sizeof(int));
+    if (!st->col_indices || !st->col_desc) return TF_ERROR;
+    for (size_t k = 0; k < st->n_cols; k++) {
+        int idx = tf_batch_col_index(st->buf, st->cols[k].name);
+        if (idx < 0) {
+            char msg[256];
+            snprintf(msg, sizeof(msg), "sort: column '%s' not found",
+                     st->cols[k].name ? st->cols[k].name : "");
+            tf_set_last_error(msg);
+            return TF_ERROR;
         }
-        st->has_schema = 1;
+        st->col_indices[k] = idx;
+        st->col_desc[k] = st->cols[k].desc;
     }
+    return TF_OK;
+}
 
-    /* Append all rows */
-    for (size_t r = 0; r < in->n_rows; r++) {
-        size_t dst_row = st->buf->n_rows;
-        if (copy_row(st->buf, dst_row, in, r) != TF_OK) return TF_ERROR;
-        st->buf->n_rows = dst_row + 1;
+static int init_schema(sort_state *st, const tf_batch *in) {
+    if (st->has_schema) return TF_OK;
+
+    st->n_schema_cols = in->n_cols;
+    st->schema_names = tf_callocarray_checked(in->n_cols ? in->n_cols : 1, sizeof(char *));
+    st->schema_types = tf_callocarray_checked(in->n_cols ? in->n_cols : 1, sizeof(tf_type));
+    if (!st->schema_names || !st->schema_types) return TF_ERROR;
+    for (size_t c = 0; c < in->n_cols; c++) {
+        st->schema_names[c] = strdup(in->col_names[c] ? in->col_names[c] : "");
+        if (!st->schema_names[c]) return TF_ERROR;
+        st->schema_types[c] = in->col_types[c];
     }
 
-    return TF_OK;
+    size_t initial = in->n_rows > 0 ? in->n_rows : 16;
+    if (st->use_spill) {
+        if (st->configured_run_rows > 0) {
+            st->run_rows = st->configured_run_rows;
+        } else if (st->spill_memory_bytes > 0) {
+            size_t row_bytes = sort_estimated_row_bytes(st);
+            st->run_rows = st->spill_memory_bytes / (row_bytes * 3);
+            if (st->run_rows < SORT_MIN_RUN_ROWS) st->run_rows = SORT_MIN_RUN_ROWS;
+        } else {
+            st->run_rows = SORT_DEFAULT_RUN_ROWS;
+        }
+        initial = st->run_rows;
+    }
+
+    st->buf = create_buffer_from_schema(st, initial);
+    if (!st->buf) return TF_ERROR;
+    st->has_schema = 1;
+    return resolve_sort_columns(st);
 }
 
-/* Comparator context for qsort_r / qsort */
 typedef struct {
     const tf_batch *batch;
     int            *col_indices;
@@ -96,21 +166,14 @@ typedef struct {
     size_t          n_sort_cols;
 } sort_ctx;
 
-static sort_ctx *g_sort_ctx; /* global for qsort (no qsort_r on all platforms) */
-
-static int compare_rows(const void *a, const void *b) {
-    size_t ra = *(const size_t *)a;
-    size_t rb = *(const size_t *)b;
-    const tf_batch *batch = g_sort_ctx->batch;
+static int sort_compare_rows(const sort_ctx *ctx, size_t ra, size_t rb) {
+    const tf_batch *batch = ctx->batch;
 
-    for (size_t k = 0; k < g_sort_ctx->n_sort_cols; k++) {
-        int ci = g_sort_ctx->col_indices[k];
-        if (ci < 0) continue;
+    for (size_t k = 0; k < ctx->n_sort_cols; k++) {
+        int ci = ctx->col_indices[k];
 
         int null_a = tf_batch_is_null(batch, ra, ci);
         int null_b = tf_batch_is_null(batch, rb, ci);
-
-        /* Nulls sort last */
         if (null_a && null_b) continue;
         if (null_a) return 1;
         if (null_b) return -1;
@@ -155,80 +218,527 @@ static int compare_rows(const void *a, const void *b) {
                 break;
         }
 
-        if (cmp != 0) {
-            return g_sort_ctx->col_desc[k] ? -cmp : cmp;
-        }
+        if (cmp != 0) return ctx->col_desc[k] ? -cmp : cmp;
     }
     return 0;
 }
 
-static int sort_flush(tf_step *self, tf_batch **out, tf_side_channels *side) {
-    (void)side;
-    sort_state *st = self->state;
-    *out = NULL;
+static int sort_compare_index(const void *ctx, size_t ra, size_t rb) {
+    int cmp = sort_compare_rows((const sort_ctx *)ctx, ra, rb);
+    if (cmp != 0) return cmp;
+    return (ra > rb) - (ra < rb);
+}
 
+static size_t *sort_batch_indices(const sort_state *st, const tf_batch *batch) {
+    size_t n = batch->n_rows;
+    size_t *indices = tf_mallocarray_checked(n ? n : 1, sizeof(size_t));
+    if (!indices) return NULL;
+    for (size_t i = 0; i < n; i++) indices[i] = i;
+    sort_ctx ctx = {
+        .batch = batch,
+        .col_indices = st->col_indices,
+        .col_desc = st->col_desc,
+        .n_sort_cols = st->n_cols,
+    };
+    tf_sort_indices(indices, n, sort_compare_index, &ctx);
+    return indices;
+}
+
+static int write_exact(FILE *f, const void *ptr, size_t len) {
+    return fwrite(ptr, 1, len, f) == len ? TF_OK : TF_ERROR;
+}
+
+static int read_exact(FILE *f, void *ptr, size_t len) {
+    return fread(ptr, 1, len, f) == len ? TF_OK : TF_ERROR;
+}
+
+static int write_cell(FILE *f, const tf_batch *b, size_t r, size_t c) {
+    uint8_t is_null = tf_batch_is_null(b, r, c) ? 1 : 0;
+    if (write_exact(f, &is_null, sizeof(is_null)) != TF_OK) return TF_ERROR;
+    if (is_null) return TF_OK;
+
+    switch (b->col_types[c]) {
+        case TF_TYPE_BOOL: {
+            uint8_t v = tf_batch_get_bool(b, r, c) ? 1 : 0;
+            return write_exact(f, &v, sizeof(v));
+        }
+        case TF_TYPE_INT64: {
+            int64_t v = tf_batch_get_int64(b, r, c);
+            return write_exact(f, &v, sizeof(v));
+        }
+        case TF_TYPE_FLOAT64: {
+            double v = tf_batch_get_float64(b, r, c);
+            return write_exact(f, &v, sizeof(v));
+        }
+        case TF_TYPE_STRING: {
+            const char *s = tf_batch_get_string(b, r, c);
+            uint64_t len = s ? (uint64_t)strlen(s) : 0;
+            if (write_exact(f, &len, sizeof(len)) != TF_OK) return TF_ERROR;
+            return len ? write_exact(f, s, (size_t)len) : TF_OK;
+        }
+        case TF_TYPE_DATE: {
+            int32_t v = tf_batch_get_date(b, r, c);
+            return write_exact(f, &v, sizeof(v));
+        }
+        case TF_TYPE_TIMESTAMP: {
+            int64_t v = tf_batch_get_timestamp(b, r, c);
+            return write_exact(f, &v, sizeof(v));
+        }
+        default:
+            return TF_OK;
+    }
+}
+
+static int append_run_path(sort_state *st, char *path) {
+    if (st->n_runs == st->cap_runs) {
+        size_t need = 0;
+        size_t new_cap = 0;
+        if (tf_size_add(st->n_runs, 1, &need) != TF_OK ||
+            tf_size_grow_pow2(st->cap_runs, need, 8, &new_cap) != TF_OK) {
+            return TF_ERROR;
+        }
+        char **tmp = tf_reallocarray_checked(st->run_paths, new_cap, sizeof(char *));
+        if (!tmp) return TF_ERROR;
+        st->run_paths = tmp;
+        st->cap_runs = new_cap;
+    }
+    st->run_paths[st->n_runs++] = path;
+    st->spill_runs_created++;
+    return TF_OK;
+}
+
+static int write_spill_run(sort_state *st) {
     if (!st->buf || st->buf->n_rows == 0) return TF_OK;
 
-    size_t n = st->buf->n_rows;
+    size_t *indices = sort_batch_indices(st, st->buf);
+    if (!indices) return TF_ERROR;
+
+    char *path = NULL;
+    FILE *f = tf_spill_open_run_file(st->spill, "sort", &path);
+    if (!f) {
+        free(indices);
+        return TF_ERROR;
+    }
+
+    for (size_t i = 0; i < st->buf->n_rows; i++) {
+        size_t r = indices[i];
+        for (size_t c = 0; c < st->buf->n_cols; c++) {
+            if (write_cell(f, st->buf, r, c) != TF_OK) {
+                tf_set_last_error("sort spill: failed writing run file");
+                fclose(f);
+                remove(path);
+                free(path);
+                free(indices);
+                return TF_ERROR;
+            }
+        }
+    }
+    long pos = ftell(f);
+    if (pos > 0) st->spilled_bytes += (size_t)pos;
+    if (fclose(f) != 0) {
+        tf_set_last_error("sort spill: failed closing run file");
+        remove(path);
+        free(path);
+        free(indices);
+        return TF_ERROR;
+    }
+
+    free(indices);
+    if (append_run_path(st, path) != TF_OK) {
+        remove(path);
+        free(path);
+        return TF_ERROR;
+    }
 
-    /* Resolve sort column indices */
-    int *col_indices = malloc(st->n_cols * sizeof(int));
-    int *col_desc = malloc(st->n_cols * sizeof(int));
-    if (!col_indices || !col_desc) { free(col_indices); free(col_desc); return TF_ERROR; }
+    tf_batch_free(st->buf);
+    st->buf = create_buffer_from_schema(st, st->run_rows);
+    return st->buf ? TF_OK : TF_ERROR;
+}
+
+static void spill_row_clear(sort_spill_row *row, const tf_type *types, size_t n_cols) {
+    if (!row || !row->cells || !row->nulls) return;
+    for (size_t c = 0; c < n_cols; c++) {
+        if (!row->nulls[c] && types[c] == TF_TYPE_STRING) {
+            free(row->cells[c].str);
+            row->cells[c].str = NULL;
+        }
+        row->nulls[c] = 1;
+    }
+}
+
+static int spill_row_init(sort_spill_row *row, size_t n_cols) {
+    row->nulls = tf_callocarray_checked(n_cols ? n_cols : 1, sizeof(uint8_t));
+    row->cells = tf_callocarray_checked(n_cols ? n_cols : 1, sizeof(sort_cell));
+    if (!row->nulls || !row->cells) {
+        free(row->nulls);
+        free(row->cells);
+        row->nulls = NULL;
+        row->cells = NULL;
+        return TF_ERROR;
+    }
+    for (size_t c = 0; c < n_cols; c++) row->nulls[c] = 1;
+    return TF_OK;
+}
+
+static void spill_row_free(sort_spill_row *row, const tf_type *types, size_t n_cols) {
+    if (!row) return;
+    spill_row_clear(row, types, n_cols);
+    free(row->nulls);
+    free(row->cells);
+    row->nulls = NULL;
+    row->cells = NULL;
+}
+
+static int read_cell_value(FILE *f, sort_spill_row *row, const tf_type *types, size_t c) {
+    switch (types[c]) {
+        case TF_TYPE_BOOL: {
+            uint8_t v = 0;
+            if (read_exact(f, &v, sizeof(v)) != TF_OK) return TF_ERROR;
+            row->cells[c].b = v;
+            return TF_OK;
+        }
+        case TF_TYPE_INT64: {
+            int64_t v = 0;
+            if (read_exact(f, &v, sizeof(v)) != TF_OK) return TF_ERROR;
+            row->cells[c].i64 = v;
+            return TF_OK;
+        }
+        case TF_TYPE_FLOAT64: {
+            double v = 0.0;
+            if (read_exact(f, &v, sizeof(v)) != TF_OK) return TF_ERROR;
+            row->cells[c].f64 = v;
+            return TF_OK;
+        }
+        case TF_TYPE_STRING: {
+            uint64_t len = 0;
+            if (read_exact(f, &len, sizeof(len)) != TF_OK) return TF_ERROR;
+            if (len > (uint64_t)SIZE_MAX - 1) return TF_ERROR;
+            char *s = malloc((size_t)len + 1);
+            if (!s) return TF_ERROR;
+            if (len && read_exact(f, s, (size_t)len) != TF_OK) {
+                free(s);
+                return TF_ERROR;
+            }
+            s[len] = '\0';
+            row->cells[c].str = s;
+            return TF_OK;
+        }
+        case TF_TYPE_DATE: {
+            int32_t v = 0;
+            if (read_exact(f, &v, sizeof(v)) != TF_OK) return TF_ERROR;
+            row->cells[c].date = v;
+            return TF_OK;
+        }
+        case TF_TYPE_TIMESTAMP: {
+            int64_t v = 0;
+            if (read_exact(f, &v, sizeof(v)) != TF_OK) return TF_ERROR;
+            row->cells[c].i64 = v;
+            return TF_OK;
+        }
+        default:
+            return TF_OK;
+    }
+}
 
+static int reader_advance(sort_state *st, sort_run_reader *reader) {
+    if (!reader || !reader->file || reader->done) return 0;
+    spill_row_clear(&reader->row, st->schema_types, st->n_schema_cols);
+
+    int first = fgetc(reader->file);
+    if (first == EOF) {
+        if (ferror(reader->file)) {
+            tf_set_last_error("sort spill: failed reading run file");
+            return -1;
+        }
+        reader->done = 1;
+        reader->has_row = 0;
+        return 0;
+    }
+
+    reader->row.nulls[0] = first ? 1 : 0;
+    if (!reader->row.nulls[0] && read_cell_value(reader->file, &reader->row, st->schema_types, 0) != TF_OK) {
+        tf_set_last_error("sort spill: corrupt run file");
+        return -1;
+    }
+    for (size_t c = 1; c < st->n_schema_cols; c++) {
+        uint8_t is_null = 1;
+        if (read_exact(reader->file, &is_null, sizeof(is_null)) != TF_OK) {
+            tf_set_last_error("sort spill: corrupt run file");
+            return -1;
+        }
+        reader->row.nulls[c] = is_null ? 1 : 0;
+        if (!reader->row.nulls[c] && read_cell_value(reader->file, &reader->row, st->schema_types, c) != TF_OK) {
+            tf_set_last_error("sort spill: corrupt run file");
+            return -1;
+        }
+    }
+    reader->has_row = 1;
+    return 1;
+}
+
+static int compare_spill_rows(const sort_state *st,
+                              const sort_spill_row *a,
+                              const sort_spill_row *b) {
     for (size_t k = 0; k < st->n_cols; k++) {
-        col_indices[k] = tf_batch_col_index(st->buf, st->cols[k].name);
-        col_desc[k] = st->cols[k].desc;
+        int ci = st->col_indices[k];
+        if (ci < 0) continue;
+        size_t c = (size_t)ci;
+        int null_a = a->nulls[c] != 0;
+        int null_b = b->nulls[c] != 0;
+        if (null_a && null_b) continue;
+        if (null_a) return 1;
+        if (null_b) return -1;
+
+        int cmp = 0;
+        switch (st->schema_types[c]) {
+            case TF_TYPE_BOOL:
+                cmp = (int)a->cells[c].b - (int)b->cells[c].b;
+                break;
+            case TF_TYPE_INT64:
+            case TF_TYPE_TIMESTAMP:
+                cmp = (a->cells[c].i64 > b->cells[c].i64) -
+                      (a->cells[c].i64 < b->cells[c].i64);
+                break;
+            case TF_TYPE_FLOAT64:
+                cmp = (a->cells[c].f64 > b->cells[c].f64) -
+                      (a->cells[c].f64 < b->cells[c].f64);
+                break;
+            case TF_TYPE_STRING:
+                cmp = strcmp(a->cells[c].str, b->cells[c].str);
+                break;
+            case TF_TYPE_DATE:
+                cmp = (a->cells[c].date > b->cells[c].date) -
+                      (a->cells[c].date < b->cells[c].date);
+                break;
+            default:
+                break;
+        }
+        if (cmp != 0) return st->col_desc[k] ? -cmp : cmp;
     }
+    return 0;
+}
 
-    /* Build index array */
-    size_t *indices = malloc(n * sizeof(size_t));
-    if (!indices) { free(col_indices); free(col_desc); return TF_ERROR; }
-    for (size_t i = 0; i < n; i++) indices[i] = i;
+static int spill_row_to_batch(const sort_state *st, tf_batch *out,
+                              size_t dst_row, const sort_spill_row *row) {
+    if (tf_batch_ensure_capacity(out, dst_row + 1) != TF_OK) return TF_ERROR;
+    for (size_t c = 0; c < st->n_schema_cols; c++) {
+        if (tf_batch_set_owned_cell_value(out, dst_row, c, st->schema_types[c],
+                                          row->nulls[c], &row->cells[c]) != TF_OK) {
+            return TF_ERROR;
+        }
+    }
+    return TF_OK;
+}
 
-    /* Sort */
-    sort_ctx ctx = {
-        .batch = st->buf,
-        .col_indices = col_indices,
-        .col_desc = col_desc,
-        .n_sort_cols = st->n_cols,
-    };
-    g_sort_ctx = &ctx;
-    qsort(indices, n, sizeof(size_t), compare_rows);
-    g_sort_ctx = NULL;
+static void close_readers(sort_state *st) {
+    if (!st->readers) return;
+    for (size_t i = 0; i < st->n_readers; i++) {
+        if (st->readers[i].file) fclose(st->readers[i].file);
+        spill_row_free(&st->readers[i].row, st->schema_types, st->n_schema_cols);
+    }
+    free(st->readers);
+    st->readers = NULL;
+    st->n_readers = 0;
+}
 
-    /* Build output batch in sorted order */
-    tf_batch *ob = tf_batch_create(st->buf->n_cols, n);
-    if (!ob) { free(indices); free(col_indices); free(col_desc); return TF_ERROR; }
-    for (size_t c = 0; c < st->buf->n_cols; c++) {
-        tf_batch_set_schema(ob, c, st->buf->col_names[c], st->buf->col_types[c]);
+static void remove_run_files(sort_state *st) {
+    for (size_t i = 0; i < st->n_runs; i++) {
+        if (st->run_paths[i]) {
+            remove(st->run_paths[i]);
+            free(st->run_paths[i]);
+            st->run_paths[i] = NULL;
+        }
     }
+    free(st->run_paths);
+    st->run_paths = NULL;
+    st->n_runs = 0;
+    st->cap_runs = 0;
+}
 
-    for (size_t i = 0; i < n; i++) {
-        if (copy_row(ob, i, st->buf, indices[i]) != TF_OK) {
-            free(indices); free(col_indices); free(col_desc);
+static int begin_merge(sort_state *st) {
+    if (st->merge_started) return TF_OK;
+    if (st->buf && st->buf->n_rows > 0 && write_spill_run(st) != TF_OK) return TF_ERROR;
+    if (st->buf) {
+        tf_batch_free(st->buf);
+        st->buf = NULL;
+    }
+
+    st->merge_started = 1;
+    if (st->n_runs == 0) {
+        st->merge_done = 1;
+        return TF_OK;
+    }
+
+    st->readers = tf_callocarray_checked(st->n_runs, sizeof(sort_run_reader));
+    if (!st->readers) return TF_ERROR;
+    st->n_readers = st->n_runs;
+    for (size_t i = 0; i < st->n_runs; i++) {
+        st->readers[i].file = fopen(st->run_paths[i], "rb");
+        if (!st->readers[i].file) {
+            tf_set_last_error("sort spill: cannot reopen run file");
+            return TF_ERROR;
+        }
+        if (spill_row_init(&st->readers[i].row, st->n_schema_cols) != TF_OK) return TF_ERROR;
+        int rc = reader_advance(st, &st->readers[i]);
+        if (rc < 0) return TF_ERROR;
+    }
+    return TF_OK;
+}
+
+static int next_best_reader(const sort_state *st) {
+    int best = -1;
+    for (size_t i = 0; i < st->n_readers; i++) {
+        const sort_run_reader *r = &st->readers[i];
+        if (!r->has_row || r->done) continue;
+        if (best < 0) {
+            best = (int)i;
+            continue;
+        }
+        int cmp = compare_spill_rows(st, &r->row, &st->readers[best].row);
+        if (cmp < 0 || (cmp == 0 && i < (size_t)best)) best = (int)i;
+    }
+    return best;
+}
+
+static int spill_next_batch(sort_state *st, tf_batch **out) {
+    *out = NULL;
+    if (begin_merge(st) != TF_OK) return TF_ERROR;
+    if (st->merge_done) return TF_OK;
+
+    tf_batch *ob = create_buffer_from_schema(st, st->output_batch_rows);
+    if (!ob) return TF_ERROR;
+
+    while (ob->n_rows < st->output_batch_rows) {
+        int best = next_best_reader(st);
+        if (best < 0) break;
+        size_t out_row = ob->n_rows;
+        if (spill_row_to_batch(st, ob, out_row, &st->readers[best].row) != TF_OK) {
+            tf_batch_free(ob);
+            return TF_ERROR;
+        }
+        if (tf_batch_expose_row(ob, out_row) != TF_OK) {
+            tf_batch_free(ob);
+            return TF_ERROR;
+        }
+        int rc = reader_advance(st, &st->readers[best]);
+        if (rc < 0) {
             tf_batch_free(ob);
             return TF_ERROR;
         }
-        ob->n_rows = i + 1;
     }
 
-    free(indices);
-    free(col_indices);
-    free(col_desc);
+    if (ob->n_rows == 0) {
+        tf_batch_free(ob);
+        close_readers(st);
+        remove_run_files(st);
+        tf_spill_cleanup(st->spill);
+        st->spill = NULL;
+        st->merge_done = 1;
+        return TF_OK;
+    }
 
+    st->spill_output_batches++;
+    st->spill_output_rows += ob->n_rows;
     *out = ob;
     return TF_OK;
 }
 
-static void sort_destroy(tf_step *self) {
+static int sort_process(tf_step *self, tf_batch *in, tf_batch **out,
+                        tf_side_channels *side) {
+    (void)side;
+    sort_state *st = self->state;
+    *out = NULL;
+
+    if (init_schema(st, in) != TF_OK) return TF_ERROR;
+    for (size_t r = 0; r < in->n_rows; r++) {
+        size_t dst_row = st->buf->n_rows;
+        if (tf_batch_copy_row(st->buf, dst_row, in, r) != TF_OK) return TF_ERROR;
+        if (tf_batch_expose_row(st->buf, dst_row) != TF_OK) return TF_ERROR;
+        if (st->use_spill && st->buf->n_rows >= st->run_rows) {
+            if (write_spill_run(st) != TF_OK) return TF_ERROR;
+        }
+    }
+    return TF_OK;
+}
+
+static int sort_flush_in_memory(tf_step *self, tf_batch **out) {
     sort_state *st = self->state;
-    if (st) {
-        if (st->buf) tf_batch_free(st->buf);
-        for (size_t i = 0; i < st->n_cols; i++) free(st->cols[i].name);
-        free(st->cols);
-        free(st);
+    *out = NULL;
+    if (!st->buf || st->buf->n_rows == 0) return TF_OK;
+
+    size_t n = st->buf->n_rows;
+    size_t *indices = sort_batch_indices(st, st->buf);
+    if (!indices) return TF_ERROR;
+
+    tf_batch *ob = create_buffer_from_schema(st, n);
+    if (!ob) { free(indices); return TF_ERROR; }
+    for (size_t i = 0; i < n; i++) {
+        if (tf_batch_copy_row(ob, i, st->buf, indices[i]) != TF_OK) {
+            free(indices);
+            tf_batch_free(ob);
+            return TF_ERROR;
+        }
+        if (tf_batch_expose_row(ob, i) != TF_OK) {
+            free(indices);
+            tf_batch_free(ob);
+            return TF_ERROR;
+        }
     }
+
+    free(indices);
+    *out = ob;
+    return TF_OK;
+}
+
+static int sort_flush(tf_step *self, tf_batch **out, tf_side_channels *side) {
+    (void)side;
+    sort_state *st = self->state;
+    *out = NULL;
+    if (!st->use_spill) return sort_flush_in_memory(self, out);
+    return spill_next_batch(st, out);
+}
+
+static int sort_flush_next(tf_step *self, tf_batch **out, tf_side_channels *side) {
+    (void)side;
+    sort_state *st = self->state;
+    *out = NULL;
+    if (!st->use_spill) return TF_OK;
+    return spill_next_batch(st, out);
+}
+
+static int sort_append_stats(tf_step *self, tf_buffer *out) {
+    if (!self || !out) return TF_ERROR;
+    sort_state *st = self->state;
+    if (!st || !st->use_spill) return TF_OK;
+
+    char buf[256];
+    snprintf(buf, sizeof(buf),
+             ",\"spill_bytes\":%zu,\"spill_runs\":%zu,"
+             "\"spill_output_batches\":%zu,\"spill_output_rows\":%zu",
+             st->spilled_bytes, st->spill_runs_created,
+             st->spill_output_batches, st->spill_output_rows);
+    return tf_buffer_write_str(out, buf);
+}
+
+static void sort_state_free(sort_state *st) {
+    if (!st) return;
+    if (st->buf) tf_batch_free(st->buf);
+    close_readers(st);
+    remove_run_files(st);
+    for (size_t i = 0; i < st->n_cols; i++) free(st->cols[i].name);
+    for (size_t i = 0; i < st->n_schema_cols; i++) free(st->schema_names ? st->schema_names[i] : NULL);
+    free(st->schema_names);
+    free(st->schema_types);
+    free(st->cols);
+    free(st->col_indices);
+    free(st->col_desc);
+    tf_spill_cleanup(st->spill);
+    free(st->spill_dir);
+    free(st);
+}
+
+static void sort_destroy(tf_step *self) {
+    if (self) sort_state_free(self->state);
     free(self);
 }
 
@@ -242,31 +752,67 @@ tf_step *tf_sort_create(const cJSON *args) {
 
     sort_state *st = calloc(1, sizeof(sort_state));
     if (!st) return NULL;
-    st->cols = calloc(n, sizeof(sort_col));
+    st->cols = tf_callocarray_checked((size_t)n, sizeof(sort_col));
     if (!st->cols) { free(st); return NULL; }
-    st->n_cols = n;
+    st->n_cols = (size_t)n;
+    st->output_batch_rows = SORT_DEFAULT_OUTPUT_ROWS;
 
     for (int i = 0; i < n; i++) {
         cJSON *item = cJSON_GetArrayItem(columns, i);
         cJSON *name_j = cJSON_GetObjectItemCaseSensitive(item, "name");
         cJSON *desc_j = cJSON_GetObjectItemCaseSensitive(item, "desc");
-        if (!cJSON_IsString(name_j)) {
-            for (int j = 0; j < i; j++) free(st->cols[j].name);
-            free(st->cols); free(st);
+        if (!cJSON_IsString(name_j) || !name_j->valuestring || name_j->valuestring[0] == '\0') {
+            tf_set_last_error("sort: column names must be non-empty strings");
+            sort_state_free(st);
             return NULL;
         }
         st->cols[i].name = strdup(name_j->valuestring);
+        if (!st->cols[i].name) {
+            sort_state_free(st);
+            return NULL;
+        }
         st->cols[i].desc = (desc_j && cJSON_IsBool(desc_j)) ? cJSON_IsTrue(desc_j) : 0;
     }
 
-    tf_step *step = malloc(sizeof(tf_step));
+    cJSON *spill_dir_j = cJSON_GetObjectItemCaseSensitive(args, "spill_dir");
+    if (cJSON_IsString(spill_dir_j) && spill_dir_j->valuestring && spill_dir_j->valuestring[0]) {
+        st->use_spill = 1;
+        st->spill_dir = strdup(spill_dir_j->valuestring);
+        if (!st->spill_dir) {
+            sort_state_free(st);
+            return NULL;
+        }
+        if (tf_spill_session_create(st->spill_dir, &st->spill) != TF_OK) {
+            sort_state_free(st);
+            return NULL;
+        }
+        size_t parsed_size = 0;
+        int has_spill_memory = tf_json_get_size_arg(args, "spill_memory_bytes",
+                                                    1, TF_MAX_SPILL_MEMORY_BYTES,
+                                                    &parsed_size, "sort");
+        if (has_spill_memory < 0) { sort_state_free(st); return NULL; }
+        if (has_spill_memory > 0) st->spill_memory_bytes = parsed_size;
+        int has_spill_rows = tf_json_get_size_arg(args, "spill_run_rows",
+                                                  1, TF_MAX_SPILL_RUN_ROWS,
+                                                  &parsed_size, "sort");
+        if (has_spill_rows < 0) { sort_state_free(st); return NULL; }
+        if (has_spill_rows > 0) st->configured_run_rows = parsed_size;
+        int has_output_rows = tf_json_get_size_arg(args, "spill_output_rows",
+                                                   1, TF_MAX_SPILL_OUTPUT_ROWS,
+                                                   &parsed_size, "sort");
+        if (has_output_rows < 0) { sort_state_free(st); return NULL; }
+        if (has_output_rows > 0) st->output_batch_rows = parsed_size;
+    }
+
+    tf_step *step = calloc(1, sizeof(tf_step));
     if (!step) {
-        for (int i = 0; i < n; i++) free(st->cols[i].name);
-        free(st->cols); free(st);
+        sort_state_free(st);
         return NULL;
     }
     step->process = sort_process;
     step->flush = sort_flush;
+    step->flush_next = st->use_spill ? sort_flush_next : NULL;
+    step->append_stats = sort_append_stats;
     step->destroy = sort_destroy;
     step->state = st;
     return step;