diff --git a/backlog/drafts/draft-22 - UDF-protocol-and-Confit-externs-serving-transformers-without-special-support.md b/backlog/drafts/draft-22 - UDF-protocol-and-Confit-externs-serving-transformers-without-special-support.md new file mode 100644 index 0000000..e98be6c --- /dev/null +++ b/backlog/drafts/draft-22 - UDF-protocol-and-Confit-externs-serving-transformers-without-special-support.md @@ -0,0 +1,203 @@ +--- +id: DRAFT-22 +title: "UDF protocol and Confit externs: serving transformers without special support" +status: Draft +type: spike +created: 2026-07-29 +--- + +## Where this came from + +Design dialogue with AmirHossein, 2026-07-29, right after loop 5 +(transformers as UDAFs v0, PR #61). The question: **how do fitted +transformers — the ones Confit has no special support for — serve through +Confit?** Converged on a full API; parked until the go. + +Two constraints fixed the shape early: + +1. **The Python fallback is permanent infrastructure**, not a stopgap. + There will always be a transformer we haven't optimized, and the + fallback doubles as the oracle that optimized implementations gate + against. +2. **Transformers can sit mid-expression** (`sc(age) OVER () + 1`, a CASE + arm, a WHERE clause — once the v0 bare-final-item refusals lift). Any + pre-SQL/python/post-SQL sandwich dies with that; the opaque call has to + be an expression node the engine evaluates in place. + +## Decision 1 — per-group weights are data, not closure state + +A fitted per-group transformer is a *dict of estimators keyed by partition +key* — which is exactly what params tables already are. So the group lookup +does NOT hide inside the callable; the estimator is a value fetched by the +existing join, and the UDF is a **pure function of its arguments**: + +```sql +-- SELECT sc(age) OVER (PARTITION BY country) + 1 AS z => serving_sql: +SELECT __cf_tf0(p0.__cf_est, t.age) + 1 AS z +FROM __THIS__ AS t +LEFT JOIN __CF_PARAMS_0__ AS p0 ON (t.country IS NOT DISTINCT FROM p0.country) +``` + +- Params table gains an instance-id column (`__cf_est: i64`); ids index the + fitted-clone registry. Both travel together: params stay Arrow, the + instance registry is the one pickle payload, the id is the link. +- **One NULL story.** Unseen group -> LEFT JOIN misses -> NULL id -> NULL + output. Same mechanism as every marginalized aggregate; no dict-miss + convention. +- **No key semantics inside the extern.** The engine's join machinery + already owns `IS NOT DISTINCT FROM`, NULL keys, key expressions. The + extern receives one nullable i64 + typed features. (The rejected + alternative — lookup-in-closure — would force Confit's Rust side to + re-implement join-key equality.) +- **Desugaring becomes a column-width decision.** For a scaler, fit writes + `mean_0, scale_0` into the same params table instead of an id, and the + call site inlines to `(t.age - p0.mean_0) / p0.scale_0`. Same join, same + NULL flow; the oracle-vs-optimized diff is confined to the projected + expression. + +## Decision 2 — the UDF protocol (sql_transform), transformer as subclass + +The base concept is a *declared scalar UDF*; the fitted transformer is the +subclass that adds instance dispatch: + +```python +class UDF: # Confit's whole world + name: str + takes: tuple[str, ...] # Confit Ty vocabulary: "i1"|"i64"|"f64"|"str" + returns: tuple[str, ...] + def __call__(self, *args): ... # scalar semantics — THE contract + def apply_batch(self, *cols: pa.Array) -> tuple[pa.Array, ...]: + ... # default: loop __call__; override to vectorize + +class PythonUDF(UDF): # any plain function + def __init__(self, name, fn, takes, returns): ... + +class PythonTransform(UDF): + instances: dict[int, Any] # instance id -> fitted object with .transform + # implicit leading arg: nullable i64 id, never written in `takes` + # NULL id -> None (unseen group); id missing from instances -> TRAP + # (params table and instances from different fits = broken artifact, not NULL) +``` + +Rules: + +- **Declared types, no probing.** `takes`/`returns` are written in Confit's + own `Ty` vocabulary so build-time checking is literal. Everything is + nullable, always. Deletes sklearn-attribute sniffing; duck-typed + transforms are first-class. `SQLProjection` generates declarations at + `_prepare` (it owns sklearn: `takes` from the resolved query schema, + `returns` from the fitted estimator) so declaration/instance drift is + impossible there; hand-rolled users are kept honest by the runtime check. +- **Two enforcement points.** Build: call-site argument expressions + type-check against `takes`, each use of output slot j against + `returns[j]` — named refusal before any row flows. Runtime: the + trampoline checks each returned tuple against `returns` and traps on + violation. +- **Inputs: scalars for semantics, Arrow for bulk, dicts nowhere.** Row + level takes positional Python scalars (None = NULL) — the boundary has no + names and the row path is the latency path. Batch level takes pyarrow + arrays — already the currency on every side (Confit statics/infer_arrow, + DuckDB vectorized UDFs zero-copy), NULLs ride the validity bitmap. +- **Default `apply_batch` loops the scalar protocol** — amortizes the + boundary (one GIL crossing, one Arrow exchange per chunk) WITHOUT + vectorizing the math. `est.transform` on an n-row block goes through + different BLAS kernels than 1-row calls; looping keeps + row ≡ batch ≡ oracle bit-exact. A `VectorizedTransform` override is the + explicit opt-in (documented ulp caveat). +- **Determinism is the contract.** Fit-time DuckDB, serve-time Confit, and + the oracle all call the same object and are compared bit-exact; a UDF + with randomness breaks the round-trip invariant unlocalizably. Docstring + contract, not enforcement code. +- **Scope fence: scalar UDFs only.** A user-defined aggregate over + `__THIS__` is what transformers already are (fit/transform + window + syntax). The protocol never grows an aggregation arm. + +## Decision 3 — the InferFn API + +```python +fn = InferFn( + serving_sql, # contains __cf_tf0(p0.__cf_est, t.age) calls + row_tables={"__THIS__": Row}, + static_tables=params, # id columns are ordinary Arrow data + udfs=[PythonTransform("__cf_tf0", instances=..., takes=("f64",), returns=("f64",))], + shape="map", +) +``` + +- Contract, restated with one parameter: **serve bit-for-bit identical to + DuckDB *with these udfs registered* (`con.create_function(u.name, u)` + each), or refuse.** Same two outcomes; the oracle stays runnable because + the udfs list is exactly what you register in DuckDB. +- Undeclared unknown function still refuses at build, verbatim. +- Confit never sees sklearn: Arrow tables + declared callables only. + Serialization stays on the Python side; the InferFn constructor args are + the complete artifact description. +- k-output call at the boundary -> `list[float] | None` field (marshaller + assembles k slots); mid-expression use must index down to a scalar. + Confit's IR stays scalar-only (`i1/i64/f64/str`) — outputs are k slots + via out-param array, the existing `len_out` pattern. + +## Confit implementation route + +The cranelift backend already routes every nontrivial op through +`extern "C"` helpers sharing code with the interpreter +(`packages/confit/src/specializer/exec/cranelift.rs`) — a `*mut Cx` context +pointer, trap flag checked after every fallible call. The UDF extern is +**one more helper family in an existing pattern**: + +```rust +extern "C" fn h_extern(p: *mut Cx, slot: i64, /* args by ABI */ out: *mut f64) -> i64; +// Cx carries the slot table: boxed trampoline (GIL + __call__) or, later, a +// native kernel. Python exception -> trap flag, like every fallible helper. +``` + +Frontend: unknown function + matching udfs entry -> extern-call IR +instruction with k outputs, type-checked like any other op. + +Engine bindings, one object, four targets: + +| engine | binding | +|---|---| +| DuckDB (batch path + oracle) | `con.create_function(u.name, u.apply_batch, type="arrow")` | +| Confit interpreter | boxed closure calling `u.__call__` | +| Confit cranelift | C-ABI trampoline over `u.__call__` | +| future optimized | native kernel in the same slot (`NativeAffine(...)` etc.) | + +The oracle gate for optimized impls is "same query, same statics, swap the +list entry." Bit-exact for scaler-class ops; matvec-class needs a +**declared per-estimator ulp tolerance** (BLAS summation order is +unpinnable) — written down per kind, not discovered per bug. + +## What this deletes / changes in sql_transform + +- `marginalize` emits the call **in place** (mid-expression works) instead + of helper columns; the `TransformSpec` splicing block in + `_projection.py::transform` deletes — batch transform becomes + register-udfs-and-run. +- The v0 refusals (bare final-level items only) lift with no serving-side + design change. +- Plain scalar UDFs in authoring SQL resolve exactly like transformers + (registry, then caller scope, guarded by the `duckdb_functions()` + catalog) — the only difference is scalar vs window call position. + +## Sequencing (when greenlit) + +1. SQL-call rewrite in `marginalize` + DuckDB UDF binding — batch path gets + mid-expression transformers with zero Confit involvement. +2. Confit `udfs=` API + extern helper family + trampoline. +3. The params-join wiring Confit already needs regardless + (`IS NOT DISTINCT FROM`, `ON ((1 = 1))`) — previously queued. + +## Rejected along the way + +- **Compile-out as the primary path** (desugar/ONNX): a fallback is + mandatory anyway (coverage + oracle); ONNX kernels aren't sklearn-exact. +- **The sandwich / staged serving plan**: mid-expression kills the single + sandwich; the staged DAG generalization works but costs N boundary + crossings, helper-schema soup, per-stage shape contracts. Named escape + hatch if the extern loop stalls, nothing more. +- **Group lookup inside the closure**: re-implements join-key semantics in + Rust; two NULL conventions. +- **Arity/type probing in Confit**: replaced by declared `takes`/`returns`. +- **Affine probing** (recover W,b numerically, serve as matvec): a third + "approximately" mode — exactly what the two-outcome contract forbids. diff --git a/docs/superpowers/specs/2026-07-29-transformers-run-the-pipeline-design.md b/docs/superpowers/specs/2026-07-29-transformers-run-the-pipeline-design.md new file mode 100644 index 0000000..254a3c2 --- /dev/null +++ b/docs/superpowers/specs/2026-07-29-transformers-run-the-pipeline-design.md @@ -0,0 +1,99 @@ +# Transformers as UDAFs, v0: just run the pipeline + +**Date:** 2026-07-29 +**Status:** loop 5 of SQLProjection; designed in dialogue with AmirHossein. +**Deliberately deferred** (AmirHossein's scoping): desugaring transformers +into SQL windows/applies (that is the *optimization* pass, soon), and the +`sklearn.*` curated-semantics namespace (needs an index of sklearn +functions). This loop runs the actual sklearn objects. + +## Surface + +A window call whose function name is **not in DuckDB's function catalog** +resolves to a transformer: first from an explicit registry, then from the +caller's Python scope (the `FROM df` replacement-scan idiom): + +```python +embed = Pipeline([("sc", StandardScaler()), ("pca", PCA(2))]) +p = SQLProjection( + "SELECT embed(* EXCLUDE (label, country)) OVER (PARTITION BY country) AS e," + " age + 1 AS a1 FROM __THIS__", + this_model=Row, # transformers={"embed": ...} also accepted +).fit(train) +out = p.transform(table) # NEW: batch apply (pa.Table -> pa.Table) +``` + +- Lookup order: `transformers=` dict beats caller scope (locals then globals, + frames walked past our own library code). The oracle-catalog guard means a + scope variable can never shadow a real SQL function. Unknown name + no + match → named refusal with a hint. Captured objects are snapshotted at + construction and exposed as `p.transformers`. +- A transformer object is accepted iff it has `fit` and `transform` + (duck-typed; sklearn `clone` used when available, else `copy.deepcopy`). +- Arguments: one bundle — a plain column, a `struct_pack(...)`, or `*` / + `COLUMNS(...)` with the loop-4 modifiers (schema mode required for + star forms; the bundle rule from the design dialogue: one call, one named + struct, model order). +- `PARTITION BY` = per-group fitting (a clone per group, NULL keys are a + group). `OVER ()` = one global fit. ORDER BY / frames on transformer + windows: named refusals. +- v0 restrictions, each a named refusal: transformer calls only in the + final level of a chain; only as a top-level select item (no nesting the + result in an expression); no transformer inside FILTER/keys/aggregates. + +## The plan: a second step kind + +`FitStep` grows `kind: "sql" | "fit"` (sql default) plus, for fit steps, +`transformer` (registry name), `features` (bundle field names, model order), +and `keys` (partition key column names in the level table). The level table +materializes the bundle as `__cf_f{n}` columns and keys as `__cf_k{m}` — +same mechanism as windows and keys today. The executor grows one branch: + +``` +__CF_LEVEL_0__ sql reads (__THIS__,) SELECT , , FROM __THIS__ +__CF_PARAMS_1__ fit reads (__CF_LEVEL_0__,) transformer=embed keys=(country,) features=(age, fare) +``` + +A fit step's "params table" is not SQL-shaped in v0 (no wide extraction — +that is the deferred optimization): `fit()` stores fitted clones on the +projection, keyed by partition-key tuple, in `p.fitted[step.name]`. The +plan's DAG order and `reads` stay exactly as loop 3 defined them. + +## Serving/apply: batch `transform(table)` + +`serving_sql` cannot express an opaque sklearn apply, so v0 adds a **batch** +apply path (row-at-a-time `infer` through Confit stays NotImplemented): + +- In `serving_sql`, each transformer item is replaced by helper columns: + its bundle features (`__cf_tf{i}_f{n}`) and its partition keys + (`__cf_tf{i}_k{m}`), all row-wise expressions the SQL layer can compute. +- `transform(table)` runs `serving_sql` through DuckDB (threads=1, params + registered), then for each transformer item: group rows by the key tuple, + look up the fitted clone (unseen group → all-NULL output, the exact-join + policy), call `.transform` on the feature block, splice the result in at + the item's position under its alias, drop helpers. +- Transformer output: sklearn returns an (n, k) array → a + `FixedSizeList(float64, k)` column in v0 (named struct fields are part of + the deferred optimization; k is discovered at fit, which is why the output + schema is fit-time knowledge — documented). + +## The gate + +The round-trip invariant splits by column kind, both halves oracle-anchored: + +- SQL columns: DuckDB differential, unchanged, bit-exact. +- Transformer columns: `p.fit(train).transform(train)` must equal an + **independent reference path** — pandas groupby over the same keys, + `clone(est).fit(X_g).transform(X_g)` per group — allclose with tight + tolerance (sklearn's own numerics are BLAS-order dependent; bit-exactness + was never on the table for this stratum, per the DRAFT-20 sizing). +- Refusal tests for every new named refusal; corpus untouched this loop + (transformers are not SQL statements; a transformer corpus can come with + the desugar pass). + +## Out of scope, explicitly + +Desugar-to-SQL (scaler family etc.), wide params extraction, `sklearn.*` +namespace and its semantics contracts, matvec/tree_ensemble instructions, +Confit wiring of applies, transformer calls in non-final levels, struct +access into transformer results, the DRAFT-20 leakage question. diff --git a/packages/sql-transform/sql_transform/__init__.py b/packages/sql-transform/sql_transform/__init__.py index aed4156..44feb8a 100644 --- a/packages/sql-transform/sql_transform/__init__.py +++ b/packages/sql-transform/sql_transform/__init__.py @@ -13,6 +13,7 @@ Marginalized, MarginalizeError, ParamsSpec, + TransformSpec, marginalize, ) from sql_transform._projection import SQLProjection @@ -23,5 +24,6 @@ "MarginalizeError", "ParamsSpec", "SQLProjection", + "TransformSpec", "marginalize", ] diff --git a/packages/sql-transform/sql_transform/_marginalize.py b/packages/sql-transform/sql_transform/_marginalize.py index a597ba9..dfbc6f8 100644 --- a/packages/sql-transform/sql_transform/_marginalize.py +++ b/packages/sql-transform/sql_transform/_marginalize.py @@ -102,12 +102,32 @@ class MarginalizeError(ValueError): @dataclass(frozen=True) class FitStep: - """One fit-time execution: register ``name`` as the result of ``sql``, - which runs against previously registered tables (``reads``).""" + """One fit-time execution. ``kind == "sql"``: register ``name`` as the + result of ``sql`` over previously registered tables (``reads``). + ``kind == "fit"``: group ``reads[0]`` by the ``keys`` columns, fit a + clone of registry transformer ``transformer`` per group on the + ``features`` columns, and store the fitted clones under ``name``.""" name: str sql: str reads: tuple[str, ...] + kind: str = "sql" + transformer: str = "" + features: tuple[str, ...] = () + keys: tuple[str, ...] = () + + +@dataclass(frozen=True) +class TransformSpec: + """One transformer column in the serving output: which fit step supplies + the fitted clones, the output column name, and the helper columns in + ``serving_sql`` carrying its features and partition keys (the apply layer + replaces the helper block with the transformed column in place).""" + + step: str + alias: str + feature_cols: tuple[str, ...] + key_cols: tuple[str, ...] @dataclass(frozen=True) @@ -128,6 +148,7 @@ class Marginalized: serving_sql: str plan: tuple[FitStep, ...] params: tuple[ParamsSpec, ...] + transforms: tuple[TransformSpec, ...] = () def _refuse(what: str, loc: int | None = None) -> None: @@ -716,10 +737,71 @@ def _params_ref(self, raw: Node, w: _Window, discriminates: bool) -> Node: # -- the walk + def _is_transformer_call(self, item: Node) -> str | None: + return _is_tf_node(item) + + def _tf_bundle(self, w_raw: Node) -> list[tuple[str, Node, Node]]: + """(field name, fit expr in level terms, serving expr) per feature.""" + children = w_raw["children"] + if len(children) == 0: + _refuse( + "a transformer call with no arguments — note fn(*) parses to" + " zero arguments; pass a column or struct_pack(...)", + w_raw.get("query_location"), + ) + if len(children) != 1: + _refuse( + "a transformer call takes exactly one bundle argument (a" + " column or struct_pack(...)); configuration belongs in the" + " registered object", + w_raw.get("query_location"), + ) + c = children[0] + quals = self.level.source_quals + if c.get("class") == "COLUMN_REF": + fname = _ColumnRef.of(c).column_names[-1] + return [(fname, _strip_quals(c, quals), dict(self.rewrite(c), alias=""))] + if ( + c.get("class") == "FUNCTION" + and c.get("function_name", "").lower() == "struct_pack" + ): + out: list[tuple[str, Node, Node]] = [] + names: set[str] = set() + for ch in c["children"]: + if not ch.get("alias"): + _refuse( + "struct_pack fields in a transformer bundle must be" + " named (f := expr)", + c.get("query_location"), + ) + if ch["alias"].lower() in names: + _refuse(f"duplicate bundle field {ch['alias']}") + names.add(ch["alias"].lower()) + _walk_row_wise(ch, "a transformer bundle") + out.append( + ( + ch["alias"], + dict(_strip_quals(ch, quals), alias=""), + dict(self.rewrite(ch), alias=""), + ) + ) + return out + _refuse( + "a transformer bundle must be a column or struct_pack(...)", + c.get("query_location"), + ) + raise AssertionError("unreachable") + def rewrite(self, x: Any) -> Any: if isinstance(x, dict): cls = x.get("class") if cls == "WINDOW": + if self._is_transformer_call(x) is not None: + _refuse( + "a transformer call is only supported as a top-level" + " select item", + x.get("query_location"), + ) if self.alias_exprs: self._check_no_lateral_in_window(x) w = _Window.of(x) @@ -858,6 +940,11 @@ def rewrite_items( names_seen: set[str] = set() explicit = not self.env.star for item in items: + tf_name = self._is_transformer_call(item) + if tf_name is not None: + helpers = self.planner.transformer_item(self, item, tf_name, is_final) + out.extend(helpers) + continue if item.get("class") == "STAR": star = _Star.of(item) if explicit: @@ -945,8 +1032,97 @@ def __init__(self) -> None: self.scalars: dict[str, int] = {} self.key_count = 0 self.window_count = 0 + self.feature_count = 0 self.window_names: dict[tuple[int, str], str] = {} self.key_names: dict[tuple[int, str], str] = {} + self.lookup: Any = None # transformer resolver: name -> object | None + self.tf_steps: list[dict] = [] + self.tf_specs: list[TransformSpec] = [] + + def key_name(self, level_i: int, key: _Key) -> str: + kid = (level_i, key.ident) + if kid not in self.key_names: + self.key_names[kid] = f"__cf_k{self.key_count}" + self.key_count += 1 + return self.key_names[kid] + + def transformer_item( + self, rw: _LevelRewriter, raw: Node, name: str, is_final: bool + ) -> list[Node]: + """One transformer window as a top-level item: validate, plan its fit + step, and return the serving helper items (features + keys).""" + if raw.get("schema"): + _refuse( + f"namespaced transformer {raw['schema']}.{name}" + " (the curated-namespace index is a later loop)", + raw.get("query_location"), + ) + if not is_final: + _refuse( + "a transformer call in a non-final level", raw.get("query_location") + ) + obj = self.lookup(name) if self.lookup is not None else None + if obj is None: + _refuse( + f"unknown window function {name} — not a DuckDB aggregate, and" + " no transformer of that name in the registry or caller scope", + raw.get("query_location"), + ) + if not (hasattr(obj, "fit") and hasattr(obj, "transform")): + _refuse(f"transformer {name} has no fit/transform") + w = _Window.of(raw) + if w.orders or w.arg_orders: + _refuse("ORDER BY on a transformer window", w.query_location) + if ( + w.start != "UNBOUNDED_PRECEDING" + or w.end != "CURRENT_ROW_RANGE" + or w.start_expr is not None + or w.end_expr is not None + ): + _refuse("a frame on a transformer window", w.query_location) + if w.filter_expr is not None or w.distinct or w.ignore_nulls: + _refuse( + "FILTER/DISTINCT/IGNORE NULLS on a transformer window", + w.query_location, + ) + keys = [rw._key_of(p) for p in w.partitions] + seen: set[str] = set() + keys = [k for k in keys if not (k.ident in seen or seen.add(k.ident))] + bundle = rw._tf_bundle(raw) + j = len(self.tf_steps) + step_name = f"__CF_FIT_{j}__" + feature_fit = [] + helpers: list[Node] = [] + for i, (fname, fit_expr, serving_expr) in enumerate(bundle): + col = f"__cf_f{self.feature_count}" + self.feature_count += 1 + feature_fit.append((fname, col, fit_expr)) + helpers.append(dict(serving_expr, alias=f"__cf_tf{j}_f{i}")) + key_cols = [] + for m, k in enumerate(keys): + self.key_name(rw.level_i, k) + helpers.append( + dict(copy.deepcopy(k.serving_expr), alias=f"__cf_tf{j}_k{m}") + ) + key_cols.append(f"__cf_tf{j}_k{m}") + self.tf_steps.append( + { + "level": rw.level_i, + "name": step_name, + "transformer": name, + "feature_fit": feature_fit, + "keys": keys, + } + ) + self.tf_specs.append( + TransformSpec( + step=step_name, + alias=raw.get("alias") or name, + feature_cols=tuple(f"__cf_tf{j}_f{i}" for i in range(len(bundle))), + key_cols=tuple(key_cols), + ) + ) + return helpers def group_for(self, level_i: int, keys: list[_Key]) -> tuple[int, _Join]: ident = (level_i, tuple(k.ident for k in keys)) @@ -1036,6 +1212,17 @@ def check(x: Any) -> None: check(sub_node) +def _is_tf_node(item: Node) -> str | None: + """The lowered name when the node is a transformer window call (an + unknown-to-DuckDB or namespaced function with OVER), else None.""" + if item.get("class") != "WINDOW" or item.get("type") != "WINDOW_AGGREGATE": + return None + name = item.get("function_name", "").lower() + if name in _aggregate_names() and not item.get("schema"): + return None + return name + + def _has_windows(x: Any) -> bool: if isinstance(x, dict): if x.get("class") == "WINDOW": @@ -1058,6 +1245,8 @@ def _level_step( quals = level.source_quals originals = [] for item in level.node["select_list"]: + if _is_tf_node(item) is not None: + continue # not executable SQL; its features ride as __cf_f cols fit_item = _strip_quals(item, quals) originals.append(fit_item) windows = [ @@ -1066,18 +1255,25 @@ def _level_step( ] keys: list[Node] = [] emitted: set[str] = set() - for join in planner.joins: - if join.level != level_i: - continue - for k in join.keys: + key_lists = [j.keys for j in planner.joins if j.level == level_i] + [ + t["keys"] for t in planner.tf_steps if t["level"] == level_i + ] + for klist in key_lists: + for k in klist: kname = planner.key_names[(level_i, k.ident)] if kname not in emitted: emitted.add(kname) keys.append(dict(copy.deepcopy(k.fit_expr), alias=kname)) + features = [ + dict(copy.deepcopy(fit_expr), alias=col) + for t in planner.tf_steps + if t["level"] == level_i + for (_, col, fit_expr) in t["feature_fit"] + ] from_table = _clone("base_table", table_name=source_name, alias="") return FitStep( name=f"__CF_LEVEL_{level_i}__", - sql=_select_doc([*originals, *windows, *keys], from_table), + sql=_select_doc([*originals, *windows, *keys, *features], from_table), reads=(source_name,), ) @@ -1147,7 +1343,11 @@ def _serving_from(joins: list[_Join]) -> Node: # --- the entry point ---------------------------------------------------------- -def marginalize(sql: str, columns: Sequence[str] | None = None) -> Marginalized: +def marginalize( + sql: str, + columns: Sequence[str] | None = None, + transformers: Any = None, +) -> Marginalized: """Parse a chain of strict projections over ``__THIS__`` and marginalize it. Pure: parses, rewrites, and plans — never executes. @@ -1213,6 +1413,7 @@ def marginalize(sql: str, columns: Sequence[str] | None = None) -> Marginalized: deepest = max(windows_at, default=-1) planner = _Planner() + planner.lookup = transformers.get if hasattr(transformers, "get") else transformers env = base_env final_items: list[Node] = [] for i, level in enumerate(levels): @@ -1225,12 +1426,32 @@ def marginalize(sql: str, columns: Sequence[str] | None = None) -> Marginalized: for idx, join in enumerate(planner.joins): if join.level == i: planner.plan.append(_collapse_step(planner, idx, join)) + for t in planner.tf_steps: + if t["level"] == i: + planner.plan.append( + FitStep( + name=t["name"], + sql="", + reads=(f"__CF_LEVEL_{i}__",), + kind="fit", + transformer=t["transformer"], + features=tuple(col for (_, col, _) in t["feature_fit"]), + keys=tuple( + planner.key_names[(i, k.ident)] for k in t["keys"] + ), + ) + ) if is_final: final_items = rewritten break env = _Env(entries) - if not planner.joins and len(levels) == 1 and base_env is _BASE_ENV: + if ( + not planner.joins + and not planner.tf_steps + and len(levels) == 1 + and base_env is _BASE_ENV + ): # No aggregates, no chain, no schema: marginalization is the identity # (modulo normalization). With a declared schema the rewrite always # canonicalizes — stars/COLUMNS expand, lateral aliases inline. @@ -1244,5 +1465,8 @@ def marginalize(sql: str, columns: Sequence[str] | None = None) -> Marginalized: for i, j in enumerate(planner.joins) ) return Marginalized( - serving_sql=_deserialize(doc), plan=tuple(planner.plan), params=specs + serving_sql=_deserialize(doc), + plan=tuple(planner.plan), + params=specs, + transforms=tuple(planner.tf_specs), ) diff --git a/packages/sql-transform/sql_transform/_projection.py b/packages/sql-transform/sql_transform/_projection.py index a036ffd..ecdd9e1 100644 --- a/packages/sql-transform/sql_transform/_projection.py +++ b/packages/sql-transform/sql_transform/_projection.py @@ -16,6 +16,19 @@ from sql_transform._marginalize import MarginalizeError, marginalize + +def _feature_matrix(table: pa.Table, cols: list[str]): + """Feature block as a 2D numpy array (float when possible, else object).""" + import numpy as np + + raw = [table.column(c).to_pylist() for c in cols] + try: + mat = np.array(raw, dtype=float).T + except TypeError, ValueError: + mat = np.array(raw, dtype=object).T + return mat + + _TODO = "SQLProjection serving is a later loop; {} has no implementation yet" @@ -34,6 +47,7 @@ def __init__( sql: str | Template, /, this_model: type[BaseModel] | None = None, + transformers: dict[str, Any] | None = None, ) -> None: """``this_model`` declares the ``__THIS__`` schema (pydantic field names, in definition order — the model is authoritative). With it, @@ -46,8 +60,27 @@ def __init__( self._columns = ( list(this_model.model_fields) if this_model is not None else None ) - self._marginalized = marginalize(sql, self._columns) + # Transformer resolution: explicit registry first, then the caller's + # scope (the `FROM df` replacement-scan idiom). Captured objects are + # snapshotted into `self.transformers` at construction. + import sys + + frame = sys._getframe(1) + self.transformers: dict[str, Any] = {} + + def _resolve(name: str): + obj = None + if transformers is not None: + obj = transformers.get(name) + if obj is None and frame is not None: + obj = frame.f_locals.get(name) or frame.f_globals.get(name) + if obj is not None: + self.transformers[name] = obj + return obj + + self._marginalized = marginalize(sql, self._columns, _resolve) self._params: dict[str, pa.Table] | None = None + self._fitted: dict[str, dict[tuple, Any]] | None = None @classmethod def from_file(cls, path: str) -> SQLProjection: @@ -85,7 +118,13 @@ def fit(self, table: pa.Table, /) -> SQLProjection: # register its result under its name. Every intermediate is # inspectable and every step is plain SQL runnable by hand. materialized: dict[str, pa.Table] = {} + self._fitted = {} for step in m.plan: + if step.kind == "fit": + self._fitted[step.name] = self._fit_step( + step, materialized[step.reads[0]] + ) + continue materialized[step.name] = con.execute(step.sql).to_arrow_table() con.register(step.name, materialized[step.name]) self._params = {spec.name: materialized[spec.name] for spec in m.params} @@ -93,6 +132,74 @@ def fit(self, table: pa.Table, /) -> SQLProjection: con.close() return self + def _fit_step(self, step, table: pa.Table) -> dict[tuple, Any]: + """Group by the key columns, fit a clone of the transformer per group + on the feature block. Unknown groups at apply time get NULLs.""" + proto = self.transformers[step.transformer] + try: + from sklearn.base import clone + except ImportError: # duck-typed objects without sklearn installed + import copy as _copy + + def clone(o): # type: ignore[no-redef] + return _copy.deepcopy(o) + + feats = _feature_matrix(table, list(step.features)) + keys = [table.column(k).to_pylist() for k in step.keys] + groups: dict[tuple, list[int]] = {} + for i in range(table.num_rows): + groups.setdefault(tuple(k[i] for k in keys), []).append(i) + fitted: dict[tuple, Any] = {} + for key, idx in groups.items(): + est = clone(proto) + est.fit(feats[idx]) + fitted[key] = est + return fitted + + def transform(self, table: pa.Table, /) -> pa.Table: + """Batch apply: run ``serving_sql`` through DuckDB, then run each + fitted transformer on its helper columns and splice the outputs in + place of the helper blocks. Row-at-a-time serving stays with Confit.""" + import duckdb + import numpy as np + + if self._params is None or self._fitted is None: + raise MarginalizeError("not fitted: call fit(table) first") + if self._columns is not None: + table = table.select(self._columns) + m = self._marginalized + con = duckdb.connect() + try: + con.execute("SET threads = 1") + con.register("__THIS__", table) + for name, params_table in self._params.items(): + con.register(name, params_table) + res = con.execute(m.serving_sql).to_arrow_table() + finally: + con.close() + for spec in m.transforms: + feats = _feature_matrix(res, list(spec.feature_cols)) + keys = [res.column(k).to_pylist() for k in spec.key_cols] + groups: dict[tuple, list[int]] = {} + for i in range(res.num_rows): + groups.setdefault(tuple(k[i] for k in keys), []).append(i) + fitted = self._fitted[spec.step] + out: list = [None] * res.num_rows + for key, idx in groups.items(): + est = fitted.get(key) + if est is None: + continue # unseen group: NULL output (exact-join policy) + block = np.asarray(est.transform(feats[idx])) + if block.ndim == 1: + block = block.reshape(-1, 1) + for row, vals in zip(idx, block, strict=True): + out[row] = [float(v) for v in vals] + first = res.column_names.index(spec.feature_cols[0]) + drop = set(spec.feature_cols) | set(spec.key_cols) + res = res.drop_columns([c for c in res.column_names if c in drop]) + res = res.add_column(first, spec.alias, pa.array(out)) + return res + @property def serving_sql(self) -> str: """The rewritten projection: params joins instead of aggregates.""" diff --git a/packages/sql-transform/sql_transform/_transformers_test.py b/packages/sql-transform/sql_transform/_transformers_test.py new file mode 100644 index 0000000..7999d35 --- /dev/null +++ b/packages/sql-transform/sql_transform/_transformers_test.py @@ -0,0 +1,130 @@ +"""Transformers-as-UDAFs v0: fit runs sklearn per group; the gate's oracle +for transformer columns is an independent clone/fit/transform reference.""" + +import numpy as np +import pyarrow as pa +import pytest +from sklearn.base import clone +from sklearn.decomposition import PCA +from sklearn.pipeline import Pipeline +from sklearn.preprocessing import StandardScaler + +from sql_transform import MarginalizeError, SQLProjection + +TRAIN = pa.table( + { + "country": ["US", "US", None, "DE", None, "DE", "FR"], + "age": [40.0, 30.0, 20.0, 25.0, 50.0, 45.0, 35.0], + "fare": [7.0, 8.0, 6.0, 9.0, 10.0, 11.0, 5.0], + "name": ["x", "y", "z", "w", "v", "u", "t"], + } +) + + +def _reference(proto, feats, keys): + """Independent oracle: clone-per-group fit_transform, row-aligned.""" + groups = {} + for i, k in enumerate(keys): + groups.setdefault(k, []).append(i) + out = [None] * len(keys) + for _k, idx in groups.items(): + est = clone(proto) + block = np.asarray(est.fit(feats[idx]).transform(feats[idx])) + if block.ndim == 1: + block = block.reshape(-1, 1) + for row, vals in zip(idx, block, strict=True): + out[row] = [float(v) for v in vals] + return out + + +def _by_name(table, value_col): + names = table.column("name").to_pylist() + vals = table.column(value_col).to_pylist() + return dict(zip(names, vals, strict=True)) + + +def test_global_scaler_from_scope(): + sc = StandardScaler() + p = SQLProjection("SELECT sc(age) OVER () AS z, name FROM __THIS__").fit(TRAIN) + (step,) = [s for s in p.plan if s.kind == "fit"] + assert step.transformer == "sc" and step.keys == () + got = _by_name(p.transform(TRAIN), "z") + feats = np.array([TRAIN.column("age").to_pylist()], dtype=float).T + ref = _reference(sc, feats, [()] * TRAIN.num_rows) + for i, n in enumerate(TRAIN.column("name").to_pylist()): + np.testing.assert_allclose(got[n], ref[i], rtol=1e-12) + + +def test_per_country_pipeline_with_struct_bundle(): + embed = Pipeline([("s", StandardScaler()), ("p", PCA(n_components=1))]) + p = SQLProjection( + "SELECT embed(struct_pack(a := age, f := fare))" + " OVER (PARTITION BY country) AS e, name FROM __THIS__", + transformers={"embed": embed}, + ).fit(TRAIN) + got = _by_name(p.transform(TRAIN), "e") + feats = np.array( + [TRAIN.column("age").to_pylist(), TRAIN.column("fare").to_pylist()], + dtype=float, + ).T + ref = _reference(embed, feats, [(c,) for c in TRAIN.column("country").to_pylist()]) + for i, n in enumerate(TRAIN.column("name").to_pylist()): + np.testing.assert_allclose(got[n], ref[i], rtol=1e-9) + + +def test_unseen_group_gets_null(): + sc = StandardScaler() + p = SQLProjection( + "SELECT sc(age) OVER (PARTITION BY country) AS z, name FROM __THIS__", + transformers={"sc": sc}, + ).fit(TRAIN) + new = pa.table({"country": ["JP"], "age": [1.0], "fare": [1.0], "name": ["q"]}) + assert p.transform(new).column("z").to_pylist() == [None] + + +def test_mixed_sql_and_transformer_columns(): + sc = StandardScaler() + p = SQLProjection( + "SELECT age - avg(age) OVER () AS c, sc(fare) OVER () AS z, name FROM __THIS__", + transformers={"sc": sc}, + ).fit(TRAIN) + out = p.transform(TRAIN) + mean = np.mean(TRAIN.column("age").to_pylist()) + got = _by_name(out, "c") + for i, n in enumerate(TRAIN.column("name").to_pylist()): + np.testing.assert_allclose(got[n], TRAIN.column("age")[i].as_py() - mean) + assert out.column_names.index("c") < out.column_names.index("z") + + +@pytest.mark.parametrize( + "sql,match", + [ + ("SELECT nope(age) OVER () FROM __THIS__", "unknown window function nope"), + ( + "SELECT sc(age) OVER (ORDER BY age) FROM __THIS__", + "ORDER BY on a transformer", + ), + ("SELECT sc(*) OVER () FROM __THIS__", "zero arguments"), + ("SELECT sc(age, fare) OVER () FROM __THIS__", "exactly one bundle"), + ("SELECT sc(struct_pack(age)) OVER () FROM __THIS__", "must be named"), + ("SELECT sc(age) OVER () + 1 FROM __THIS__", "top-level select item"), + ( + "WITH a AS (SELECT sc(age) OVER () AS z FROM __THIS__) SELECT z FROM a", + "non-final level", + ), + ("SELECT sk.sc(age) OVER () FROM __THIS__", "namespaced transformer"), + ], +) +def test_transformer_refusals(sql, match): + sc = StandardScaler() + with pytest.raises(MarginalizeError, match=match): + SQLProjection(sql, transformers={"sc": sc}) + + +def test_registry_beats_scope(): + sc = StandardScaler() # noqa: F841 — the scope decoy + p = SQLProjection( + "SELECT sc(age) OVER () AS z FROM __THIS__", + transformers={"sc": Pipeline([("s", StandardScaler())])}, + ) + assert isinstance(p.transformers["sc"], Pipeline)