|
| 1 | +"""The row-shape contract flag (TASK-58): shape="map" | "filter" | "many". |
| 2 | +
|
| 3 | +"map" is a STATIC build-time proof of exactly one output row per input row |
| 4 | +(out[i] <-> in[i]) — the serving-path guarantee. "filter" is the default |
| 5 | +0..1 behavior, byte-identical to before the flag existed. "many" is |
| 6 | +reserved for join multiplicity (stage B) and is the only shape under which |
| 7 | +those constructs will ever build. |
| 8 | +""" |
| 9 | + |
| 10 | +from __future__ import annotations |
| 11 | + |
| 12 | +import pyarrow as pa |
| 13 | +import pytest |
| 14 | +from pydantic import create_model |
| 15 | + |
| 16 | +from sql_transform._interpreter import DuckDBInferFn |
| 17 | + |
| 18 | +T = create_model("T", a=(int, ...), s=(str | None, None)) |
| 19 | +DIM = pa.table({"id": [1, 2], "v": [10, 20]}) |
| 20 | + |
| 21 | + |
| 22 | +def build(sql, shape=None, statics=None): |
| 23 | + kwargs = {"output": "dict"} |
| 24 | + if shape is not None: |
| 25 | + kwargs["shape"] = shape |
| 26 | + return DuckDBInferFn( |
| 27 | + sql, row_tables={"__THIS__": T}, static_tables=statics or {}, **kwargs |
| 28 | + ) |
| 29 | + |
| 30 | + |
| 31 | +def test_map_serves_projections_and_left_joins(): |
| 32 | + fn = build("SELECT a + 1 AS b, upper(s) AS u FROM __THIS__", shape="map") |
| 33 | + assert fn.shape == "map" |
| 34 | + got = fn.infer({"__THIS__": [T(a=1, s="x"), T(a=2, s=None)]}) |
| 35 | + assert [r["b"] for r in got] == [2, 3] # exactly one out per in, in order |
| 36 | + |
| 37 | + fn = build( |
| 38 | + "SELECT a, v FROM __THIS__ LEFT JOIN d ON a = d.id", |
| 39 | + shape="map", |
| 40 | + statics={"d": DIM}, |
| 41 | + ) |
| 42 | + got = fn.infer({"__THIS__": [T(a=1), T(a=99)]}) |
| 43 | + assert [r["v"] for r in got] == [10, None] # a miss maps, never drops |
| 44 | + |
| 45 | + |
| 46 | +def test_map_rejects_row_dropping_constructs(): |
| 47 | + with pytest.raises(ValueError, match="shape='map'.*WHERE"): |
| 48 | + build("SELECT a FROM __THIS__ WHERE a > 0", shape="map") |
| 49 | + with pytest.raises(ValueError, match="shape='map'.*INNER JOIN 'd'"): |
| 50 | + build( |
| 51 | + "SELECT a, v FROM __THIS__ JOIN d ON a = d.id", |
| 52 | + shape="map", |
| 53 | + statics={"d": DIM}, |
| 54 | + ) |
| 55 | + # A static-only query emits fixed rows unrelated to the input. |
| 56 | + with pytest.raises(ValueError, match="shape='map'.*static-tables-only"): |
| 57 | + build("SELECT max(id) FROM d", shape="map", statics={"d": DIM}) |
| 58 | + |
| 59 | + |
| 60 | +def test_filter_default_unchanged(): |
| 61 | + for kwargs in [{}, {"shape": "filter"}]: |
| 62 | + fn = build("SELECT a FROM __THIS__ WHERE a > 1", **kwargs) |
| 63 | + assert fn.shape == "filter" |
| 64 | + got = fn.infer({"__THIS__": [T(a=1), T(a=2)]}) |
| 65 | + assert [r["a"] for r in got] == [2] |
| 66 | + |
| 67 | + |
| 68 | +def test_many_is_reserved_and_bad_values_are_named(): |
| 69 | + with pytest.raises(ValueError, match="stage B"): |
| 70 | + build("SELECT a FROM __THIS__", shape="many") |
| 71 | + with pytest.raises(ValueError, match="must be 'map', 'filter', or 'many'"): |
| 72 | + build("SELECT a FROM __THIS__", shape="projection") |
0 commit comments