diff --git a/README.md b/README.md index c95c54c..13dfaa5 100644 --- a/README.md +++ b/README.md @@ -72,8 +72,8 @@ sql = "SELECT target / MEAN(target) OVER (PARTITION BY city) AS enc FROM __THIS_ - **Typed I/O**: Pydantic models for the input row and output, validated when the transformer is built and again at call time. Output is a typed model; the input schema is auto-synthesized or user-supplied. -- **Transformer references**: interpolate a fitted sklearn transformer directly - into a t-string and apply it to columns (see below). +- **Transform references**: interpolate another `SQLTransform` into a t-string to + compose transforms (see below). See [docs/SQL_SUPPORT.md](docs/SQL_SUPPORT.md) for the feature-by-feature tracker. @@ -104,57 +104,33 @@ query engine once; inference pays only for a lean interpreter walking a plan. Th separation of **fit** (compute statistics) and **transform/infer** (apply them) is the standard ML pattern — fit on training data, apply to training and serving. -## Referencing a fitted sklearn transformer +## Referencing another SQLTransform -Interpolate a fitted transformer into a t-string and apply it to columns: +Interpolate a `SQLTransform` into a t-string to apply it to a column. The +reference inlines to a scalar expression at `fit`, so it composes with window +aggregates freely: ```python -import pandas as pd import pyarrow as pa -from sklearn.preprocessing import StandardScaler from sql_transform import SQLTransform -train_df = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} -) -table = pa.Table.from_pandas(train_df) +table = pa.table({"age": [10.0, 20.0, 30.0, 40.0]}) -sc = StandardScaler().fit(train_df) # fit on a DataFrame -> records feature_names_in_ -t = SQLTransform(t"SELECT {sc}(age, income) AS scaled FROM __THIS__").fit(table) -``` - -**Output is a single Arrow struct column**, not one column per feature: - -```python -t.transform(table).schema # scaled: struct -t.transform(table).flatten().schema.names # ['scaled.age', 'scaled.income'] -``` - -Call `.flatten()` to get flat columns for an sklearn handoff. - -**Column binding** depends on how the transformer was fitted: - -| fitted with | `feature_names_in_` | binding | -|---|---|---| -| `fit(DataFrame)` | recorded | by **name** — call order is free, and is validated against the names | -| `fit(ndarray)` | absent | by **position**, in call order — only the count is checked | - -With positional binding, `{sc}(income, age)` against a transformer fitted as -`[age, income]` silently swaps the features. Fit on a DataFrame when you can. - -Aggregating over a transformer's output (`AVG({sc}(age)) OVER ()`) is not -supported — it is inherently two-stage and needs a subquery. Aggregate over an -input column, or use a `SQLTransform` reference, which inlines to a scalar and -composes with aggregates freely: - -```python norm = SQLTransform("SELECT age / MEAN(age) OVER () AS a FROM __THIS__").fit(table) -t2 = SQLTransform( +t = SQLTransform( t"SELECT AVG({norm.transform}(age)) OVER () AS m FROM __THIS__" ).fit(table) -t2.transform(table).column("m").to_pylist() # [1.0, 1.0, 1.0, 1.0] +t.transform(table).column("m").to_pylist() # [1.0, 1.0, 1.0, 1.0] ``` +`{norm.transform}` reuses the referenced transform's already-frozen state. A bare +`{norm}` instead re-fits that transform's *definition* into the outer fit, leaving +the referenced object untouched. References nest — `{a}({b}(age))` — and resolve +innermost-first. + +> **Note:** referencing a fitted **sklearn** transformer is not currently +> supported. That surface was removed in TASK-40 and may return later. + ## Development ```bash diff --git "a/backlog/tasks/task-34 - codegen-transformer-support-resolve_transformers-\342\206\222-Expr-Transform-native-only-today.md" "b/backlog/archive/tasks/task-34 - codegen-transformer-support-resolve_transformers-\342\206\222-Expr-Transform-native-only-today.md" similarity index 84% rename from "backlog/tasks/task-34 - codegen-transformer-support-resolve_transformers-\342\206\222-Expr-Transform-native-only-today.md" rename to "backlog/archive/tasks/task-34 - codegen-transformer-support-resolve_transformers-\342\206\222-Expr-Transform-native-only-today.md" index d8ba3fb..106530d 100644 --- "a/backlog/tasks/task-34 - codegen-transformer-support-resolve_transformers-\342\206\222-Expr-Transform-native-only-today.md" +++ "b/backlog/archive/tasks/task-34 - codegen-transformer-support-resolve_transformers-\342\206\222-Expr-Transform-native-only-today.md" @@ -7,7 +7,7 @@ status: To Do assignee: - Ritchie created_date: '2026-07-19 16:08' -updated_date: '2026-07-23 14:32' +updated_date: '2026-07-25 01:23' labels: - codegen - transformer @@ -77,4 +77,14 @@ created: 2026-07-23 04:41 --- PRE-AUTHORIZED (2026-07-23): AmirHossein green-lit dispatching TASK-34 to Ritchie automatically once TASK-29 lands. No further approval needed — PM dispatches on TASK-29 completion. (Dependency TASK-29 still In Progress.) --- + +author: Claude (TASK-40) +created: 2026-07-25 01:23 +--- +OBSOLETED BY TASK-40 (2026-07-25), archiving. This ticket adds codegen support for the transformer-ref feature; TASK-40 removes that feature from both engines, so there is nothing left to port. The `__tfm_` UnsupportedInCodegen branch this ticket was meant to replace is deleted too. + +Its standing PRE-AUTHORIZATION to auto-dispatch to Ritchie on TASK-29 completion is retired -- do NOT dispatch when TASK-29 lands. + +Recoverable: archived, not deleted. If transformer support returns, write a fresh ticket against the new surface rather than un-archiving this one (its AC#2/#3 reference resolve_transformers and TASK-31's parity test, both of which are gone). +--- diff --git a/backlog/tasks/task-3 - Transformer-refs-Part-2-authoring-surface-review-follow-ups.md b/backlog/tasks/task-3 - Transformer-refs-Part-2-authoring-surface-review-follow-ups.md index 35c26ba..bad8fea 100644 --- a/backlog/tasks/task-3 - Transformer-refs-Part-2-authoring-surface-review-follow-ups.md +++ b/backlog/tasks/task-3 - Transformer-refs-Part-2-authoring-surface-review-follow-ups.md @@ -1,11 +1,11 @@ --- id: TASK-3 title: Transformer-refs (Part-2 authoring surface) review follow-ups -status: In Progress +status: Done assignee: - Wren created_date: '2026-07-18 13:44' -updated_date: '2026-07-23 14:32' +updated_date: '2026-07-25 01:23' labels: - python - transformer-refs @@ -69,4 +69,24 @@ created: 2026-07-23 13:09 --- Dispatched to Wren (2026-07-23) per the standing pre-authorization, on TASK-2 completion. 6 ACs — mostly DX/guardrails on the shipped authoring surface. Note AC#6 is a README/docs item and AC#4 is self-described low value; Wren to flag if any AC looks not-worth-it rather than padding it out. --- + +author: Claude (TASK-40) +created: 2026-07-25 01:23 +--- +CLOSED BY TASK-40 (2026-07-25). The five fixes that were completed landed on master (commits d47d2df..a83b742: bind names on both ref paths, probe in CALL order, nested refs declare the struct they receive, probe in fitted order + engine tolerance, shared _rows_equal). + +ACs #1-#6 are now VOID, not deferred: TASK-40 removes the {transformer}(col) authoring surface entirely, so the double-.transform() probe, the feature_names_in_ footgun, the aggregate-over-output error, and the struct-output README note all describe code that no longer exists. Re-open a fresh ticket against the new surface if transformer support returns. + +Important: this ticket carried a standing PRE-AUTHORIZATION to dispatch to Wren. Closing it retires that authorization -- do not dispatch. +--- + +## Final Summary + + +Partially completed, then superseded by TASK-40 (remove transformer support). + +LANDED (on master): probe/binding correctness fixes for the transformer-ref path -- bind names on both ref paths with an actionable error when unsettable; probe in CALL order so the UDF signature matches the named_struct the SQL builds; nested refs declare the struct they actually receive; probe in fitted order and compare engines with a tolerance; test cleanup to the shared _rows_equal helper. + +NOT DONE, AND NOW VOID: ACs #1-#6. All six target the {transformer}(col) authoring surface, which TASK-40 deletes from both the Python front-end and the native interpreter. They are not deferred work -- the code they describe is gone. + diff --git a/backlog/tasks/task-31 - Test-transformer-ref-nested-inside-a-CASE-branch-native-codegen-parity.md b/backlog/tasks/task-31 - Test-transformer-ref-nested-inside-a-CASE-branch-native-codegen-parity.md index c7ef548..7bd15e0 100644 --- a/backlog/tasks/task-31 - Test-transformer-ref-nested-inside-a-CASE-branch-native-codegen-parity.md +++ b/backlog/tasks/task-31 - Test-transformer-ref-nested-inside-a-CASE-branch-native-codegen-parity.md @@ -5,7 +5,7 @@ status: Done assignee: - Ritchie created_date: '2026-07-19 15:35' -updated_date: '2026-07-23 00:53' +updated_date: '2026-07-25 01:24' labels: - test - sql-surface @@ -46,6 +46,12 @@ created: 2026-07-19 15:53 --- Done per Ritchie's merge cb8c350. Mutation-verified (broke the recursion, saw the test go red on non-first branches, restored) — meets the 'must catch the regression' bar exactly. Answering Ritchie's codegen question: your call is correct — transformers are native-only, so a native parity test + a codegen defers-loudly (UnsupportedInCodegen) assertion is the right shape of 'both engines'. Nothing to change; don't fabricate codegen computation for a native-only feature. AC #2 accepted under that corrected premise. --- + +author: Claude (TASK-40) +created: 2026-07-25 01:24 +--- +NOTE (2026-07-25): this ticket's test file, tests/test_transformer_case.py, was deleted by TASK-40 along with the transformer-ref feature it exercised. Ticket stays Done -- the work was real and correct when it landed; the surface it covered no longer exists. Native/codegen parity for CASE itself remains covered by tests/test_diff_case.py and tests/test_case_window_integration.py. +--- ## Final Summary diff --git a/backlog/tasks/task-40 - Remove-transformer-support-from-SQLTransformer-and-the-native-interpreter.md b/backlog/tasks/task-40 - Remove-transformer-support-from-SQLTransformer-and-the-native-interpreter.md new file mode 100644 index 0000000..757d21b --- /dev/null +++ b/backlog/tasks/task-40 - Remove-transformer-support-from-SQLTransformer-and-the-native-interpreter.md @@ -0,0 +1,96 @@ +--- +id: TASK-40 +title: Remove transformer support from SQLTransformer and the native interpreter +status: In Progress +assignee: + - '@Claude' +created_date: '2026-07-25 01:09' +updated_date: '2026-07-25 01:25' +labels: + - epic +dependencies: [] +references: + - 'https://github.com/ahrzb/sql-transforms/pull/22' +priority: high +type: task +ordinal: 35000 +--- + +## Description + + +Epic: rip out the transformer-ref / transformer-callout feature for now (v0, no backward compat — delete outright, no deprecation shims). + +Scope, two removals: + +1. **SQLTransformer surface (Python)** — remove the ability to reference fitted sklearn transformers inside a composition: + - `sql_transform/_transformer_ref.py`, `sql_transform/_transformer_udf.py` (delete) + - `is_transformer` branches in `sql_transform/_compose.py` (Ref.is_transformer, unfitted-transformer guard, callout build path) + - transformer paths in `sql_transform/_batch.py`, `sql_transform/_codegen/engine.py`, `sql_transform/__init__.py` exports +2. **Native interpreter (Rust)** — remove the transformer callout / UDF machinery: + - `Expr::Transform` and related plumbing in `src/expr.rs`, `src/plan.rs`, `src/types.rs`, `src/schema.rs`, `src/lib.rs` + - regenerate `sql_transform/_interpreter.pyi` + +Tests to delete: `tests/test_transformer_udf.py`, `tests/test_transformer_ref.py`, `tests/test_transformer_case.py`, `tests/test_transformer_callout_infer.py`, `tests/test_diff_transformer_callout.py`, plus transformer cases in `sql_transform/_compose_test.py` and `tests/test_diff_errors.py`. + +Ticket interactions to resolve before/while dispatching: +- TASK-3 (In Progress, transformer-refs follow-ups) works on the exact surface being deleted — needs a stop/land decision. +- TASK-34 (To Do, codegen transformer support) becomes obsolete — archive when this epic lands. +- TASK-31's parity test gets deleted with the feature. + +Docs: update SQL_SUPPORT.md / README where transformer refs are mentioned; the design specs under docs/superpowers/specs stay as history. + + +## Acceptance Criteria + +- [x] #1 No transformer-ref / callout code paths remain in sql_transform (grep for is_transformer, transformer_ref, transformer_udf comes back empty) +- [x] #2 No Expr::Transform / callout machinery remains in src/*.rs; native extension rebuilds clean +- [x] #3 All transformer tests removed; full test suite passes on both engines (codegen + native) +- [x] #4 TASK-3 resolved (stopped or landed) and TASK-34 archived or re-scoped +- [x] #5 Docs no longer advertise transformer refs + + +## Implementation Notes + + +Removal done on branch `claude/remove-transformer-support-deb837`; PR open, ticket stays In Progress until it merges. + +PYTHON SURFACE +- Deleted `_transformer_ref.py` (239 lines: is_transformer, _bind_names, _probe, resolve_transformer_refs) and `_transformer_udf.py` (84 lines: the DataFusion ScalarUDF wrapper). +- `_compose.py`: dropped `Ref.is_transformer`, the `is_transformer(v)` branch, and the transformer-shaped-but-unfitted guard. A sklearn object now falls through to the existing TypeError, which names the type: "interpolation {sc} must be a SQLTransform or its .transform, got StandardScaler". SQLTransform-to-SQLTransform composition is untouched. +- `__init__.py`: dropped `_udf_specs`, the sqlt/tfm ref split, and the `transformers=` kwarg to InferFn; `inline_references` now takes `self._refs` whole. +- `_batch.py`: dropped the `transformers` param and UDF registration loop. +- `_codegen/engine.py`: dropped the `__tfm_` UnsupportedInCodegen branch (an unknown function is now uniformly a ValueError). + +NATIVE INTERPRETER +- `expr.rs`: removed the `Expr::Transform` variant and its ~90-line eval arm (the GIL-attach / feature reorder / .transform() / .tolist() marshalling path). +- `lib.rs`: removed `ResolvedTransformer`, `read_feature_names_in`, the 70-line `resolve_transformers` tree rewrite, and the `transformers` constructor param. +- `types.rs`, `plan.rs`: removed the `Expr::Transform` arms from infer_type and validate_expr. +- `schema.rs`: removed `arrow_schema_to_ordered_fields` (its only caller was the transformer path). +- Unused imports cleaned: `Arc` + `FieldType` in expr.rs, `HashSet` in types.rs. + +VERIFICATION +- `cargo check`: clean, zero warnings. +- `uv run pytest`: 523 passed, 12 xfailed (the xfails are pre-existing, unrelated to transformers). +- `ruff check` + `ruff format --check`: clean. +- Ran the rewritten README example end to end: transform() gives [1.0, 1.0, 1.0, 1.0], infer() gives m=1.0, a bare `{t}` re-fit ref gives [0.4, 0.8, 1.2, 1.6], and a fitted StandardScaler is rejected with the expected TypeError. +- `cargo fmt --check` reports 11 diffs, all pre-existing (master has 13 in the same files); cargo fmt is not in the project's `mise run fmt` pipeline, so the crate was never fmt-clean. Not touched -- reformatting would bury the deletion in noise. + +TICKET FALLOUT +- TASK-3 closed Done: its five landed commits are on master, ACs #1-#6 are void (they describe deleted code). Its standing pre-authorization to dispatch to Wren is retired. +- TASK-34 archived: it exists to port this feature to codegen, so its premise is gone. Its pre-authorization to auto-dispatch to Ritchie when TASK-29 lands is retired -- that was the urgent part, otherwise a dev gets sent at deleted code. +- TASK-31's parity test (tests/test_transformer_case.py) was deleted with the feature; TASK-31 stays Done with a note. + +DOCS +- README: replaced the "Referencing a fitted sklearn transformer" section with "Referencing another SQLTransform" (the surviving composition feature, with a verified runnable example), and noted the sklearn surface is gone for now. Design specs under docs/superpowers/specs and the backlog docs/decisions/drafts are left as history. + + +## Comments + + +author: Claude +created: 2026-07-25 01:25 +--- +PR open: https://github.com/ahrzb/sql-transforms/pull/22 (+152 -1616 across 24 files). All 5 ACs met and verified; ticket stays In Progress until the PR merges. +--- + diff --git a/sql_transform/__init__.py b/sql_transform/__init__.py index 7f2566d..7ea4de5 100644 --- a/sql_transform/__init__.py +++ b/sql_transform/__init__.py @@ -18,7 +18,6 @@ from sql_transform._schema import synthesize_this_model from sql_transform._sql import find_window_aggregates, parse_and_validate from sql_transform._state import build_state_tables -from sql_transform._transformer_ref import resolve_transformer_refs __all__ = ["InferFn", "SQLTransform"] @@ -64,7 +63,6 @@ def __init__( self._state_tables: dict[str, pa.Table] | None = None self._rewritten_sql: str | None = None self._infer_fn: InferFn | None = None - self._udf_specs: dict[str, tuple] = {} @classmethod def from_file(cls, path: str) -> SQLTransform: @@ -84,11 +82,7 @@ def fit( ctx = datafusion.SessionContext() ctx.from_arrow(table, name="__THIS__") - sqlt_refs = {n: r for n, r in self._refs.items() if not r.is_transformer} - tfm_refs = {n: r.transform for n, r in self._refs.items() if r.is_transformer} - self._udf_specs = resolve_transformer_refs(tree, tfm_refs, table) - - inline = inline_references(tree, sqlt_refs, ctx, table) + inline = inline_references(tree, self._refs, ctx, table) windows = find_window_aggregates(tree) own_state = build_state_tables( @@ -102,9 +96,6 @@ def fit( self._rewritten_sql, row_tables={"__THIS__": this_model}, static_tables=self._state_tables, - transformers={ - n: (obj, out_s) for n, (obj, in_s, out_s) in self._udf_specs.items() - }, ) return self @@ -116,7 +107,7 @@ def transform(self, table: pa.Table, /) -> pa.Table: the Rust engine instead.""" if self._infer_fn is None: raise RuntimeError("Must call fit() before transform") - out = run_batch(self._rewritten_sql, table, self._state_tables, self._udf_specs) + out = run_batch(self._rewritten_sql, table, self._state_tables) if self._output == "dense": return _table_to_dense(out) return out diff --git a/sql_transform/_batch.py b/sql_transform/_batch.py index 8fc574e..a7580eb 100644 --- a/sql_transform/_batch.py +++ b/sql_transform/_batch.py @@ -12,8 +12,6 @@ import pyarrow as pa import sqlglot -from sql_transform._transformer_udf import _transformer_udf - _ROW_ID = "__row_id__" @@ -21,7 +19,6 @@ def run_batch( rewritten_sql: str, table: pa.Table, state_tables: dict[str, pa.Table], - transformers: dict[str, tuple[object, pa.Schema, pa.Schema]] | None = None, ) -> pa.Table: """Execute `rewritten_sql` against `table` (as __THIS__) and every state table (registered by name) via DataFusion, returning a pyarrow Table. @@ -48,8 +45,6 @@ def run_batch( ctx.from_arrow(numbered_table, name="__THIS__") for name, state_table in state_tables.items(): ctx.from_arrow(state_table, name=name) - for name, (obj, in_schema, out_schema) in (transformers or {}).items(): - ctx.register_udf(_transformer_udf(obj, in_schema, out_schema, name)) df = ctx.sql(tree.sql()) # collect() returns [] for a zero-row result, so pass the schema explicitly. out = pa.Table.from_batches(df.collect(), schema=df.schema()) diff --git a/sql_transform/_codegen/engine.py b/sql_transform/_codegen/engine.py index d5ace53..66a5546 100644 --- a/sql_transform/_codegen/engine.py +++ b/sql_transform/_codegen/engine.py @@ -94,12 +94,6 @@ def _emit_expr(e: Any, env: dict) -> str: if isinstance(e, cp.Func): fn = _BUILTINS.get(e.name) if fn is None: - if e.name.startswith("__tfm_"): - # A transformer placeholder call: codegen has no transformer - # support (native-only, TASK-34) -- deferred surface, not an error. - raise UnsupportedInCodegen( - f"{e.name}() is not supported in codegen yet" - ) raise ValueError(f"Unknown function: {e.name}") return f"{fn}({', '.join(_emit_expr(a, env) for a in e.args)})" if isinstance(e, cp.Case): diff --git a/sql_transform/_compose.py b/sql_transform/_compose.py index 367ccea..bd20e83 100644 --- a/sql_transform/_compose.py +++ b/sql_transform/_compose.py @@ -26,15 +26,13 @@ require_in_projection, ) from sql_transform._state import build_state_tables -from sql_transform._transformer_ref import is_transformer @dataclass(frozen=True) class Ref: - transform: object # a SQLTransform, or a fitted transformer if is_transformer + transform: object # a SQLTransform frozen: bool # True for {a.transform}; False for bare {a} expr_text: str # interpolation source, for error messages - is_transformer: bool = False @dataclass(frozen=True) @@ -62,17 +60,6 @@ def desugar_template(template: Template) -> tuple[str, dict[str, Ref]]: ref = Ref(v.__self__, frozen=True, expr_text=item.expression) elif isinstance(v, SQLTransform): ref = Ref(v, frozen=False, expr_text=item.expression) - elif is_transformer(v): - ref = Ref(v, frozen=False, expr_text=item.expression, is_transformer=True) - elif hasattr(v, "transform") and not hasattr(v, "n_features_in_"): - # Transformer-shaped but unfitted. Without this the generic TypeError - # below blames the interpolation's TYPE, which sends users looking in - # entirely the wrong place. - raise ValueError( - f"interpolation {{{item.expression}}}: {type(v).__name__} is not " - f"fitted (or does not expose n_features_in_) -- call .fit(...) " - f"before referencing it" - ) else: raise TypeError( f"interpolation {{{item.expression}}} must be a SQLTransform or " @@ -246,9 +233,8 @@ def remap(n): def _find_call(select: exp.Select, name: str, ref: Ref) -> exp.Anonymous: for n in select.find_all(exp.Anonymous): if str(n.this).upper() == name: - # Same build-time guard as the transformer-callout path (TASK-2 - # AC#1): a ref parked in QUALIFY/SORT BY/... inlines into a clause - # the native engine never resolves, so the engines would disagree. + # A ref parked in QUALIFY/SORT BY/... inlines into a clause the + # native engine never resolves, so the engines would disagree. require_in_projection( select, n, f"referenced transform {{{ref.expr_text}}}" ) diff --git a/sql_transform/_compose_test.py b/sql_transform/_compose_test.py index bf1ba89..4840588 100644 --- a/sql_transform/_compose_test.py +++ b/sql_transform/_compose_test.py @@ -44,7 +44,7 @@ def test_referenced_transform_outside_projection_raises(sql_of): # TASK-2 AC#1, sibling path: a {a.transform} ref inlines into whatever # clause it sits in. The native engine only resolves the projection, so a # ref in QUALIFY/SORT BY/CLUSTER BY meant fit() accepted a query the two - # engines then disagreed about. Same guard as the transformer-callout path. + # engines then disagreed about. train = pa.table({"age": [10.0, 20.0, 30.0, 40.0]}) inner = SQLTransform("SELECT age / MEAN(age) OVER () AS a FROM __THIS__").fit(train) with pytest.raises(ValueError, match="projection"): @@ -129,16 +129,3 @@ def test_reference_not_applied_to_column_errors(): def test_non_transform_interpolation_errors(): with pytest.raises(TypeError, match="SQLTransform"): SQLTransform(t"SELECT {42}(age) AS s FROM __THIS__") - - -def test_unfitted_transformer_raises_not_fitted_error(): - # Before: TypeError blaming the interpolation TYPE ("must be a SQLTransform"), - # which hides the real cause. An unfitted transformer has .transform but no - # n_features_in_, so we can name the actual problem. - from sklearn.preprocessing import StandardScaler - - train = pa.table({"age": [10.0, 20.0], "income": [1.0, 2.0]}) - with pytest.raises(ValueError, match="not fitted"): - SQLTransform(t"SELECT {StandardScaler()}(age, income) AS o FROM __THIS__").fit( - train - ) diff --git a/sql_transform/_interpreter.pyi b/sql_transform/_interpreter.pyi index 725d546..1dabb40 100644 --- a/sql_transform/_interpreter.pyi +++ b/sql_transform/_interpreter.pyi @@ -12,7 +12,6 @@ class InferFn: row_tables: dict[str, type[BaseModel]], static_tables: dict[str, pa.Table], output_model: type[BaseModel] | None = None, - transformers: dict[str, tuple[object, pa.Schema]] | None = None, ) -> None: ... def infer( self, diff --git a/sql_transform/_transformer_ref.py b/sql_transform/_transformer_ref.py deleted file mode 100644 index 6bb851b..0000000 --- a/sql_transform/_transformer_ref.py +++ /dev/null @@ -1,239 +0,0 @@ -"""Resolve transformer-ref placeholder calls in a desugared SELECT. - -A transformer ref desugars to a __COMPOSE_i__(arg...) call. Here we wrap a leaf -call's column args into a single named_struct (the struct arg the engines' -opaque callout expects), and derive the transformer's in/out schema by probing -.transform on the training batch. numpy lives here, mirroring _transformer_udf. -""" - -from __future__ import annotations - -import copy - -import numpy as np -import pyarrow as pa -from sqlglot import exp - -from sql_transform._sql import require_in_projection - - -def is_transformer(obj: object) -> bool: - """A FITTED sklearn-style transformer. - - Keys off `n_features_in_`, not `feature_names_in_`: sklearn sets - `n_features_in_` on any successful fit, but `feature_names_in_` only when - fitted on named data (a DataFrame). Gating on the latter would reject - perfectly good ndarray-fit transformers. Absence of `n_features_in_` means - "not fitted", which is what lets us give that its own error. - """ - return hasattr(obj, "transform") and hasattr(obj, "n_features_in_") - - -def _bind_names(obj: object, cols: list[str], name: str) -> tuple[object, list[str]]: - """Resolve the input column order the transformer expects. - - Names are metadata: they ride the named_struct as Arrow field names and BOTH - engines align on them. A transformer fitted on a DataFrame recorded its own; - one fitted on an ndarray did not, so we synthesise them from the call site - onto a copy (never mutating the user's object -- doc-8's clone contract). - - Some estimators expose feature_names_in_ as a read-only property delegating - to an inner step (Pipeline), so the synthesis can fail. Say so, instead of - letting a raw AttributeError escape from the middle of fit(). - """ - feat = getattr(obj, "feature_names_in_", None) - if feat is not None: - return obj, [str(n) for n in feat] - if len(cols) != obj.n_features_in_: - raise ValueError( - f"{name} takes {obj.n_features_in_} columns (fitted without names, so " - f"arguments bind positionally in call order), got {len(cols)}: {cols}" - ) - obj = copy.copy(obj) - try: - obj.feature_names_in_ = np.array(cols) - except AttributeError: - raise ValueError( - f"{name}: {type(obj).__name__} was fitted without feature names and does " - f"not allow them to be set. Re-fit it on a pandas DataFrame so it records " - f"feature_names_in_." - ) from None - return obj, cols - - -def _in_window_agg(node: exp.Expression) -> bool: - """Is `node` inside a window aggregate's argument?""" - p = node.parent - while p is not None: - if isinstance(p, exp.Window): - return True - p = p.parent - return False - - -def _find_call(select: exp.Select, name: str) -> exp.Anonymous: - for n in select.find_all(exp.Anonymous): - if str(n.this).upper() == name: - require_in_projection(select, n, f"transformer ref {name}") - if _in_window_agg(n): - raise ValueError( - f"{name} output cannot feed a window aggregate: aggregating over " - f"transformer output is inherently two-stage (materialise the " - f"output, then aggregate it), which needs a subquery -- " - f"SQLTransform's single-SELECT surface has none. Aggregate over an " - f"input column instead, or use a SQLTransform reference, which " - f"inlines to a scalar." - ) - return n - raise ValueError( - f"transformer ref {name} must be applied to columns, e.g. {{t}}(a, b)" - ) - - -def _named_struct(cols: list[exp.Column]) -> exp.Anonymous: - """named_struct('name', , ...) keyed by the real column name, with each - value the user's ORIGINAL column ref carried through verbatim. Its quoting is - preserved (not force-quoted), so an unquoted CamelCase ref folds in DataFusion - exactly as a hand-written query would -- matching the oracle bug-for-bug - (TASK-28). Users quote a case-sensitive column themselves: {t}("MSZoning").""" - args: list[exp.Expression] = [] - for c in cols: - args.append(exp.Literal.string(c.name)) - args.append(c.copy()) - return exp.Anonymous(this="named_struct", expressions=args) - - -def _probe( - obj: object, cols: list[str], table: pa.Table -) -> tuple[pa.Schema, pa.Schema, np.ndarray]: - """in_schema from `cols`; out_schema by probing .transform; plus the probe's - own output `y`, so a caller that needs the materialised table can build it - without running .transform a second time. - - `cols` is the STRUCT order -- it must describe the named_struct the SQL - builds, because the UDF is registered with an Exact struct signature and - Arrow struct field order is part of that type. But `.transform` is fed in - the transformer's OWN feature_names_in_ order, which is how both engines - feed it at runtime; probing in struct order would run the transform on - feature-swapped data and make a validating transformer reject a query that - both engines execute correctly. _bind_names guarantees the two orders hold - the same column set, so the reorder always resolves. - """ - in_schema = pa.schema([(c, table.schema.field(c).type) for c in cols]) - feat = [str(n) for n in obj.feature_names_in_] - x = np.column_stack([table.column(c).to_numpy(zero_copy_only=False) for c in feat]) - y = np.asarray(obj.transform(x)) - names = [str(n) for n in obj.get_feature_names_out()] - if y.ndim != 2 or y.shape[1] != len(names): - raise ValueError( - f"cannot derive out_schema for {type(obj).__name__}: expected 2-D width " - f"{len(names)}, got shape {y.shape}" - ) - out_schema = pa.schema([(n, pa.from_numpy_dtype(y.dtype)) for n in names]) - return in_schema, out_schema, y - - -def _table_from_probe(y: np.ndarray, out_schema: pa.Schema) -> pa.Table: - """The probe's output as a pa.Table shaped like out_schema, so an outer - transformer can probe on real data. Reuses `y` -- no second .transform.""" - arrays = [ - pa.array(y[:, i], type=out_schema.field(i).type) for i in range(len(out_schema)) - ] - return pa.table(arrays, schema=out_schema) - - -def resolve_transformer_refs( - select: exp.Select, tfm_refs: dict[str, object], table: pa.Table -) -> dict[str, tuple[object, pa.Schema, pa.Schema]]: - """Wrap each leaf transformer call's args into a named_struct and derive its - schema. A call whose single arg is another transformer-ref call is left - unwrapped instead -- its outer schema is probed on the inner's materialized - output, resolving innermost-first. Returns - {placeholder_name.lower(): (obj, in_schema, out_schema)}.""" - registry: dict[str, tuple[object, pa.Schema, pa.Schema]] = {} - materialized: dict[ - str, pa.Table - ] = {} # name -> this ref's output; ONLY for refs an outer consumes - resolved: set[str] = set() # every processed ref -- `materialized` is now partial - - def call_arg_ref(call: exp.Anonymous) -> str | None: - """If the call's single arg is another transformer-ref call, its name.""" - if len(call.expressions) == 1 and isinstance( - call.expressions[0], exp.Anonymous - ): - inner = str(call.expressions[0].this).upper() - if inner in tfm_refs: - return inner - return None - - # Which refs are consumed as another ref's argument? Must be computed BEFORE - # any resolution: resolve() rewrites call args into a named_struct, which - # destroys the nested-call signal call_arg_ref() reads. - consumed = { - inner - for n in tfm_refs - if (inner := call_arg_ref(_find_call(select, n))) is not None - } - - def resolve(name: str) -> None: - if name in resolved: - return - call = _find_call(select, name) - obj = tfm_refs[name] - inner = call_arg_ref(call) - if inner is not None: - resolve(inner) # innermost first - in_tbl = materialized[inner] # inner's output, real data to probe on - # The inner's materialised output order IS the natural positional - # order, so an ndarray-fit outer binds here for free. - # in_schema must describe the struct the outer ACTUALLY receives -- the - # inner's output, left unwrapped in the SQL -- not the outer's fitted - # order. The UDF is registered with an Exact struct signature and Arrow - # struct field order is part of the type, so declaring the fitted order - # desyncs it from what the SQL builds: DataFusion refuses the call while - # the native engine, which binds by name, keeps returning rows. - cols = list(in_tbl.schema.names) - obj, feat = _bind_names(obj, cols, name) - if set(cols) != set(feat): - raise ValueError( - f"{name} columns {cols} must match feature_names_in_ {feat}" - ) - in_schema, out_schema, y = _probe(obj, cols, in_tbl) - # arg is the inner call node; leave it unwrapped. - else: - if not all(isinstance(a, exp.Column) for a in call.expressions): - raise ValueError( - f"{name} args must be plain columns or another transformer " - f"ref, e.g. {{t}}(a, b) or {{t}}({{u}}(a, b))" - ) - cols = [a.name for a in call.expressions] - obj, feat = _bind_names(obj, cols, name) - if set(cols) != set(feat): - raise ValueError( - f"{name} columns {cols} must match feature_names_in_ {feat}" - ) - # in_schema describes the named_struct built on the next line from - # call.expressions, so it MUST be in CALL order. The UDF is - # registered with an Exact struct signature and Arrow struct field - # order is part of the type, so probing in fitted order silently - # desyncs the declared signature from the struct the SQL builds -- - # DataFusion then refuses the call while the native engine, which - # binds by name, keeps working. - in_schema, out_schema, y = _probe(obj, cols, table) - call.set("expressions", [_named_struct(call.expressions)]) - # Both engines fold an unquoted function-call name to lowercase before - # resolving it (DataFusion's ANSI identifier folding; Rust's - # expr_build::convert_function does `.to_lowercase()` explicitly), and - # sqlglot's generator upper-normalizes the placeholder's text on - # print regardless of its stored case. Register under the lowercase - # form so both engines' lookups actually hit it. - registry[name.lower()] = (obj, in_schema, out_schema) - resolved.add(name) - if name in consumed: - # Only an outer ref's probe reads this. A leaf has no next stage, so - # building it would be a discarded table. - materialized[name] = _table_from_probe(y, out_schema) - - for name in tfm_refs: - resolve(name) - return registry diff --git a/sql_transform/_transformer_udf.py b/sql_transform/_transformer_udf.py deleted file mode 100644 index 903e21d..0000000 --- a/sql_transform/_transformer_udf.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Wrap a fitted sklearn transformer as a vectorized DataFusion UDF. - -The oracle side of the engine transformer-callout capability: `transform` -registers this UDF so a query can call a fitted transformer by name, struct in -/ struct out. The row engine (`InferFn`) performs the identical alignment and -marshalling in Rust; the differential harness proves the two agree. - -Input is aligned to the object's `feature_names_in_` order; output is built -from the caller-declared `out_schema` (no introspection). numpy lives here and -here only -- the Rust engine imports no numpy. -""" - -from __future__ import annotations - -import numpy as np -import pyarrow as pa -from datafusion import ScalarUDF, udf - - -def check_out_schema_natural(obj: object, y: np.ndarray, out_fields: list) -> None: - """Enforce out_schema == the transform's NATURAL output dtype (TASK-2 AC#3). - - The two engines reach the declared type by different coercion rules -- - this one via a pyarrow cast (`pa.array(..., type=)`), the native `infer` - engine via pydantic coercion at model-validate. Those rules only agree - when no coercion happens, i.e. when the declared type already IS what - `.transform` produces. A mismatch is therefore a silent cross-engine - divergence (declaring int64 over a float64 transform truncates on one - side only), so refuse it rather than let the engines drift apart. - """ - natural = pa.from_numpy_dtype(y.dtype) - bad = [f.name for f in out_fields if f.type != natural] - if bad: - declared = [f.type for f in out_fields if f.name in bad] - raise ValueError( - f"{type(obj).__name__} out_schema declares {declared} for {bad}, but " - f"the natural .transform dtype is {natural}. The engines coerce to a " - f"declared type differently (pyarrow cast vs pydantic), so they only " - f"agree when the declared type IS the natural one." - ) - - -def _transformer_udf( - obj: object, - in_schema: pa.Schema, - out_schema: pa.Schema, - name: str, -) -> ScalarUDF: - """Build a vectorized struct-in/struct-out DataFusion scalar UDF. - - obj: a fitted sklearn transformer/Pipeline exposing `.transform` and - `feature_names_in_`. - in_schema: names+types of the struct the SQL feeds in (the authored - `named_struct`). Its field-name set must equal `feature_names_in_`. - out_schema: names+types of the returned struct (declared, not introspected). - name: the reserved SQL identifier the UDF is registered and called under. - - Invariant: `out_schema` must equal the transformer's NATURAL output dtype. - This engine coerces to it via a pyarrow cast (`pa.array(..., type=...)`); - the Rust `infer` engine reaches the same declared type by pydantic - coercion at model-validate time. Those two coercion rules only agree when - no real coercion happens -- i.e. when the declared type already matches - what `.transform` produces. Declaring a mismatched dtype would diverge. - """ - feature_names = [str(n) for n in obj.feature_names_in_] - in_type = pa.struct(list(in_schema)) - out_type = pa.struct(list(out_schema)) - out_fields = list(out_schema) - - def _apply(struct_array: pa.Array) -> pa.Array: - # DataFusion hands the whole batch's StructArray in one call. - cols = [ - struct_array.field(fname).to_numpy(zero_copy_only=False) - for fname in feature_names - ] - x = np.column_stack(cols) - y = np.asarray(obj.transform(x)) - check_out_schema_natural(obj, y, out_fields) - out_cols = [ - pa.array(y[:, i], type=out_fields[i].type) for i in range(len(out_fields)) - ] - return pa.StructArray.from_arrays(out_cols, fields=out_fields) - - return udf(_apply, [in_type], out_type, "immutable", name=name) diff --git a/src/expr.rs b/src/expr.rs index a08298b..4a18004 100644 --- a/src/expr.rs +++ b/src/expr.rs @@ -1,10 +1,6 @@ use pyo3::prelude::*; use pyo3::types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString}; -use std::sync::Arc; - -use crate::types::FieldType; - #[derive(Debug)] pub enum Value { Int(i64), @@ -288,17 +284,6 @@ pub enum Expr { base: Box, field: String, }, - /// Callout to an opaque, already-fitted Python transformer. `obj` is the - /// fitted sklearn object; `input_features` is its `feature_names_in_` - /// (the struct arg is reordered to this before `.transform`); - /// `output_fields` is the caller-declared output schema (marshalling - /// target, order significant); `arg` evaluates to the input `Value::Struct`. - Transform { - obj: Arc>, - input_features: Vec, - output_fields: Vec<(String, FieldType)>, - arg: Box, - }, /// SQL CASE. `arms` are (condition, result) pairs evaluated left to right; /// the first arm whose condition is `Bool(true)` wins. `default` is the ELSE /// result, or `None` when there is no ELSE (an unmatched row yields NULL). @@ -392,94 +377,6 @@ pub fn eval(expr: &Expr, row: &crate::plan::Row) -> Result { - let fields = match eval(arg, row)? { - Value::Struct(f) => f, - Value::Null => return Ok(Value::Null), - other => { - return Err(crate::plan::InterpError::Eval(format!( - "transformer argument must be a struct, got a {} value", - type_name(&other) - ))) - } - }; - // infer() already holds the GIL; attach is cheap and re-entrant, so - // eval() stays a pure-Rust signature (no py token threaded through). - Python::attach(|py| -> Result { - // Reorder the struct's fields to feature_names_in_ order, then - // build a 1-row Python list-of-lists. sklearn's check_array - // coerces it -- no numpy on the Rust side. - let mut ordered: Vec> = Vec::with_capacity(input_features.len()); - for feat in input_features { - let value = fields - .iter() - .find(|(n, _)| n == feat) - .map(|(_, v)| v) - .ok_or_else(|| { - crate::plan::InterpError::Eval(format!( - "transformer input struct is missing field '{feat}'" - )) - })?; - let py_val = value.to_pyobject(py).map_err(|e| { - crate::plan::InterpError::Eval(format!( - "marshalling transformer input failed: {e}" - )) - })?; - ordered.push(py_val); - } - let row_list = PyList::new(py, ordered).map_err(|e| { - crate::plan::InterpError::Eval(format!( - "building transformer input row failed: {e}" - )) - })?; - let x = PyList::new(py, [row_list]).map_err(|e| { - crate::plan::InterpError::Eval(format!( - "building transformer input matrix failed: {e}" - )) - })?; - let y = obj.bind(py).call_method1("transform", (x,)).map_err(|e| { - crate::plan::InterpError::Eval(format!("transformer.transform failed: {e}")) - })?; - // .tolist() turns numpy scalars into Python builtins so - // from_pyobject sees float/int/str, not opaque numpy objects. - let y_list = y.call_method0("tolist").map_err(|e| { - crate::plan::InterpError::Eval(format!( - "transformer output .tolist() failed: {e}" - )) - })?; - let y0 = y_list.get_item(0).map_err(|e| { - crate::plan::InterpError::Eval(format!( - "transformer produced no output row: {e}" - )) - })?; - // Marshal each output position through the declared field - // type. Parity invariant: the declared output dtype must equal - // the transform's natural dtype -- here it drives pydantic - // coercion at model-validate time, and the DataFusion oracle - // reaches the same type via a pyarrow cast; the two agree only - // when no real coercion is needed (see _transformer_udf.py). - let mut out = Vec::with_capacity(output_fields.len()); - for (i, (fname, ft)) in output_fields.iter().enumerate() { - let elem = y0.get_item(i).map_err(|e| { - crate::plan::InterpError::Eval(format!( - "transformer output missing position {i} for field '{fname}': {e}" - )) - })?; - let val = Value::from_pyobject_typed(&elem, &ft.base).map_err(|e| { - crate::plan::InterpError::Eval(format!( - "marshalling transformer output field '{fname}' failed: {e}" - )) - })?; - out.push((fname.clone(), val)); - } - Ok(Value::Struct(out)) - }) - } Expr::Case { arms, default } => { // Short-circuit: evaluate conditions left to right, stop at the first // Bool(true), and evaluate ONLY that arm's result. A non-matching diff --git a/src/lib.rs b/src/lib.rs index 96d9ff9..40fd6da 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -89,120 +89,16 @@ fn validate_output_model( Ok(()) } -/// A `transformers` registry entry resolved at build time: the fitted object, -/// its `feature_names_in_` (input alignment order), and the caller-declared -/// output field list (marshalling target, order significant). -struct ResolvedTransformer { - obj: std::sync::Arc>, - input_features: Vec, - output_fields: Vec<(String, types::FieldType)>, -} - -/// Reads `obj.feature_names_in_.tolist()`. Absence is a clear build error: -/// the object was fit on bare arrays, so we cannot align inputs by name. -fn read_feature_names_in(py: Python<'_>, obj: &Py) -> Result, plan::InterpError> { - let bound = obj.bind(py); - let attr = bound.getattr("feature_names_in_").map_err(|_| { - plan::InterpError::Build( - "transformer has no `feature_names_in_`; fit it on named data \ - (e.g. a pandas DataFrame) so input columns can be aligned by name" - .to_string(), - ) - })?; - attr.call_method0("tolist") - .and_then(|l| l.extract::>()) - .map_err(|e| plan::InterpError::Build(format!("could not read feature_names_in_: {e}"))) -} - -/// Rewrites every `Expr::Function` whose name is a registered transformer into -/// an `Expr::Transform`, recursing through the whole expression tree so a -/// transformer call nested inside arithmetic is still resolved. A transformer -/// call must have exactly one argument. -fn resolve_transformers( - expr: Expr, - resolved: &HashMap, -) -> Result { - match expr { - Expr::Function { name, args } => { - let mut new_args = Vec::with_capacity(args.len()); - for a in args { - new_args.push(resolve_transformers(a, resolved)?); - } - if let Some(rt) = resolved.get(&name) { - if new_args.len() != 1 { - return Err(plan::InterpError::Build(format!( - "transformer '{name}' takes exactly one argument, got {}", - new_args.len() - ))); - } - let arg = new_args.into_iter().next().unwrap(); - return Ok(Expr::Transform { - obj: rt.obj.clone(), - input_features: rt.input_features.clone(), - output_fields: rt.output_fields.clone(), - arg: Box::new(arg), - }); - } - Ok(Expr::Function { name, args: new_args }) - } - Expr::BinaryOp { op, left, right } => Ok(Expr::BinaryOp { - op, - left: Box::new(resolve_transformers(*left, resolved)?), - right: Box::new(resolve_transformers(*right, resolved)?), - }), - Expr::Not(inner) => Ok(Expr::Not(Box::new(resolve_transformers(*inner, resolved)?))), - Expr::Cast { expr, target } => Ok(Expr::Cast { - expr: Box::new(resolve_transformers(*expr, resolved)?), - target, - }), - Expr::Struct(fields) => { - let mut out = Vec::with_capacity(fields.len()); - for (k, v) in fields { - out.push((k, resolve_transformers(v, resolved)?)); - } - Ok(Expr::Struct(out)) - } - Expr::List(items) => { - let mut out = Vec::with_capacity(items.len()); - for e in items { - out.push(resolve_transformers(e, resolved)?); - } - Ok(Expr::List(out)) - } - Expr::FieldAccess { base, field } => Ok(Expr::FieldAccess { - base: Box::new(resolve_transformers(*base, resolved)?), - field, - }), - Expr::Case { arms, default } => { - let mut new_arms = Vec::with_capacity(arms.len()); - for (cond, result) in arms { - new_arms.push(( - resolve_transformers(cond, resolved)?, - resolve_transformers(result, resolved)?, - )); - } - let default = match default { - Some(d) => Some(Box::new(resolve_transformers(*d, resolved)?)), - None => None, - }; - Ok(Expr::Case { arms: new_arms, default }) - } - // Column, Literal, and an already-built Transform pass through. - other => Ok(other), - } -} - #[pymethods] impl InferFn { #[new] - #[pyo3(signature = (sql, row_tables, static_tables, output_model=None, transformers=None))] + #[pyo3(signature = (sql, row_tables, static_tables, output_model=None))] fn new( py: Python<'_>, sql: String, row_tables: HashMap>, static_tables: HashMap>, output_model: Option>, - transformers: Option, Py)>>, ) -> PyResult { let raw_plan = plan::build_plan(&sql)?; let row_table_names: HashSet = row_tables.keys().cloned().collect(); @@ -225,34 +121,6 @@ impl InferFn { &static_schemas, )?; - // Resolve registered transformers AFTER column validation (so - // validate_columns sees the plain Expr::Function and its named_struct - // arg -- no Transform arm needed there) and BEFORE output-model - // synthesis (so infer_type sees Expr::Transform and returns the - // declared output struct type). Reads feature_names_in_ here (with py). - let transformers = transformers.unwrap_or_default(); - if !transformers.is_empty() { - let mut resolved: HashMap = HashMap::new(); - for (name, (obj, out_schema_obj)) in &transformers { - let input_features = read_feature_names_in(py, obj)?; - let output_fields = schema::arrow_schema_to_ordered_fields(py, out_schema_obj)?; - resolved.insert( - name.clone(), - ResolvedTransformer { - obj: std::sync::Arc::new(obj.clone_ref(py)), - input_features, - output_fields, - }, - ); - } - let projection = std::mem::take(&mut optimized_plan.projection); - let mut new_projection = Vec::with_capacity(projection.len()); - for (alias, expr) in projection { - new_projection.push((alias, resolve_transformers(expr, &resolved)?)); - } - optimized_plan.projection = new_projection; - } - let output_model = match output_model { Some(supplied) => { validate_output_model( diff --git a/src/plan.rs b/src/plan.rs index 027ddcd..880a775 100644 --- a/src/plan.rs +++ b/src/plan.rs @@ -1173,19 +1173,6 @@ fn validate_expr( } Ok(()) } - // Compiler-mandated: `Expr` exhaustiveness is checked at compile time - // and this match has no catch-all. Transformer resolution runs AFTER - // validate_columns, so no `Transform` node reaches here today; recurse - // into `arg` anyway so column validation stays correct if that ordering - // ever changes. - Expr::Transform { arg, .. } => validate_expr( - arg, - resolved, - row_schemas, - static_schemas, - effective_schemas, - used_columns, - ), Expr::FieldAccess { base, field } => { validate_expr( base, diff --git a/src/schema.rs b/src/schema.rs index 2a41384..4376452 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -77,37 +77,6 @@ pub fn from_arrow_table(py: Python<'_>, table: &Py) -> Result, - schema_obj: &Py, -) -> Result, InterpError> { - let bound = schema_obj.bind(py); - let names: Vec = bound - .getattr("names") - .and_then(|n| n.extract()) - .map_err(|e| { - InterpError::Build(format!("transformer output schema is not a pyarrow.Schema: {e}")) - })?; - let pa_types = PyModule::import(py, "pyarrow.types") - .map_err(|e| InterpError::Build(format!("Failed to import pyarrow.types: {e}")))?; - let mut out = Vec::with_capacity(names.len()); - for name in names { - let field = bound - .call_method1("field", (name.as_str(),)) - .map_err(|e| InterpError::Build(format!("Failed to read output field '{name}': {e}")))?; - let ft = arrow_field_to_field_type(&pa_types, &field) - .map_err(|e| InterpError::Build(format!("Failed to read type of output field '{name}': {e}")))?; - out.push((name, ft)); - } - Ok(out) -} - /// Recursively resolves a `pyarrow.Field` (name/nullable/type) into a /// `FieldType`, walking `pa.StructType`/`pa.ListType` children rather than /// prefix-matching the type's string form — needed since a struct/list type's diff --git a/src/types.rs b/src/types.rs index 48adad7..59b57fd 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; #[derive(Clone, PartialEq, Eq, Debug)] pub enum Base { @@ -101,38 +101,6 @@ pub fn infer_type( ))), } } - Expr::Transform { - input_features, - output_fields, - arg, - .. - } => { - let arg_ty = infer_type(arg, schemas)?; - match &arg_ty.base { - Base::Struct(fields) => { - let got: HashSet<&String> = fields.iter().map(|(n, _)| n).collect(); - let want: HashSet<&String> = input_features.iter().collect(); - if got != want { - return Err(InterpError::Build(format!( - "transformer input struct fields {:?} do not match \ - feature_names_in_ {:?}", - fields.iter().map(|(n, _)| n).collect::>(), - input_features - ))); - } - } - _ => { - return Err(InterpError::Build( - "transformer argument must be a struct (e.g. named_struct(...))" - .to_string(), - )) - } - } - Ok(FieldType { - base: Base::Struct(output_fields.clone()), - nullable: false, - }) - } Expr::Case { arms, default } => { let mut branch_types: Vec = arms .iter() diff --git a/tests/test_diff_errors.py b/tests/test_diff_errors.py index 47c16e5..d9d47e1 100644 --- a/tests/test_diff_errors.py +++ b/tests/test_diff_errors.py @@ -32,9 +32,9 @@ def test_int_mod_by_zero_rejected_by_both(): def test_unknown_function_rejected_by_both(): - # foobar() is invalid SQL -- DataFusion errors -- and is NOT a `__tfm_` - # transformer placeholder, so codegen must raise a hard error too (not - # UnsupportedInCodegen, which the harness would skip as a deferred surface). + # foobar() is invalid SQL -- DataFusion errors -- so codegen must raise a + # hard error too, not UnsupportedInCodegen (which the harness would skip as + # a deferred surface). check_both_raise( "SELECT foobar(x) AS y FROM t", {"t": rows({"x": "int"}, [{"x": 1}])}, diff --git a/tests/test_diff_transformer_callout.py b/tests/test_diff_transformer_callout.py deleted file mode 100644 index 5811d24..0000000 --- a/tests/test_diff_transformer_callout.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Differential parity: transform (DataFusion UDF) == infer (Rust Expr::Transform).""" - -import numpy as np -import pandas as pd -import pyarrow as pa -import pytest -from datafusion import SessionContext -from differential import _rows_equal -from sklearn.pipeline import Pipeline -from sklearn.preprocessing import OrdinalEncoder, PowerTransformer, StandardScaler - -from sql_transform._interpreter import InferFn -from sql_transform._schema import synthesize_this_model -from sql_transform._transformer_udf import _transformer_udf - -# See test_transformer_udf.py: the nameless-input warning is a known false -# positive; both the oracle UDF and the Rust infer path emit it. -pytestmark = pytest.mark.filterwarnings( - "ignore:X does not have valid feature names:UserWarning" -) - - -def _parity(sql, table, obj, in_schema, out_schema, name="__tfm_0__"): - # Oracle: DataFusion with the transformer registered as a UDF. - ctx = SessionContext() - ctx.from_arrow(table, name="__THIS__") - ctx.register_udf(_transformer_udf(obj, in_schema, out_schema, name)) - df = ctx.sql(sql) - oracle = pa.Table.from_batches(df.collect(), schema=df.schema()).to_pylist() - - # Rust infer with the same object in the transformers registry. - model = synthesize_this_model(table.schema) - rows = [model(**r) for r in table.to_pylist()] - fn = InferFn( - sql, - row_tables={"__THIS__": model}, - static_tables={}, - transformers={name: (obj, out_schema)}, - ) - actual = [r.model_dump() for r in fn.infer({"__THIS__": rows})] - - assert _rows_equal(actual, oracle), (sql, actual, oracle) - return oracle - - -def test_standard_scaler_parity(): - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) - schema = pa.schema([("age", pa.float64()), ("income", pa.float64())]) - sql = ( - "SELECT __tfm_0__(named_struct('age', age, 'income', income)) " - "AS s FROM __THIS__" - ) - _parity(sql, pa.Table.from_pandas(train), sc, schema, schema) - - -def test_struct_field_order_independence_parity(): - # The named_struct is authored in the OPPOSITE order to feature_names_in_ - # (income, age vs age, income). Both engines must align by NAME, not - # position -- a positional bug in either would diverge here. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) # feature_names_in_ == [age, income] - in_schema = pa.schema([("income", pa.float64()), ("age", pa.float64())]) - out_schema = pa.schema([("age", pa.float64()), ("income", pa.float64())]) - sql = ( - "SELECT __tfm_0__(named_struct('income', income, 'age', age)) " - "AS s FROM __THIS__" - ) - _parity(sql, pa.Table.from_pandas(train), sc, in_schema, out_schema) - - -def test_whole_pipeline_parity(): - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - pipe = Pipeline([("sc", StandardScaler()), ("pt", PowerTransformer())]).fit(train) - schema = pa.schema([("age", pa.float64()), ("income", pa.float64())]) - sql = ( - "SELECT __tfm_0__(named_struct('age', age, 'income', income)) " - "AS s FROM __THIS__" - ) - _parity(sql, pa.Table.from_pandas(train), pipe, schema, schema) - - -def test_single_field_one_in_one_out_parity(): - # TASK-2 AC#2. Every other case here is 2-in/2-out; the marshalling on both - # sides indexes y[:, i], which assumes .transform returns 2-D. A 1-in/1-out - # transformer is the degenerate width -- (n,1), not (n,) -- and is the case - # a positional/squeeze bug in either engine would break first. - train = pd.DataFrame({"age": [10.0, 20.0, 30.0, 40.0]}) - sc = StandardScaler().fit(train) - schema = pa.schema([("age", pa.float64())]) - sql = "SELECT __tfm_0__(named_struct('age', age)) AS s FROM __THIS__" - oracle = _parity(sql, pa.Table.from_pandas(train), sc, schema, schema) - assert np.allclose([[r["s"]["age"]] for r in oracle], sc.transform(train)) - - -def test_ordinal_encoder_non_float_in_and_out_parity(): - train = pd.DataFrame( - {"color": ["red", "green", "blue", "red"], "size": ["S", "M", "L", "M"]} - ) - enc = OrdinalEncoder(dtype=np.int64).fit(train) # string in, int out - in_schema = pa.schema([("color", pa.string()), ("size", pa.string())]) - out_schema = pa.schema([("color", pa.int64()), ("size", pa.int64())]) - # pa.Table.from_pandas defaults object-dtype columns to large_string, - # which doesn't match the Utf8 in_schema the UDF is registered with; cast - # to align the physical table type with the declared schema. - table = pa.Table.from_pandas(train).cast(in_schema) - sql = ( - "SELECT __tfm_0__(named_struct('color', color, 'size', size)) " - "AS s FROM __THIS__" - ) - _parity(sql, table, enc, in_schema, out_schema) diff --git a/tests/test_transformer_callout_infer.py b/tests/test_transformer_callout_infer.py deleted file mode 100644 index 9837625..0000000 --- a/tests/test_transformer_callout_infer.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Direct `infer` tests for the Expr::Transform callout (no oracle). - -Parity with DataFusion is proven separately in test_diff_transformer_callout.py; -these pin the Rust marshalling to hand-computed values and the build-time errors. -""" - -import numpy as np -import pandas as pd -import pyarrow as pa -import pytest -from sklearn.preprocessing import StandardScaler - -from sql_transform._interpreter import InferFn -from sql_transform._schema import synthesize_this_model - -# See test_transformer_udf.py: the nameless-input warning is a known false positive. -pytestmark = pytest.mark.filterwarnings( - "ignore:X does not have valid feature names:UserWarning" -) - -_THIS = pa.schema([("age", pa.float64()), ("income", pa.float64())]) -_OUT = pa.schema([("age", pa.float64()), ("income", pa.float64())]) -_SQL = "SELECT __tfm_0__(named_struct('age', age, 'income', income)) AS s FROM __THIS__" - - -def _fitted_scaler(): - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - return StandardScaler().fit(train), train - - -def _infer(sql, transformers, rows): - model = synthesize_this_model(_THIS) - fn = InferFn( - sql, - row_tables={"__THIS__": model}, - static_tables={}, - transformers=transformers, - ) - return [r.model_dump() for r in fn.infer({"__THIS__": [model(**r) for r in rows]})] - - -def test_standard_scaler_infer_hand_computed(): - sc, _ = _fitted_scaler() - out = _infer(_SQL, {"__tfm_0__": (sc, _OUT)}, [{"age": 10.0, "income": 1.0}]) - # population std (ddof=0): age mean 25, std 11.18034; income mean 2.5, std 1.11803 - assert abs(out[0]["s"]["age"] - ((10.0 - 25.0) / 11.180339887498949)) < 1e-9 - assert abs(out[0]["s"]["income"] - ((1.0 - 2.5) / 1.1180339887498949)) < 1e-9 - - -def test_missing_feature_names_in_is_build_error(): - sc = StandardScaler().fit(np.array([[10.0, 1.0], [20.0, 2.0]])) # bare array - with pytest.raises(ValueError, match="feature_names_in_"): - _infer(_SQL, {"__tfm_0__": (sc, _OUT)}, [{"age": 10.0, "income": 1.0}]) - - -def test_non_struct_argument_is_build_error(): - sc, _ = _fitted_scaler() - sql = "SELECT __tfm_0__(age) AS s FROM __THIS__" - with pytest.raises(ValueError, match="must be a struct"): - _infer(sql, {"__tfm_0__": (sc, _OUT)}, [{"age": 10.0, "income": 1.0}]) - - -def test_field_name_mismatch_is_build_error(): - sc, _ = _fitted_scaler() - sql = ( - "SELECT __tfm_0__(named_struct('age', age, 'wrong', income)) AS s FROM __THIS__" - ) - with pytest.raises(ValueError, match="feature_names_in_"): - _infer(sql, {"__tfm_0__": (sc, _OUT)}, [{"age": 10.0, "income": 1.0}]) diff --git a/tests/test_transformer_case.py b/tests/test_transformer_case.py deleted file mode 100644 index 68d87ad..0000000 --- a/tests/test_transformer_case.py +++ /dev/null @@ -1,65 +0,0 @@ -"""TASK-31: a transformer reference nested inside a CASE branch resolves. - -`resolve_transformers` (src/lib.rs) recurses through the expression tree to -rewrite each `foo(...)` transformer call into an `Expr::Transform`. Its CASE arm -has a *silent* failure mode: if it ever stopped recursing into every branch, a -transformer call in a non-first branch would simply never be rewritten, and the -engine would then fail (or worse, silently mishandle) at eval time. - -The guard here puts the transformer in a NON-first WHEN arm AND in the ELSE -(both non-first branches, exercised by the arms loop beyond index 0 and by the -separate default recursion). If either recursion regressed, the g=2 row (arm 1) -or the g=3 row (ELSE) would raise "Unknown function __tfm_0__" instead of -matching the DataFusion oracle -- so this test genuinely FAILS on the -regression, not just passes on the happy path. - -Transformers are a native-only feature (codegen has no transformer support and -`SQLTransform` runs on the native InferFn), so parity is asserted native-vs-oracle; -the codegen side is covered by asserting it defers the construct loudly rather -than mishandling it silently. -""" - -import pandas as pd -import pyarrow as pa -import pytest -from sklearn.preprocessing import StandardScaler -from test_diff_transformer_callout import _parity - -from sql_transform._codegen import CodegenFn, UnsupportedInCodegen -from sql_transform._schema import synthesize_this_model - -# Same known false-positive warning as the other transformer tests: both the -# oracle UDF and the native infer path emit "X does not have valid feature names". -pytestmark = pytest.mark.filterwarnings( - "ignore:X does not have valid feature names:UserWarning" -) - -# A transformer in a non-first WHEN arm (g=2) and in the ELSE (g=3); a plain -# named_struct in the first arm (g=1) so the transformer is genuinely not the -# first branch. All three branches are struct{age: float}, so the CASE type-checks. -_SQL = ( - "SELECT CASE " - "WHEN g = 1 THEN named_struct('age', 0.0) " - "WHEN g = 2 THEN __tfm_0__(named_struct('age', age)) " - "ELSE __tfm_0__(named_struct('age', age)) " - "END AS s FROM __THIS__" -) - - -def test_transformer_in_nonfirst_case_branch_parity_native(): - # Rows hit each branch: g=1 -> plain arm, g=2 -> transformer in arm 1, - # g=3 -> transformer in ELSE. If resolve_transformers stopped recursing into - # any branch, g=2 or g=3 would raise instead of matching the oracle. - train = pd.DataFrame({"g": [1, 2, 3], "age": [10.0, 20.0, 30.0]}) - sc = StandardScaler().fit(train[["age"]]) - schema = pa.schema([("age", pa.float64())]) - _parity(_SQL, pa.Table.from_pandas(train), sc, schema, schema) - - -def test_transformer_in_case_defers_on_codegen(): - # Codegen has no transformer support, so the __tfm_0__(...) call defers - # (named_struct itself is supported now); the construct must defer loudly - # (UnsupportedInCodegen), never silently mishandle. - model = synthesize_this_model(pa.schema([("g", pa.int64()), ("age", pa.float64())])) - with pytest.raises(UnsupportedInCodegen): - CodegenFn(_SQL, row_tables={"__THIS__": model}, static_tables={}) diff --git a/tests/test_transformer_ref.py b/tests/test_transformer_ref.py deleted file mode 100644 index a867bc4..0000000 --- a/tests/test_transformer_ref.py +++ /dev/null @@ -1,540 +0,0 @@ -"""Differential parity for fitted transformers referenced as {ref} in a t-string.""" - -import numpy as np -import pandas as pd -import pyarrow as pa -import pytest -from differential import _rows_equal -from sklearn.decomposition import PCA -from sklearn.preprocessing import OrdinalEncoder, StandardScaler - -from sql_transform import SQLTransform, _transformer_ref - -# The nameless-input warning is a known false positive (see test_transformer_udf). -pytestmark = pytest.mark.filterwarnings( - "ignore:X does not have valid feature names:UserWarning" -) - - -def _both_engines(t, test_df): - """transform (DataFusion) and infer (Rust) as plain dicts; assert equal via - the shared _rows_equal helper. - - transform (DataFusion) and infer_batch (Rust) take different numerical - paths -- transform stacks the whole batch into one matrix, infer_batch runs - row-at-a-time -- so their floats can differ in the last bit or two (pure - ULP noise, e.g. a non-collinear PCA fixture measured at ~2e-16). _rows_equal - is the project's one place that encodes both the tolerance for that ULP - noise AND the type-strictness (bool != int, int != float) that cross-engine - comparison requires -- exact `==` would flake on the ULP noise, and a naive - tolerant comparison would silently let a real type divergence (a genuine - bug class here) through as a pass. - """ - batch = t.transform(pa.Table.from_pandas(test_df)).to_pylist() - infer = [r.model_dump() for r in t.infer_batch(test_df.to_dict("records"))] - assert _rows_equal(infer, batch), (infer, batch) - return batch - - -def test_single_scaler_ref_parity(): - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) - t = SQLTransform(t"SELECT {sc}(age, income) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - - test = pd.DataFrame({"age": [25.0, 35.0], "income": [2.5, 3.5]}) - batch = _both_engines(t, test) - - expected = sc.transform(test) - got = np.array([[b["out"]["age"], b["out"]["income"]] for b in batch]) - assert np.allclose(got, expected) - - -@pytest.mark.parametrize( - "sql_of", - [ - pytest.param( - lambda sc: t"SELECT age FROM __THIS__ QUALIFY {sc}(age) > 1", id="qualify" - ), - pytest.param( - lambda sc: t"SELECT DISTINCT ON ({sc}(age)) age FROM __THIS__", - id="distinct_on", - ), - pytest.param( - lambda sc: t"SELECT age FROM __THIS__ CLUSTER BY {sc}(age)", id="cluster_by" - ), - pytest.param( - lambda sc: t"SELECT age FROM __THIS__ SORT BY {sc}(age)", id="sort_by" - ), - pytest.param( - lambda sc: ( - t"SELECT AVG(age) OVER w AS a FROM __THIS__ " - t"WINDOW w AS (PARTITION BY {sc}(age))" - ), - id="window_clause", - ), - ], -) -def test_transformer_call_outside_projection_raises(sql_of): - # TASK-2 AC#1. The native engine resolves transformer calls only over the - # projection (src/lib.rs); DataFusion plans the whole statement. These five - # clauses survive parse_and_validate, so before this guard fit() ACCEPTED - # them and the disagreement surfaced late and asymmetrically: transform() - # raised a DataFusion planning error while infer() silently ignored the - # clause and returned rows (WINDOW crashed at fit with a bare - # AttributeError). Reject at build time with one clear message instead. - # ORDER BY/WHERE/GROUP BY/HAVING/LIMIT/JOIN are already rejected upstream - # by parse_and_validate, so they cannot reach this guard. - train = pd.DataFrame({"age": [10.0, 20.0, 30.0, 40.0]}) - sc = StandardScaler().fit(train) - t = SQLTransform(sql_of(sc)) - with pytest.raises(ValueError, match="projection"): - t.fit(pa.Table.from_pandas(train)) - - -def test_nested_threading_parity(): - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) - # Wrap as a DataFrame (not the bare ndarray sc.transform returns) so PCA's - # fit records feature_names_in_ == sc.get_feature_names_out() -- required - # for is_transformer(pca) to hold and for the nested schema match. - scaled = pd.DataFrame(sc.transform(train), columns=sc.get_feature_names_out()) - pca = PCA(n_components=1).fit(scaled) - t = SQLTransform(t"SELECT {pca}({sc}(age, income)) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - - test = pd.DataFrame({"age": [25.0, 35.0], "income": [2.5, 3.5]}) - batch = _both_engines(t, test) - - expected = pca.transform(sc.transform(test)) - out_names = [str(n) for n in pca.get_feature_names_out()] - got = np.array([[b["out"][n] for n in out_names] for b in batch]) - assert np.allclose(got, expected) - - -def test_ordinal_encoder_ref_string_in_int_out(): - train = pd.DataFrame( - {"color": ["red", "green", "blue", "red"], "size": ["S", "M", "L", "M"]} - ) - enc = OrdinalEncoder(dtype=np.int64).fit(train) - t = SQLTransform(t"SELECT {enc}(color, size) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - - test = pd.DataFrame({"color": ["blue", "red"], "size": ["L", "S"]}) - batch = _both_engines(t, test) - - expected = enc.transform(test) - got = np.array([[b["out"]["color"], b["out"]["size"]] for b in batch]) - assert (got == expected).all() - - -def test_transformer_and_native_window_agg_coexist(): - # A native window agg over __THIS__ alongside a transformer call: proves the - # two compose without either engine choking. No agg reads the transformer output. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train[["age", "income"]]) - t = SQLTransform( - t"SELECT {sc}(age, income) AS out, age / (MEAN(age) OVER ()) AS an FROM __THIS__" # noqa: E501 - ).fit(pa.Table.from_pandas(train)) - - test = pd.DataFrame({"age": [25.0, 35.0], "income": [2.5, 3.5]}) - batch = _both_engines(t, test) - assert np.allclose( - [b["an"] for b in batch], test["age"].to_numpy() / train["age"].mean() - ) - - -def test_camelcase_columns_quoted_compose(): - # A transformer ref on case-sensitive columns works when the user quotes them - # (DataFusion keeps a quoted identifier case-exact). The ref carries the - # quoting through verbatim, so both engines resolve `"LotArea"` identically. - train = pd.DataFrame( - {"LotArea": [1.0, 2.0, 3.0, 4.0], "YearBuilt": [1990.0, 2000.0, 2010.0, 2020.0]} - ) - sc = StandardScaler().fit(train) - t = SQLTransform(t'SELECT {sc}("LotArea", "YearBuilt") AS out FROM __THIS__').fit( - pa.Table.from_pandas(train) - ) - - test = pd.DataFrame({"LotArea": [2.5], "YearBuilt": [2005.0]}) - batch = _both_engines(t, test) - - expected = sc.transform(test) - got = np.array([[b["out"]["LotArea"], b["out"]["YearBuilt"]] for b in batch]) - assert np.allclose(got, expected) - - -def test_camelcase_columns_unquoted_folds_and_fails(): - # Unquoted CamelCase folds to lowercase in DataFusion (the oracle) and misses, - # so the transformer-ref path must NOT silently make it work (the reverted - # TASK-25 force-quote did). The user quotes instead (see the quoted test). - # TASK-28. - train = pd.DataFrame( - {"LotArea": [1.0, 2.0, 3.0, 4.0], "YearBuilt": [1990.0, 2000.0, 2010.0, 2020.0]} - ) - sc = StandardScaler().fit(train) - t = SQLTransform(t"SELECT {sc}(LotArea, YearBuilt) AS out FROM __THIS__") - with pytest.raises(Exception, match="(?i)lotarea|no field|schema"): - fitted = t.fit(pa.Table.from_pandas(train)) - fitted.transform(pa.Table.from_pandas(train)) - - -def test_ndarray_fit_transformer_binds_positionally(): - # sklearn records feature_names_in_ only for DataFrame fit. An ndarray-fit - # transformer has no names, so arguments bind positionally in call order. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train.to_numpy()) - assert not hasattr(sc, "feature_names_in_") - - t = SQLTransform(t"SELECT {sc}(age, income) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - batch = _both_engines(t, train) - expected = sc.transform(train.to_numpy()) - got = np.array([[b["out"]["age"], b["out"]["income"]] for b in batch]) - assert np.allclose(got, expected) - - # clone contract: the user's object must be left untouched - assert not hasattr(sc, "feature_names_in_") - - -def test_ndarray_fit_arity_mismatch_raises(): - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train.to_numpy()) # n_features_in_ == 2 - t = SQLTransform(t"SELECT {sc}(age) AS out FROM __THIS__") - with pytest.raises(ValueError, match="bind positionally"): - t.fit(pa.Table.from_pandas(train)) - - -def test_named_fit_column_mismatch_still_raises(): - # The named path is unchanged: names are validated as a set. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) - t = SQLTransform(t"SELECT {sc}(age) AS out FROM __THIS__") - with pytest.raises(ValueError, match="must match feature_names_in_"): - t.fit(pa.Table.from_pandas(train)) - - -def test_ndarray_fit_outer_in_nested_ref_works(): - # The nested branch used to read feature_names_in_ unconditionally, so an - # ndarray-fit OUTER died with a raw AttributeError. The inner's materialised - # output order is the natural positional order, so binding it works. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - a = StandardScaler().fit(train) - a_out = pd.DataFrame(a.transform(train), columns=a.get_feature_names_out()) - b = PCA(n_components=1).fit(a_out.to_numpy()) # ndarray fit: no names - assert not hasattr(b, "feature_names_in_") - - t = SQLTransform(t"SELECT {b}({a}(age, income)) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - batch = _both_engines(t, train) - expected = b.transform(a.transform(train)) - out_names = [str(n) for n in b.get_feature_names_out()] - got = np.array([[r["out"][n] for n in out_names] for r in batch]) - assert np.allclose(got, expected) - assert not hasattr(b, "feature_names_in_") # clone contract - - -def test_unsettable_feature_names_gives_actionable_error(): - # Pipeline.feature_names_in_ is a read-only property delegating to steps[0]. - # Synthesising names onto it raises AttributeError; the user must see an - # actionable message, not a raw attribute error from inside fit(). - from sklearn.pipeline import Pipeline - - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - p = Pipeline([("sc", StandardScaler())]).fit(train.to_numpy()) - t = SQLTransform(t"SELECT {p}(age, income) AS o FROM __THIS__") - with pytest.raises(ValueError, match="Re-fit it on a pandas DataFrame"): - t.fit(pa.Table.from_pandas(train)) - - -class _SpyTransformer: - """Wraps a fitted transformer and counts .transform() calls.""" - - def __init__(self, obj): - self._obj = obj - self.calls = 0 - self.feature_names_in_ = obj.feature_names_in_ - self.n_features_in_ = obj.n_features_in_ - - def transform(self, x): - self.calls += 1 - return self._obj.transform(x) - - def get_feature_names_out(self): - return self._obj.get_feature_names_out() - - -def test_leaf_ref_probes_transform_once(): - # A leaf ref's materialised output is only ever read by an OUTER ref's probe. - # With no outer, materialising it is a discarded .transform() call. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - spy = _SpyTransformer(StandardScaler().fit(train)) - - SQLTransform(t"SELECT {spy}(age, income) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - assert spy.calls == 1, f"expected a single probe, got {spy.calls}" - - -def test_nested_refs_probe_once_each(): - # The inner IS consumed, so it must still be materialised -- but from the - # probe's own output, not a second .transform() call. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) - scaled = pd.DataFrame(sc.transform(train), columns=sc.get_feature_names_out()) - inner = _SpyTransformer(sc) - outer = _SpyTransformer(PCA(n_components=1).fit(scaled)) - - SQLTransform(t"SELECT {outer}({inner}(age, income)) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - assert (inner.calls, outer.calls) == (1, 1), (inner.calls, outer.calls) - - -def test_leaf_ref_skips_materialization(monkeypatch): - # The `consumed` gate saves a _table_from_probe call (a table allocation), - # not a .transform() call -- _table_from_probe reuses the probe's `y` and - # never calls .transform. So the spy tests above cannot see this gate: they - # count .transform(), which is identical whether or not the gate exists. - # Count _table_from_probe calls directly instead. - calls = [] - orig = _transformer_ref._table_from_probe - - def counting(y, out_schema): - calls.append(1) - return orig(y, out_schema) - - monkeypatch.setattr(_transformer_ref, "_table_from_probe", counting) - - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) - - # Leaf-only: no outer consumer, so the materialised table is never built. - SQLTransform(t"SELECT {sc}(age, income) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - assert len(calls) == 0, f"expected 0 materializations for a leaf, got {len(calls)}" - - calls.clear() - scaled = pd.DataFrame(sc.transform(train), columns=sc.get_feature_names_out()) - pca = PCA(n_components=1).fit(scaled) - - # Nested: the inner ref IS consumed by the outer, so it must be built once. - SQLTransform(t"SELECT {pca}({sc}(age, income)) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - assert len(calls) == 1, ( - f"expected 1 materialization for the inner, got {len(calls)}" - ) - - -def test_aggregate_over_transformer_output_raises(): - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) - t = SQLTransform(t"SELECT AVG({sc}(age, income)) OVER () AS m FROM __THIS__") - with pytest.raises(ValueError, match="two-stage"): - t.fit(pa.Table.from_pandas(train)) - - -def test_aggregate_over_sqltransform_ref_still_works(): - # The guard must NOT overreach. A SQLTransform ref inlines to a scalar, so an - # aggregate over it is ordinary flat SQL -- a shipped, documented capability. - train = pa.table({"age": [10.0, 20.0, 30.0, 40.0]}) - inner = SQLTransform("SELECT age / MEAN(age) OVER () AS a FROM __THIS__").fit( - pa.table({"age": [1.0, 2.0, 3.0]}) - ) - t = SQLTransform( - t"SELECT AVG({inner.transform}(age)) OVER () AS m FROM __THIS__" - ).fit(train) - assert t.transform(train).column("m").to_pylist() == [12.5, 12.5, 12.5, 12.5] - - -def test_mixed_leaf_and_nested_args_raises(): - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) - scaled = pd.DataFrame(sc.transform(train), columns=sc.get_feature_names_out()) - pca = PCA(n_components=2).fit(scaled) - t = SQLTransform(t"SELECT {pca}({sc}(age, income), age) AS o FROM __THIS__") - with pytest.raises(ValueError, match="plain columns or another transformer ref"): - t.fit(pa.Table.from_pandas(train)) - - -def test_transformer_alongside_partitioned_window_agg(): - # Lock-in: a transformer callout and a PARTITION BY window agg coexist, and - # both engines agree. Currently works; nothing must silently break it. - train = pd.DataFrame( - { - "age": [10.0, 20.0, 30.0, 40.0], - "income": [1.0, 2.0, 3.0, 4.0], - "city": ["a", "a", "b", "b"], - } - ) - sc = StandardScaler().fit(train[["age", "income"]]) - t = SQLTransform( - t"SELECT {sc}(age, income) AS o, AVG(age) OVER (PARTITION BY city) AS m " - t"FROM __THIS__" - ).fit(pa.Table.from_pandas(train)) - - batch = t.transform(pa.Table.from_pandas(train)).to_pylist() - infer = [r.model_dump() for r in t.infer_batch(train.to_dict("records"))] - # `o` is StandardScaler output -- same batch-column_stack-vs-row-at-a-time - # risk class as PCA -- so compare with the tolerance helper, not exact - # equality, even though this particular fixture happens to land exact. - assert _rows_equal(infer, batch), (infer, batch) - assert [r["m"] for r in batch] == [15.0, 15.0, 35.0, 35.0] - - -def test_unfit_ref_is_fitted_once_globally_not_per_partition(): - # An unfit ref under an outer PARTITION BY is fitted ONCE over all rows; the - # partitioning applies to the outer aggregate over its output. Matches - # sklearn, where a Pipeline step is fitted once on all training data. - # Per-group fitting is a separate feature (DRAFT-14), not this. - train = pa.table({"age": [10.0, 20.0, 30.0, 50.0], "city": ["a", "a", "b", "b"]}) - norm = SQLTransform("SELECT age / MEAN(age) OVER () AS a FROM __THIS__") - t = SQLTransform( - t"SELECT AVG({norm}(age)) OVER (PARTITION BY city) AS m FROM __THIS__" - ).fit(train) - - # global mean 27.5, NOT per-city 15.0 / 40.0 - assert t._state_tables["__STATE_R0__"].column("avg_age").to_pylist() == [27.5] - got = t.transform(train).column("m").to_pylist() - assert np.allclose(got, [0.5454545, 0.5454545, 1.4545455, 1.4545455]) - - -def test_three_level_nesting_parity(): - # AC#4. Load-bearing after Task 2: the `consumed` set decides materialisation - # by nesting position, and 3 levels is the only shape where a ref is - # simultaneously consumed AND a consumer. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - a = StandardScaler().fit(train) - a_out = pd.DataFrame(a.transform(train), columns=a.get_feature_names_out()) - b = StandardScaler().fit(a_out) - b_out = pd.DataFrame(b.transform(a_out), columns=b.get_feature_names_out()) - c = PCA(n_components=1).fit(b_out) - - t = SQLTransform(t"SELECT {c}({b}({a}(age, income))) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - batch = _both_engines(t, train) - expected = c.transform(b.transform(a.transform(train))) - out_names = [str(n) for n in c.get_feature_names_out()] - got = np.array([[r["out"][n] for n in out_names] for r in batch]) - assert np.allclose(got, expected) - - -def test_named_fit_call_order_is_free(): - # A DataFrame-fitted transformer binds BY NAME, so the SQL may list columns - # in any order -- a capability the README documents. The regression this - # guards: in_schema was probed in FITTED order while the named_struct was - # built in CALL order, so the UDF's Exact struct signature stopped matching - # what the SQL built. DataFusion then refused the call while the native - # engine (which binds by name) still returned rows -- the two engines - # disagreeing about one fitted object. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train) # fitted order: [age, income] - t = SQLTransform(t"SELECT {sc}(income, age) AS out FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) # CALL order: [income, age] - batch = _both_engines(t, train) - got = np.array([[b["out"]["age"], b["out"]["income"]] for b in batch]) - assert np.allclose(got, sc.transform(train)) - - -def test_call_order_free_for_a_validating_transformer(): - # The fit-time probe must feed .transform in feature_names_in_ order, the - # way both engines feed it at runtime. Probing in call order ran the - # transform on feature-swapped data, so a transformer that VALIDATES its - # input (OrdinalEncoder with handle_unknown="error") rejected the probe and - # fit() raised -- for a query both engines execute correctly. - df = pd.DataFrame({"cat": ["a", "b", "a", "b"], "grp": ["x", "y", "x", "y"]}) - oe = OrdinalEncoder(handle_unknown="error").fit(df) # names: [cat, grp] - - t = SQLTransform(t"SELECT {oe}(grp, cat) AS o FROM __THIS__").fit( - pa.Table.from_pandas(df) - ) # CALL order reversed - batch = _both_engines(t, df) - expected = oe.transform(df) - got = np.array([[r["o"]["cat"], r["o"]["grp"]] for r in batch]) - assert np.allclose(got, expected) - - -def test_nested_outer_fitted_in_permuted_order_parity(): - # The outer's fitted order is a PERMUTATION of the inner's output names. The - # struct the outer actually receives is the inner's output, so in_schema must - # describe THAT, not the outer's fitted order. Declaring the fitted order - # desynced the UDF's Exact struct signature from what the SQL builds: - # fit() accepted, the DataFusion oracle refused to plan, and the native - # engine still returned rows -- the two engines disagreeing about one query. - # n_components=1 (not 2, like the fit is otherwise wide enough for): with - # this fixture's exactly-collinear age/income, a 2nd component would be a - # true zero-variance direction where the two engines' float noise disagrees - # in the low digits -- a real but unrelated sensitivity this test must not - # trip over. The permutation bug being tested lives in the INPUT struct - # order, unaffected by output width. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - inner = StandardScaler().fit(train) # emits struct - mid = pd.DataFrame(inner.transform(train), columns=inner.get_feature_names_out()) - outer = PCA(n_components=1).fit(mid[["income", "age"]]) # fitted [income, age] - - t = SQLTransform(t"SELECT {outer}({inner}(age, income)) AS o FROM __THIS__").fit( - pa.Table.from_pandas(train) - ) - batch = _both_engines(t, train) - expected = outer.transform(mid[["income", "age"]]) - out_names = [str(n) for n in outer.get_feature_names_out()] - got = np.array([[r["o"][n] for n in out_names] for r in batch]) - assert np.allclose(got, expected) - - -def test_nested_outer_column_name_mismatch_raises(): - # Outer fitted on DIFFERENT NAMES than the inner emits. Previously this died - # with a raw KeyError from inside _probe, mid-fit(). The nested branch now - # validates the name set like the leaf branch does. - train = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - inner = StandardScaler().fit(train) # emits age, income - mid = pd.DataFrame(inner.transform(train), columns=inner.get_feature_names_out()) - renamed = mid.rename(columns={"age": "c1", "income": "c2"}) - outer = PCA(n_components=2).fit(renamed) # feature_names_in_ = [c1, c2] - - t = SQLTransform(t"SELECT {outer}({inner}(age, income)) AS o FROM __THIS__") - with pytest.raises(ValueError, match="must match feature_names_in_"): - t.fit(pa.Table.from_pandas(train)) diff --git a/tests/test_transformer_udf.py b/tests/test_transformer_udf.py deleted file mode 100644 index f761d2b..0000000 --- a/tests/test_transformer_udf.py +++ /dev/null @@ -1,85 +0,0 @@ -"""Standalone oracle test: the DataFusion transformer UDF must match sklearn.""" - -import numpy as np -import pandas as pd -import pyarrow as pa -import pytest -from datafusion import SessionContext -from sklearn.preprocessing import StandardScaler - -from sql_transform._transformer_udf import _transformer_udf - -# Both engines deliberately feed sklearn positionally-aligned nameless arrays -# (we reorder to feature_names_in_ ourselves), so sklearn's redundant -# "X does not have valid feature names" warning is a known false positive here. -pytestmark = pytest.mark.filterwarnings( - "ignore:X does not have valid feature names:UserWarning" -) - - -def _collect(df) -> list[dict]: - return pa.Table.from_batches(df.collect(), schema=df.schema()).to_pylist() - - -def test_out_schema_mismatching_natural_dtype_raises(): - # TASK-2 AC#3. The two engines reach the declared out_schema by DIFFERENT - # coercion rules -- DataFusion casts via pa.array(..., type=), the native - # engine coerces at pydantic model-validate. They only agree when no - # coercion happens, i.e. when the declared type already IS the natural - # .transform dtype. Declaring int64 over a float64 transform silently - # truncates on one side; make it a build-time error instead. - train_df = pd.DataFrame({"age": [10.0, 20.0, 30.0, 40.0]}) - sc = StandardScaler().fit(train_df) - in_schema = pa.schema([("age", pa.float64())]) - bad_out = pa.schema([("age", pa.int64())]) # natural is float64 - - ctx = SessionContext() - ctx.from_arrow(pa.Table.from_pandas(train_df), name="__THIS__") - ctx.register_udf(_transformer_udf(sc, in_schema, bad_out, "__tfm_0__")) - - with pytest.raises(Exception, match="natural"): - _collect( - ctx.sql("SELECT __tfm_0__(named_struct('age', age)) AS s FROM __THIS__") - ) - - -def test_out_schema_matching_natural_dtype_is_accepted(): - # The guard must not fire on the legitimate case: declared == natural. - train_df = pd.DataFrame({"age": [10.0, 20.0, 30.0, 40.0]}) - sc = StandardScaler().fit(train_df) - schema = pa.schema([("age", pa.float64())]) - - ctx = SessionContext() - ctx.from_arrow(pa.Table.from_pandas(train_df), name="__THIS__") - ctx.register_udf(_transformer_udf(sc, schema, schema, "__tfm_0__")) - - got = _collect( - ctx.sql("SELECT __tfm_0__(named_struct('age', age)) AS s FROM __THIS__") - ) - assert np.allclose([r["s"]["age"] for r in got], sc.transform(train_df).ravel()) - - -def test_standard_scaler_udf_matches_sklearn(): - train_df = pd.DataFrame( - {"age": [10.0, 20.0, 30.0, 40.0], "income": [1.0, 2.0, 3.0, 4.0]} - ) - sc = StandardScaler().fit(train_df) # fit on NAMED data -> feature_names_in_ - in_schema = pa.schema([("age", pa.float64()), ("income", pa.float64())]) - out_schema = pa.schema([("age", pa.float64()), ("income", pa.float64())]) - - table = pa.Table.from_pandas(train_df) - ctx = SessionContext() - ctx.from_arrow(table, name="__THIS__") - ctx.register_udf(_transformer_udf(sc, in_schema, out_schema, "__tfm_0__")) - - q = ( - "SELECT __tfm_0__(named_struct('age', age, 'income', income)) " - "AS s FROM __THIS__" - ) - got = _collect(ctx.sql(q)) - - expected = sc.transform(train_df) - assert len(got) == 4 - for i, r in enumerate(got): - assert abs(r["s"]["age"] - expected[i][0]) < 1e-9 - assert abs(r["s"]["income"] - expected[i][1]) < 1e-9