tranfi 0.1.2 → 0.2.0

package/csrc/op_lead.c

@@ -1,10 +1,11 @@
 /*
- * op_lead.c — Lookahead: access value N rows ahead.
+ * op_lead.c - Bounded lookahead.
  *
  * Config: {"column": "price", "offset": 1, "result": "next_price"}
  *
- * Buffers the last `offset` rows across batch boundaries.
- * On flush, emits remaining rows with NULL lead values.
+ * Delays output by `offset` rows so the appended value can come from rows ahead
+ * of the current row. On flush, emits the remaining pending rows with NULL lead
+ * values. Pending state is bounded by `offset` rows.
  */
 
 #include "internal.h"
@@ -14,136 +15,178 @@
 #include <stdio.h>
 
 typedef struct {
-    char   *column;
-    char   *result;
-    size_t  offset;
-    /* Pending rows: stored as a batch that hasn't been emitted yet */
+    char     *column;
+    char     *result;
+    size_t    offset;
     tf_batch *pending;
 } lead_state;
 
-static double get_numeric(const tf_batch *b, size_t r, int ci) {
-    if (b->col_types[ci] == TF_TYPE_INT64) return (double)tf_batch_get_int64(b, r, ci);
-    if (b->col_types[ci] == TF_TYPE_FLOAT64) return tf_batch_get_float64(b, r, ci);
-    return 0;
+static int lead_set_error(tf_side_channels *side, const char *msg) {
+    return tf_side_write_error(side, msg);
+}
+
+static const tf_batch *source_at(const lead_state *st, const tf_batch *in,
+                                 size_t idx, size_t pend_count, size_t *row) {
+    if (idx < pend_count) {
+        *row = idx;
+        return st->pending;
+    }
+    *row = idx - pend_count;
+    return in;
+}
+
+static int copy_source_row(tf_batch *dst, size_t dst_row, const lead_state *st,
+                           const tf_batch *in, size_t idx, size_t pend_count) {
+    size_t src_row = 0;
+    const tf_batch *src = source_at(st, in, idx, pend_count, &src_row);
+    return tf_batch_copy_row(dst, dst_row, src, src_row);
 }
 
 static int lead_process(tf_step *self, tf_batch *in, tf_batch **out,
                         tf_side_channels *side) {
-    (void)side;
     lead_state *st = self->state;
     *out = NULL;
 
+    int ci = tf_batch_col_index(in, st->column);
+    if (ci < 0) {
+        char msg[256];
+        snprintf(msg, sizeof(msg), "lead: column '%s' not found", st->column);
+        if (lead_set_error(side, msg) != TF_OK) return TF_ERROR;
+        return TF_ERROR;
+    }
+
     size_t pend_count = st->pending ? st->pending->n_rows : 0;
     size_t total = pend_count + in->n_rows;
 
     if (total <= st->offset) {
-        /* Not enough rows yet — append all to pending */
         tf_batch *new_pend = tf_batch_create(in->n_cols, total);
         if (!new_pend) return TF_ERROR;
-        for (size_t c = 0; c < in->n_cols; c++)
-            tf_batch_set_schema(new_pend, c, in->col_names[c], in->col_types[c]);
+        if (tf_batch_clone_schema(new_pend, in) != TF_OK) {
+            tf_batch_free(new_pend);
+            return TF_ERROR;
+        }
         size_t row = 0;
-        if (st->pending) {
-            for (size_t r = 0; r < pend_count; r++)
-                tf_batch_copy_row(new_pend, row++, st->pending, r);
-            tf_batch_free(st->pending);
+        for (size_t r = 0; r < pend_count; r++) {
+            if (tf_batch_copy_row(new_pend, row++, st->pending, r) != TF_OK) {
+                tf_batch_free(new_pend);
+                return TF_ERROR;
+            }
+        }
+        for (size_t r = 0; r < in->n_rows; r++) {
+            if (tf_batch_copy_row(new_pend, row++, in, r) != TF_OK) {
+                tf_batch_free(new_pend);
+                return TF_ERROR;
+            }
+        }
+        if (total > 0 && tf_batch_expose_row(new_pend, total - 1) != TF_OK) {
+            tf_batch_free(new_pend);
+            return TF_ERROR;
         }
-        for (size_t r = 0; r < in->n_rows; r++)
-            tf_batch_copy_row(new_pend, row++, in, r);
-        new_pend->n_rows = total;
+        if (st->pending) tf_batch_free(st->pending);
         st->pending = new_pend;
         return TF_OK;
     }
 
     size_t emit_count = total - st->offset;
+    const char *extra_names[1] = {st->result};
+    tf_type extra_types[1] = {in->col_types[ci]};
     tf_batch *ob = tf_batch_create(in->n_cols + 1, emit_count);
     if (!ob) return TF_ERROR;
-    for (size_t c = 0; c < in->n_cols; c++)
-        tf_batch_set_schema(ob, c, in->col_names[c], in->col_types[c]);
-    tf_batch_set_schema(ob, in->n_cols, st->result, TF_TYPE_FLOAT64);
+    if (tf_batch_clone_with_extra_cols(ob, in, extra_names, extra_types, 1) != TF_OK) {
+        tf_batch_free(ob);
+        return TF_ERROR;
+    }
 
-    /* Emit rows: for each emitted row i, lead value comes from row i+offset */
     for (size_t i = 0; i < emit_count; i++) {
-        /* Source row for base data */
-        if (i < pend_count) {
-            tf_batch_copy_row(ob, i, st->pending, i);
-        } else {
-            tf_batch_copy_row(ob, i, in, i - pend_count);
-        }
-
-        /* Lead value from row i + offset */
-        size_t lead_idx = i + st->offset;
-        tf_batch *lead_src;
-        size_t lead_row;
-        if (lead_idx < pend_count) {
-            lead_src = st->pending;
-            lead_row = lead_idx;
-        } else {
-            lead_src = in;
-            lead_row = lead_idx - pend_count;
+        if (copy_source_row(ob, i, st, in, i, pend_count) != TF_OK) {
+            tf_batch_free(ob);
+            return TF_ERROR;
         }
 
+        size_t lead_row = 0;
+        const tf_batch *lead_src = source_at(st, in, i + st->offset, pend_count, &lead_row);
         int lead_ci = tf_batch_col_index(lead_src, st->column);
-        if (lead_ci >= 0 && !tf_batch_is_null(lead_src, lead_row, lead_ci)) {
-            tf_batch_set_float64(ob, i, in->n_cols, get_numeric(lead_src, lead_row, lead_ci));
-        } else {
-            tf_batch_set_null(ob, i, in->n_cols);
+        if (lead_ci < 0) {
+            tf_batch_free(ob);
+            if (lead_set_error(side, "lead: buffered schema lost selected column") != TF_OK) return TF_ERROR;
+            return TF_ERROR;
+        }
+        if (tf_batch_copy_cell(ob, i, in->n_cols, lead_src, lead_row, (size_t)lead_ci) != TF_OK) {
+            tf_batch_free(ob);
+            return TF_ERROR;
+        }
+        if (tf_batch_expose_row(ob, i) != TF_OK) {
+            tf_batch_free(ob);
+            return TF_ERROR;
         }
     }
-    ob->n_rows = emit_count;
 
-    /* Store remaining rows as new pending */
-    if (st->pending) {
-        tf_batch_free(st->pending);
-        st->pending = NULL;
+    tf_batch *new_pend = tf_batch_create(in->n_cols, st->offset);
+    if (!new_pend) { tf_batch_free(ob); return TF_ERROR; }
+    if (tf_batch_clone_schema(new_pend, in) != TF_OK) {
+        tf_batch_free(new_pend);
+        tf_batch_free(ob);
+        return TF_ERROR;
     }
-    size_t new_pend_count = st->offset;
-    if (new_pend_count > 0) {
-        tf_batch *new_pend = tf_batch_create(in->n_cols, new_pend_count);
-        if (!new_pend) { tf_batch_free(ob); return TF_ERROR; }
-        for (size_t c = 0; c < in->n_cols; c++)
-            tf_batch_set_schema(new_pend, c, in->col_names[c], in->col_types[c]);
-        for (size_t i = 0; i < new_pend_count; i++) {
-            size_t src_idx = total - st->offset + i;
-            if (src_idx < pend_count) {
-                /* This shouldn't happen since we emit at least pend_count rows when total > offset */
-                tf_batch_copy_row(new_pend, i, st->pending, src_idx);
-            } else {
-                tf_batch_copy_row(new_pend, i, in, src_idx - pend_count);
-            }
+    for (size_t i = 0; i < st->offset; i++) {
+        size_t src_idx = total - st->offset + i;
+        if (copy_source_row(new_pend, i, st, in, src_idx, pend_count) != TF_OK) {
+            tf_batch_free(new_pend);
+            tf_batch_free(ob);
+            return TF_ERROR;
+        }
+        if (tf_batch_expose_row(new_pend, i) != TF_OK) {
+            tf_batch_free(new_pend);
+            tf_batch_free(ob);
+            return TF_ERROR;
         }
-        new_pend->n_rows = new_pend_count;
-        st->pending = new_pend;
     }
 
+    if (st->pending) tf_batch_free(st->pending);
+    st->pending = new_pend;
     *out = ob;
     return TF_OK;
 }
 
 static int lead_flush(tf_step *self, tf_batch **out, tf_side_channels *side) {
-    (void)side;
     lead_state *st = self->state;
     *out = NULL;
 
     if (!st->pending || st->pending->n_rows == 0) return TF_OK;
 
-    /* Emit remaining pending rows with NULL lead values */
     tf_batch *pend = st->pending;
+    int ci = tf_batch_col_index(pend, st->column);
+    if (ci < 0) {
+        if (lead_set_error(side, "lead: buffered schema lost selected column") != TF_OK) return TF_ERROR;
+        return TF_ERROR;
+    }
+
+    const char *extra_names[1] = {st->result};
+    tf_type extra_types[1] = {pend->col_types[ci]};
     tf_batch *ob = tf_batch_create(pend->n_cols + 1, pend->n_rows);
     if (!ob) return TF_ERROR;
-    for (size_t c = 0; c < pend->n_cols; c++)
-        tf_batch_set_schema(ob, c, pend->col_names[c], pend->col_types[c]);
-    tf_batch_set_schema(ob, pend->n_cols, st->result, TF_TYPE_FLOAT64);
+    if (tf_batch_clone_with_extra_cols(ob, pend, extra_names, extra_types, 1) != TF_OK) {
+        tf_batch_free(ob);
+        return TF_ERROR;
+    }
 
     for (size_t r = 0; r < pend->n_rows; r++) {
-        tf_batch_copy_row(ob, r, pend, r);
-        tf_batch_set_null(ob, r, pend->n_cols);
+        if (tf_batch_copy_row(ob, r, pend, r) != TF_OK) {
+            tf_batch_free(ob);
+            return TF_ERROR;
+        }
+        if (tf_batch_set_null(ob, r, pend->n_cols) != TF_OK) {
+            tf_batch_free(ob);
+            return TF_ERROR;
+        }
+        if (tf_batch_expose_row(ob, r) != TF_OK) {
+            tf_batch_free(ob);
+            return TF_ERROR;
+        }
     }
-    ob->n_rows = pend->n_rows;
 
     tf_batch_free(st->pending);
     st->pending = NULL;
-
     *out = ob;
     return TF_OK;
 }
@@ -162,26 +205,35 @@ static void lead_destroy(tf_step *self) {
 tf_step *tf_lead_create(const cJSON *args) {
     if (!args) return NULL;
     cJSON *col_j = cJSON_GetObjectItemCaseSensitive(args, "column");
-    if (!cJSON_IsString(col_j)) return NULL;
+    if (!cJSON_IsString(col_j) || !col_j->valuestring[0]) return NULL;
 
-    lead_state *st = calloc(1, sizeof(lead_state));
+    lead_state *st = tf_callocarray_checked(1, sizeof(lead_state));
     if (!st) return NULL;
-    st->column = strdup(col_j->valuestring);
+    st->column = tf_strdup_checked(col_j->valuestring);
+    if (!st->column) { free(st); return NULL; }
 
-    cJSON *off_j = cJSON_GetObjectItemCaseSensitive(args, "offset");
-    st->offset = (cJSON_IsNumber(off_j) && off_j->valueint > 0) ? (size_t)off_j->valueint : 1;
+    size_t offset = 1;
+    int has_offset = tf_json_get_size_arg(args, "offset",
+                                          1, TF_MAX_WINDOW_SIZE,
+                                          &offset, "lead");
+    if (has_offset < 0) { free(st->column); free(st); return NULL; }
+    st->offset = offset;
 
     cJSON *res_j = cJSON_GetObjectItemCaseSensitive(args, "result");
-    if (cJSON_IsString(res_j)) {
-        st->result = strdup(res_j->valuestring);
+    if (cJSON_IsString(res_j) && res_j->valuestring[0]) {
+        st->result = tf_strdup_checked(res_j->valuestring);
     } else {
-        char buf[256];
-        snprintf(buf, sizeof(buf), "%s_lead", st->column);
-        st->result = strdup(buf);
+        st->result = tf_string_append_suffix_checked(st->column, "_lead");
     }
+    if (!st->result) { free(st->column); free(st); return NULL; }
 
-    tf_step *step = malloc(sizeof(tf_step));
-    if (!step) { free(st->column); free(st->result); free(st); return NULL; }
+    tf_step *step = tf_callocarray_checked(1, sizeof(tf_step));
+    if (!step) {
+        free(st->column);
+        free(st->result);
+        free(st);
+        return NULL;
+    }
     step->process = lead_process;
     step->flush = lead_flush;
     step->destroy = lead_destroy;