tranfi 0.1.2 → 0.2.1

package/csrc/op_onehot.c

@@ -2,7 +2,11 @@
  * op_onehot.c — One-hot encoding of a categorical column.
  * Expands a single column into N binary (0/1) columns.
  *
- * Config: {"column": "city", "drop": false}
+ * Config: {"column": "city", "drop": false,
+ *          "categories": ["Paris", "London"],
+ *          "max_categories": 1000,
+ *          "max_state_bytes": 1048576,
+ *          "unknown": "error" | "other" | "null"}
  */
 
 #include "internal.h"
@@ -11,19 +15,68 @@
 #include <string.h>
 #include <stdio.h>
 
+typedef enum {
+    TF_CAT_UNKNOWN_ADD,
+    TF_CAT_UNKNOWN_ERROR,
+    TF_CAT_UNKNOWN_OTHER,
+    TF_CAT_UNKNOWN_NULL
+} tf_cat_unknown_policy;
+
 typedef struct {
     char  *value;     /* category value string */
     char  *col_name;  /* generated column name: "column_value" */
 } onehot_category;
 
 typedef struct {
-    char             *column;
-    int               drop;    /* drop original column */
-    onehot_category  *cats;
-    size_t            n_cats;
-    size_t            cap;
+    char                  *column;
+    int                    drop;    /* drop original column */
+    size_t                 max_categories; /* 0 = unlimited */
+    size_t                 max_state_bytes; /* 0 = unlimited */
+    int                    categories_declared;
+    tf_cat_unknown_policy  unknown;
+    onehot_category       *cats;
+    size_t                 n_cats;
+    size_t                 cap;
 } onehot_state;
 
+static const char *OTHER_CATEGORY = "__other__";
+
+static size_t onehot_retained_state_bytes(const onehot_state *st);
+
+static int onehot_write_error(tf_side_channels *side, const char *msg) {
+    return tf_side_write_error(side, msg);
+}
+
+static int onehot_unknown_error(const onehot_state *st, const char *val,
+                                tf_side_channels *side) {
+    char msg[256];
+    snprintf(msg, sizeof(msg),
+             "onehot: unknown category '%s' for column '%s'",
+             val ? val : "", st->column ? st->column : "");
+    return onehot_write_error(side, msg);
+}
+
+static int onehot_limit_error(const onehot_state *st, const char *val,
+                              tf_side_channels *side) {
+    char msg[256];
+    snprintf(msg, sizeof(msg),
+             "onehot: max_categories=%zu exceeded while tracking category '%s'",
+             st->max_categories, val ? val : "");
+    return onehot_write_error(side, msg);
+}
+
+static int onehot_check_state_bytes(onehot_state *st, tf_side_channels *side) {
+    if (!st || st->max_state_bytes == 0) return TF_OK;
+    size_t retained = onehot_retained_state_bytes(st);
+    if (retained <= st->max_state_bytes) return TF_OK;
+    char msg[192];
+    snprintf(msg, sizeof(msg),
+             "onehot: max_state_bytes=%zu exceeded while tracking categories (%zu bytes retained)",
+             st->max_state_bytes, retained);
+    if (onehot_write_error(side, msg) != TF_OK) return TF_ERROR;
+    return TF_ERROR;
+}
+
 static const char *get_string_value(const tf_batch *b, size_t r, int ci, char *buf, size_t bufsz) {
     if (tf_batch_is_null(b, r, ci)) return NULL;
     switch (b->col_types[ci]) {
@@ -40,127 +93,272 @@ static const char *get_string_value(const tf_batch *b, size_t r, int ci, char *b
     }
 }
 
-static int find_or_add_category(onehot_state *st, const char *val) {
+static int find_category(const onehot_state *st, const char *val) {
     for (size_t i = 0; i < st->n_cats; i++) {
-        if (strcmp(st->cats[i].value, val) == 0)
-            return (int)i;
+        if (strcmp(st->cats[i].value, val) == 0) return (int)i;
     }
-    /* Add new category */
+    return -1;
+}
+
+static int add_category(onehot_state *st, const char *val) {
+    int existing = find_category(st, val);
+    if (existing >= 0) return existing;
     if (st->n_cats >= st->cap) {
-        size_t newcap = st->cap ? st->cap * 2 : 16;
-        onehot_category *tmp = realloc(st->cats, newcap * sizeof(onehot_category));
+        size_t min_cap = 0, newcap = 0;
+        if (tf_size_add(st->n_cats, 1, &min_cap) != TF_OK ||
+            tf_size_grow_pow2(st->cap, min_cap, 16, &newcap) != TF_OK) {
+            return -1;
+        }
+        onehot_category *tmp = tf_reallocarray_checked(st->cats, newcap,
+                                                       sizeof(onehot_category));
         if (!tmp) return -1;
         st->cats = tmp;
         st->cap = newcap;
     }
-    st->cats[st->n_cats].value = strdup(val);
-    char namebuf[512];
-    snprintf(namebuf, sizeof(namebuf), "%s_%s", st->column, val);
-    st->cats[st->n_cats].col_name = strdup(namebuf);
+    char *value = strdup(val);
+    if (!value) return -1;
+    size_t column_len = 0, value_len = 0, name_len = 0, name_cap = 0;
+    if (tf_string_length_bounded(st->column, TF_MAX_COLUMN_NAME_BYTES,
+                                 &column_len, "onehot", "column") != TF_OK ||
+        tf_string_length_bounded(val, TF_MAX_COLUMN_NAME_BYTES,
+                                 &value_len, "onehot", "category value") != TF_OK ||
+        tf_size_add(column_len, 1, &name_len) != TF_OK ||
+        tf_size_add(name_len, value_len, &name_len) != TF_OK ||
+        tf_check_byte_limit(name_len, TF_MAX_COLUMN_NAME_BYTES,
+                            "onehot", "generated column name") != TF_OK ||
+        tf_size_add(name_len, 1, &name_cap) != TF_OK) {
+        free(value);
+        return -1;
+    }
+    char *col_name = tf_mallocarray_checked(name_cap, sizeof(char));
+    if (!col_name) {
+        free(value);
+        return -1;
+    }
+    memcpy(col_name, st->column, column_len);
+    col_name[column_len] = '_';
+    memcpy(col_name + column_len + 1, val, value_len);
+    col_name[name_len] = '\0';
+    st->cats[st->n_cats].value = value;
+    st->cats[st->n_cats].col_name = col_name;
     st->n_cats++;
     return (int)(st->n_cats - 1);
 }
 
+static int resolve_other(onehot_state *st, int *idx, tf_side_channels *side) {
+    int other = find_category(st, OTHER_CATEGORY);
+    if (other >= 0) { *idx = other; return 0; }
+    if (st->max_categories > 0 && st->n_cats >= st->max_categories) {
+        if (onehot_limit_error(st, OTHER_CATEGORY, side) != TF_OK) return -1;
+        return -1;
+    }
+    other = add_category(st, OTHER_CATEGORY);
+    if (other < 0) return -1;
+    if (onehot_check_state_bytes(st, side) != TF_OK) return -1;
+    *idx = other;
+    return 0;
+}
+
+static int resolve_unknown(onehot_state *st, const char *val, int *idx,
+                           tf_side_channels *side) {
+    switch (st->unknown) {
+        case TF_CAT_UNKNOWN_NULL:
+            *idx = -1;
+            return 0;
+        case TF_CAT_UNKNOWN_OTHER:
+            return resolve_other(st, idx, side);
+        case TF_CAT_UNKNOWN_ERROR:
+            if (onehot_unknown_error(st, val, side) != TF_OK) return -1;
+            return -1;
+        case TF_CAT_UNKNOWN_ADD:
+        default:
+            break;
+    }
+    return 1;
+}
+
+static int resolve_category(onehot_state *st, const char *val, int *idx,
+                            tf_side_channels *side) {
+    *idx = -1;
+    if (!val) return 0;
+
+    int existing = find_category(st, val);
+    if (existing >= 0) { *idx = existing; return 0; }
+
+    if (st->categories_declared) {
+        int rc = resolve_unknown(st, val, idx, side);
+        return rc == 0 ? 0 : -1;
+    }
+
+    if (st->max_categories > 0 && st->unknown == TF_CAT_UNKNOWN_OTHER) {
+        int other = find_category(st, OTHER_CATEGORY);
+        size_t reserve = other >= 0 ? 0 : 1;
+        if (st->n_cats + reserve >= st->max_categories) {
+            return resolve_other(st, idx, side);
+        }
+    }
+
+    if (st->max_categories > 0 && st->n_cats >= st->max_categories) {
+        int rc = resolve_unknown(st, val, idx, side);
+        if (rc == 0) return 0;
+        if (onehot_limit_error(st, val, side) != TF_OK) return -1;
+        return -1;
+    }
+
+    int added = add_category(st, val);
+    if (added < 0) return -1;
+    if (onehot_check_state_bytes(st, side) != TF_OK) return -1;
+    *idx = added;
+    return 0;
+}
+
 static int onehot_process(tf_step *self, tf_batch *in, tf_batch **out,
                           tf_side_channels *side) {
-    (void)side;
     onehot_state *st = self->state;
     *out = NULL;
 
     int ci = tf_batch_col_index(in, st->column);
+    size_t match_count = in->n_rows ? in->n_rows : 1;
+    int *matches = tf_mallocarray_checked(match_count, sizeof(int));
+    if (!matches) return TF_ERROR;
+    for (size_t r = 0; r < in->n_rows; r++) matches[r] = -1;
 
-    /* First pass: discover any new categories in this batch */
-    size_t cats_before = st->n_cats;
     if (ci >= 0) {
         char buf[64];
         for (size_t r = 0; r < in->n_rows; r++) {
             const char *val = get_string_value(in, r, ci, buf, sizeof(buf));
-            if (val) find_or_add_category(st, val);
+            if (resolve_category(st, val, &matches[r], side) != 0) {
+                free(matches);
+                return TF_ERROR;
+            }
         }
     }
-    (void)cats_before;
-
-    /* Compute output column count */
-    size_t out_cols = st->drop ? (in->n_cols - 1 + st->n_cats) : (in->n_cols + st->n_cats);
 
+    size_t input_cols = (st->drop && ci >= 0 && in->n_cols > 0) ? in->n_cols - 1 : in->n_cols;
+    size_t out_cols = 0;
+    if (tf_size_add(input_cols, st->n_cats, &out_cols) != TF_OK) {
+        free(matches);
+        return TF_ERROR;
+    }
     tf_batch *ob = tf_batch_create(out_cols, in->n_rows);
-    if (!ob) return TF_ERROR;
+    if (!ob) { free(matches); return TF_ERROR; }
 
-    /* Set schema: copy input cols (optionally skipping target), append onehot cols */
     size_t oc = 0;
     for (size_t c = 0; c < in->n_cols; c++) {
         if (st->drop && ci >= 0 && c == (size_t)ci) continue;
-        tf_batch_set_schema(ob, oc, in->col_names[c], in->col_types[c]);
+        if (tf_batch_set_schema(ob, oc, in->col_names[c], in->col_types[c]) != TF_OK) goto fail;
         oc++;
     }
     for (size_t i = 0; i < st->n_cats; i++) {
-        tf_batch_set_schema(ob, oc + i, st->cats[i].col_name, TF_TYPE_INT64);
+        if (tf_batch_set_schema(ob, oc + i, st->cats[i].col_name, TF_TYPE_INT64) != TF_OK) goto fail;
     }
 
-    /* Fill rows */
-    char buf[64];
     for (size_t r = 0; r < in->n_rows; r++) {
-        /* Copy input columns */
         oc = 0;
         for (size_t c = 0; c < in->n_cols; c++) {
             if (st->drop && ci >= 0 && c == (size_t)ci) continue;
-            if (tf_batch_is_null(in, r, c)) {
-                tf_batch_set_null(ob, r, oc);
-            } else {
-                switch (in->col_types[c]) {
-                    case TF_TYPE_STRING:
-                        tf_batch_set_string(ob, r, oc, tf_batch_get_string(in, r, c)); break;
-                    case TF_TYPE_INT64:
-                        tf_batch_set_int64(ob, r, oc, tf_batch_get_int64(in, r, c)); break;
-                    case TF_TYPE_FLOAT64:
-                        tf_batch_set_float64(ob, r, oc, tf_batch_get_float64(in, r, c)); break;
-                    case TF_TYPE_BOOL:
-                        tf_batch_set_bool(ob, r, oc, tf_batch_get_bool(in, r, c)); break;
-                    case TF_TYPE_DATE:
-                        tf_batch_set_date(ob, r, oc, tf_batch_get_date(in, r, c)); break;
-                    case TF_TYPE_TIMESTAMP:
-                        tf_batch_set_timestamp(ob, r, oc, tf_batch_get_timestamp(in, r, c)); break;
-                    default: tf_batch_set_null(ob, r, oc); break;
-                }
-            }
+            if (tf_batch_copy_cell(ob, r, oc, in, r, c) != TF_OK) goto fail;
             oc++;
         }
 
-        /* Set onehot columns */
-        const char *val = (ci >= 0) ? get_string_value(in, r, ci, buf, sizeof(buf)) : NULL;
-        int match = -1;
-        if (val) {
-            for (size_t i = 0; i < st->n_cats; i++) {
-                if (strcmp(st->cats[i].value, val) == 0) { match = (int)i; break; }
-            }
-        }
+        int match = matches[r];
         for (size_t i = 0; i < st->n_cats; i++) {
-            tf_batch_set_int64(ob, r, oc + i, (int)i == match ? 1 : 0);
+            if (tf_batch_set_int64(ob, r, oc + i, (int)i == match ? 1 : 0) != TF_OK) goto fail;
         }
-
-        ob->n_rows = r + 1;
+        if (tf_batch_expose_row(ob, r) != TF_OK) goto fail;
     }
 
+    free(matches);
     *out = ob;
     return TF_OK;
+
+fail:
+    free(matches);
+    tf_batch_free(ob);
+    return TF_ERROR;
 }
 
 static int onehot_flush(tf_step *self, tf_batch **out, tf_side_channels *side) {
     (void)self; (void)side; *out = NULL; return TF_OK;
 }
 
-static void onehot_destroy(tf_step *self) {
+static size_t onehot_category_value_bytes(const onehot_state *st) {
+    size_t total = 0;
+    if (!st) return 0;
+    for (size_t i = 0; i < st->n_cats; i++) {
+        if (st->cats[i].value) total += strlen(st->cats[i].value) + 1;
+    }
+    return total;
+}
+
+static size_t onehot_category_output_name_bytes(const onehot_state *st) {
+    size_t total = 0;
+    if (!st) return 0;
+    for (size_t i = 0; i < st->n_cats; i++) {
+        if (st->cats[i].col_name) total += strlen(st->cats[i].col_name) + 1;
+    }
+    return total;
+}
+
+static size_t onehot_retained_state_bytes(const onehot_state *st) {
+    if (!st) return 0;
+    size_t total = st->cap * sizeof(onehot_category);
+    total += onehot_category_value_bytes(st);
+    total += onehot_category_output_name_bytes(st);
+    if (st->column) total += strlen(st->column) + 1;
+    return total;
+}
+
+static int onehot_append_stats(tf_step *self, tf_buffer *out) {
+    if (!self || !self->state || !out) return TF_ERROR;
     onehot_state *st = self->state;
-    if (st) {
-        for (size_t i = 0; i < st->n_cats; i++) {
-            free(st->cats[i].value);
-            free(st->cats[i].col_name);
-        }
-        free(st->cats);
-        free(st->column);
-        free(st);
+    char buf[320];
+    snprintf(buf, sizeof(buf),
+             ",\"tracked_categories\":%zu,\"category_value_bytes\":%zu,"
+             "\"category_output_name_bytes\":%zu,\"retained_state_bytes\":%zu,"
+             "\"max_state_bytes\":%zu",
+             st->n_cats, onehot_category_value_bytes(st),
+             onehot_category_output_name_bytes(st), onehot_retained_state_bytes(st),
+             st->max_state_bytes);
+    return tf_buffer_write_str(out, buf);
+}
+
+static void onehot_state_free(onehot_state *st) {
+    if (!st) return;
+    for (size_t i = 0; i < st->n_cats; i++) {
+        free(st->cats[i].value);
+        free(st->cats[i].col_name);
+    }
+    free(st->cats);
+    free(st->column);
+    free(st);
+}
+
+static void onehot_destroy(tf_step *self) {
+    if (self) {
+        onehot_state_free(self->state);
+        free(self);
     }
-    free(self);
+}
+
+static int parse_unknown_policy(const cJSON *args, tf_cat_unknown_policy *policy,
+                                int *specified) {
+    cJSON *unknown_j = cJSON_GetObjectItemCaseSensitive(args, "unknown");
+    *specified = 0;
+    if (!unknown_j) return 0;
+    if (!cJSON_IsString(unknown_j)) {
+        tf_set_last_error("onehot: unknown must be one of error, other, null");
+        return -1;
+    }
+    *specified = 1;
+    if (strcmp(unknown_j->valuestring, "error") == 0) *policy = TF_CAT_UNKNOWN_ERROR;
+    else if (strcmp(unknown_j->valuestring, "other") == 0) *policy = TF_CAT_UNKNOWN_OTHER;
+    else if (strcmp(unknown_j->valuestring, "null") == 0) *policy = TF_CAT_UNKNOWN_NULL;
+    else {
+        tf_set_last_error("onehot: unknown must be one of error, other, null");
+        return -1;
+    }
+    return 0;
 }
 
 tf_step *tf_onehot_create(const cJSON *args) {
@@ -171,14 +369,88 @@ tf_step *tf_onehot_create(const cJSON *args) {
     onehot_state *st = calloc(1, sizeof(onehot_state));
     if (!st) return NULL;
     st->column = strdup(col_j->valuestring);
+    if (!st->column) { onehot_state_free(st); return NULL; }
+    st->unknown = TF_CAT_UNKNOWN_ADD;
 
     cJSON *drop_j = cJSON_GetObjectItemCaseSensitive(args, "drop");
     st->drop = cJSON_IsBool(drop_j) && cJSON_IsTrue(drop_j) ? 1 : 0;
 
-    tf_step *step = malloc(sizeof(tf_step));
-    if (!step) { free(st->column); free(st); return NULL; }
+    size_t parsed_size = 0;
+    int has_max_categories = tf_json_get_size_arg(args, "max_categories",
+                                                  1, TF_MAX_COUNT_ARG,
+                                                  &parsed_size, "onehot");
+    if (has_max_categories < 0) { onehot_state_free(st); return NULL; }
+    if (has_max_categories > 0) st->max_categories = parsed_size;
+
+    int has_max_state = tf_json_get_size_arg(args, "max_state_bytes",
+                                             1, TF_MAX_STATE_BYTES,
+                                             &parsed_size, "onehot");
+    if (has_max_state < 0) { onehot_state_free(st); return NULL; }
+    if (has_max_state > 0) st->max_state_bytes = parsed_size;
+
+    int unknown_specified = 0;
+    if (parse_unknown_policy(args, &st->unknown, &unknown_specified) != 0) {
+        onehot_state_free(st);
+        return NULL;
+    }
+
+    cJSON *cats_j = cJSON_GetObjectItemCaseSensitive(args, "categories");
+    if (cats_j) {
+        if (!cJSON_IsArray(cats_j)) {
+            tf_set_last_error("onehot: categories must be an array");
+            onehot_state_free(st);
+            return NULL;
+        }
+        st->categories_declared = 1;
+        cJSON *item = NULL;
+        cJSON_ArrayForEach(item, cats_j) {
+            if (!cJSON_IsString(item)) {
+                tf_set_last_error("onehot: categories must contain strings");
+                onehot_state_free(st);
+                return NULL;
+            }
+            if (find_category(st, item->valuestring) < 0 &&
+                st->max_categories > 0 && st->n_cats >= st->max_categories) {
+                tf_set_last_error("onehot: categories exceed max_categories");
+                onehot_state_free(st);
+                return NULL;
+            }
+            if (add_category(st, item->valuestring) < 0) {
+                onehot_state_free(st);
+                return NULL;
+            }
+            if (onehot_check_state_bytes(st, NULL) != TF_OK) {
+                onehot_state_free(st);
+                return NULL;
+            }
+        }
+    }
+
+    if (st->categories_declared && !unknown_specified)
+        st->unknown = TF_CAT_UNKNOWN_ERROR;
+
+    if (st->categories_declared && st->unknown == TF_CAT_UNKNOWN_OTHER &&
+        find_category(st, OTHER_CATEGORY) < 0) {
+        if (st->max_categories > 0 && st->n_cats >= st->max_categories) {
+            tf_set_last_error("onehot: max_categories leaves no room for other category");
+            onehot_state_free(st);
+            return NULL;
+        }
+        if (add_category(st, OTHER_CATEGORY) < 0) {
+            onehot_state_free(st);
+            return NULL;
+        }
+        if (onehot_check_state_bytes(st, NULL) != TF_OK) {
+            onehot_state_free(st);
+            return NULL;
+        }
+    }
+
+    tf_step *step = calloc(1, sizeof(tf_step));
+    if (!step) { onehot_state_free(st); return NULL; }
     step->process = onehot_process;
     step->flush = onehot_flush;
+    step->append_stats = onehot_append_stats;
     step->destroy = onehot_destroy;
     step->state = st;
     return step;