From ca29f7aa47fc7237b210495298d6990f1513aded Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philipp=20Gro=C3=9Fer?= Date: Thu, 27 Aug 2026 18:18:58 +0200 Subject: [PATCH 1/2] added functionality --- docs/api/plan.md | 91 ++++++++++++++ docs/roadmap.md | 106 ++++++++++++++++ mkdocs.yml | 1 + tests/test_column.py | 80 ++++++++++++ tests/test_core.py | 168 +++++++++++++++++++++++++ tests/test_duckdb.py | 47 +++++++ tests/test_map.py | 23 ++++ tests/test_math.py | 23 ++++ tests/test_protocol.py | 93 ++++++++++++++ tests/test_string.py | 24 ++++ transformplan/backends/base.py | 16 +++ transformplan/backends/duckdb.py | 27 +--- transformplan/backends/polars.py | 10 +- transformplan/core.py | 209 +++++++++++++++++++++++++++++-- transformplan/dtypes.py | 177 ++++++++++++++++++++++++++ transformplan/ops/_common.py | 25 ++++ transformplan/ops/column.py | 61 +++++++-- transformplan/ops/map.py | 50 +++++--- transformplan/ops/math.py | 76 +++++++---- transformplan/ops/string.py | 90 ++++++++----- transformplan/protocol.py | 71 ++++++++++- transformplan/validation.py | 15 ++- 22 files changed, 1368 insertions(+), 115 deletions(-) create mode 100644 docs/roadmap.md create mode 100644 transformplan/dtypes.py create mode 100644 transformplan/ops/_common.py diff --git a/docs/api/plan.md b/docs/api/plan.md index 6aa713f..8edba90 100644 --- a/docs/api/plan.md +++ b/docs/api/plan.md @@ -50,12 +50,103 @@ See [Backends](backends.md) for details on each backend. - validate - validate_chunked - dry_run + - pipe + - extend + - section + - because - to_dict - from_dict - to_json - from_json - to_python +## Structuring Long Plans + +Plans with dozens of steps become hard to read as a flat chain. Four methods +give them structure without changing what they do. + +### pipe + +Apply a function that takes and returns a plan, so repeated blocks can live in +their own function without breaking the chain. + +```python +def decimal_hour(plan, source, target): + return ( + plan.dt_format(source, "%H", target) + .col_cast(target, "Float64") + ) + +plan = ( + TransformPlan() + .dt_diff_days("admitted", "discharged", "stay_days") + .pipe(decimal_hour, "admitted_at", "admit_hour") + .pipe(decimal_hour, "discharged_at", "discharge_hour") +) +``` + +`pipe` registers no step of its own — the function registers ordinary steps, so +the protocol is unchanged. + +### extend + +Append another plan's steps, so a reusable block can be defined once as its own +plan. Steps are deep-copied, so a block can be reused across plans without them +sharing state. + +```python +CLEAN_NAMES = TransformPlan().str_strip("name").str_upper("name") + +plan = TransformPlan().col_drop("temp").extend(CLEAN_NAMES) +combined = plan_a + plan_b # same thing, as a new plan +``` + +Unlike `pipe`, the block's steps become part of the plan and are serialized with it. + +### section + +Group the following steps under a label. The protocol then reports a balance per +section instead of one flat list. + +```python +plan = ( + TransformPlan() + .section("Row filters") + .rows_drop_nulls("shipped_at") + .rows_drop_nulls("order_id") + .section("Split product code") + .col_duplicate("code", "code_group") + .str_slice("code_group", 0, 3) +) +``` + +``` +SECTIONS + + Row filters 1500 → 49 rows (2 steps) + Split product code 49 → 49 rows (2 steps, +1 cols) +``` + +Pass `None` to end a section. Plans without sections render exactly as before. + +### because + +Attach a reason to the step just registered. The protocol records that rows were +removed; the reason records why — the question an audit actually asks. + +```python +plan = ( + TransformPlan() + .rows_drop(Col("provider") != "MDK02") + .because("cases where the reviewer awards the insurer nothing") + .rows_drop_nulls("drg_code") + .because("PEPP cases carry no DRG") +) +``` + +When the preceding call covered several columns, the reason is attached to every +step it produced. + ## Execution Methods ### process diff --git a/docs/roadmap.md b/docs/roadmap.md new file mode 100644 index 0000000..bbeace3 --- /dev/null +++ b/docs/roadmap.md @@ -0,0 +1,106 @@ +# Roadmap + +Planned additions that are deliberately **not** in the current release, with the +reasoning and the groundwork already in place. + +Both items below add new `@abstractmethod` entries to the `Backend` ABC +(`transformplan/backends/base.py`). That breaks any backend implemented outside +this repository, so they belong in a minor release (0.3.0), not a patch. + +--- + +## Aggregation — `group_by` with aggregates + +**Status:** planned, largest known gap. + +There is `rows_pivot` and `rows_melt`, but no way to group with aggregation. Any +pipeline whose job is to condense onto an entity — one row per case, customer, or +transaction — has to leave the plan at exactly that point. The consequence is that +the most complex part of such a pipeline is the one part with **no protocol**, +while the simpler row and column operations are documented in full. + +### Proposed API + +```python +.group_by("case_id", { + "n_positions": ("*", "count"), + "total_amount": ("amount", "sum"), + "n_categories": ("field", "n_unique"), +}) +``` + +### Groundwork already present + +Two pieces exist and should be reused rather than rebuilt: + +- **`AggFunction`** (`backends/base.py`) already defines the aggregate vocabulary + `Literal["first", "sum", "mean", "median", "min", "max", "count"]`, used today by + `rows_pivot` and `math_diff_from_agg`. `group_by` should extend this literal + (`n_unique`, `last`, `std`) rather than introduce a second vocabulary. +- **`ChunkMode.GROUP_DEPENDENT` with `group_param`** (`chunking.py`) already models + "needs all rows of a group together", used by `math_cumsum`, `math_rank` and + `rows_unique`. Registering `group_by` as + `OperationMeta(ChunkMode.GROUP_DEPENDENT, group_param="by")` makes it correct + under `process_chunked()` whenever `partition_key == by` — so aggregation fits + *inside* the chunking model instead of standing beside it. + +### Where the real cost is + +Not in the backends. The expensive part is `SchemaTracker` in `validation.py`: +every existing validator *advances* the schema, whereas `group_by` must **replace** +it — after the step, only the grouping columns and the aggregate outputs exist. +This is the one place where the current forward-propagation logic does not carry +over. Beyond that: DuckDB SQL generation and `dry_run` output. + +### Constraint to hold + +Keep the aggregate vocabulary **closed** — names only, never arbitrary callables. +A lambda cannot be serialized, and serializability is the library's core promise. + +--- + +## Time components — `dt_hour`, `dt_minute`, `dt_second` + +**Status:** planned, small. + +The `dt_` family covers date components thoroughly but has no time components. +Turning `16:30` into `16.5` currently takes six steps and a detour through text +formatting and back: + +```python +.dt_format("admitted_at", "%H", "admit_h") +.dt_format("admitted_at", "%M", "admit_m") +.col_cast("admit_h", "Float64") +.col_cast("admit_m", "Float64") +.math_divide("admit_m", 60) +.math_add_columns("admit_h", "admit_m", "admit_hour") +``` + +With `dt_hour` and `dt_minute` this becomes three steps and no string round trip. + +### Notes for implementation + +- DuckDB is trivial here: `EXTRACT(hour FROM …)`. +- Each new operation touches **nine files**: `ops/datetime.py`, `backends/base.py`, + `backends/polars.py`, `backends/duckdb.py`, `validation.py` (validator plus + registry), `chunking.py` (registry), `tests/test_datetime.py`, + `tests/test_duckdb.py`, and `docs/api/ops/datetime.md`. Budget for that, not + just for the operation itself. +- **`dt_decimal_hour` is deliberately excluded.** With `dt_hour`/`dt_minute` the + block is already down to three steps, and the library has no other convenience + combinations. Adding one opens a door to an unbounded set of special cases. + +--- + +## Delivered in this release + +For reference, the related suggestions that **are** implemented: + +| Item | What shipped | +|---|---| +| `.pipe()` | Apply a function that takes and returns a plan | +| `.extend()` / `+` | Compose plans from reusable blocks | +| `.section()` | Named sections, grouped with a balance in the protocol | +| `.because()` | Reasons recorded per step in the protocol | +| Column sequences | 24 in-place operations accept `str \| Sequence[str]` | +| `col_cast` names | Canonical dtype names keep plans serializable | diff --git a/mkdocs.yml b/mkdocs.yml index fbec8e7..02f04f9 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -74,6 +74,7 @@ markdown_extensions: nav: - Home: index.md + - Roadmap: roadmap.md - Getting Started: - Installation: getting-started/installation.md - Quickstart: getting-started/quickstart.md diff --git a/tests/test_column.py b/tests/test_column.py index f1f3aa6..c3b6638 100644 --- a/tests/test_column.py +++ b/tests/test_column.py @@ -1,6 +1,7 @@ """Tests for column operations (ops/column.py).""" import polars as pl +import pytest from transformplan import TransformPlan @@ -61,9 +62,88 @@ def test_col_rename_to_existing_raises(self, basic_df: pl.DataFrame) -> None: assert "already exists" in str(result.errors[0]) +class TestColumnSequences: + """Tests for column operations accepting a sequence of columns.""" + + def test_col_drop_sequence(self, basic_df: pl.DataFrame) -> None: + """Test dropping several columns in one call.""" + plan = TransformPlan().col_drop(["age", "salary"]) + result, _ = plan.process(basic_df) + assert "age" not in result.columns + assert "salary" not in result.columns + assert len(plan) == 2 + + def test_col_drop_single_unchanged(self, basic_df: pl.DataFrame) -> None: + """Test that a single column name still registers one step.""" + plan = TransformPlan().col_drop("age") + result, _ = plan.process(basic_df) + assert "age" not in result.columns + assert len(plan) == 1 + + def test_col_fill_null_sequence(self, df_with_nulls: pl.DataFrame) -> None: + """Test filling nulls in several columns in one call.""" + cols = [c for c in ("name", "city") if c in df_with_nulls.columns] + plan = TransformPlan().col_fill_null(cols, value="N/A") + result, _ = plan.process(df_with_nulls) + for col in cols: + assert result[col].null_count() == 0 + + def test_sequence_steps_are_individually_protocolled( + self, basic_df: pl.DataFrame + ) -> None: + """Test that each column of a sequence becomes its own protocol step.""" + plan = TransformPlan().col_drop(["age", "salary"]) + _, protocol = plan.process(basic_df) + assert len(protocol) == 2 + steps = protocol.to_dict()["steps"] + assert [s["params"]["column"] for s in steps] == ["age", "salary"] + + class TestColCast: """Tests for col_cast operation.""" + def test_col_cast_canonical_name(self, basic_df: pl.DataFrame) -> None: + """Test casting with a canonical dtype name.""" + plan = TransformPlan().col_cast("age", "Float64") + result, _ = plan.process(basic_df) + assert result["age"].dtype == pl.Float64 + + def test_col_cast_alias_name(self, basic_df: pl.DataFrame) -> None: + """Test that Utf8 is accepted as an alias for String.""" + plan = TransformPlan().col_cast("id", "Utf8") + result, _ = plan.process(basic_df) + assert result["id"].dtype == pl.String + + def test_col_cast_python_type(self, basic_df: pl.DataFrame) -> None: + """Test casting with a builtin Python type.""" + plan = TransformPlan().col_cast("age", float) + result, _ = plan.process(basic_df) + assert result["age"].dtype == pl.Float64 + + def test_col_cast_normalizes_to_name(self) -> None: + """Test that a Polars dtype is stored as its canonical name.""" + plan = TransformPlan().col_cast("age", pl.Float64) + assert plan.to_dict()["steps"][0]["params"]["dtype"] == "Float64" + + def test_col_cast_unknown_name_raises_on_build(self) -> None: + """Test that a bad dtype name fails when the plan is built.""" + with pytest.raises(ValueError, match="Unknown dtype name"): + TransformPlan().col_cast("age", "Floot64") + + def test_col_cast_exotic_dtype_passes_through(self, basic_df: pl.DataFrame) -> None: + """Test that a parameterized dtype keeps working unchanged.""" + plan = TransformPlan().col_cast("name", pl.Categorical) + result, _ = plan.process(basic_df) + assert result["name"].dtype == pl.Categorical + + def test_col_cast_sequence(self, basic_df: pl.DataFrame) -> None: + """Test casting several columns in one call.""" + plan = TransformPlan().col_cast(["age", "id"], "Float64") + result, _ = plan.process(basic_df) + assert result["age"].dtype == pl.Float64 + assert result["id"].dtype == pl.Float64 + assert len(plan) == 2 + def test_col_cast_int_to_float(self, basic_df: pl.DataFrame) -> None: """Test casting int to float.""" plan = TransformPlan().col_cast("age", pl.Float64) diff --git a/tests/test_core.py b/tests/test_core.py index ae2800d..7b339de 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -216,6 +216,159 @@ def test_roundtrip_serialization(self, basic_df: pl.DataFrame) -> None: assert original_result.equals(restored_result) +class TestPipe: + """Tests for pipe() composition.""" + + def test_pipe_keeps_the_chain(self, basic_df: pl.DataFrame) -> None: + """Test that pipe applies a function without breaking the chain.""" + + def rename_block( + plan: TransformPlan, source: str, target: str + ) -> TransformPlan: + return plan.col_rename(source, target) + + plan = ( + TransformPlan() + .col_drop("age") + .pipe(rename_block, "name", "full_name") + .col_drop("salary") + ) + result, _ = plan.process(basic_df) + assert "full_name" in result.columns + assert len(plan) == 3 + + def test_pipe_forwards_keyword_arguments(self, basic_df: pl.DataFrame) -> None: + """Test that keyword arguments reach the piped function.""" + + def drop_it(plan: TransformPlan, *, column: str) -> TransformPlan: + return plan.col_drop(column) + + plan = TransformPlan().pipe(drop_it, column="age") + result, _ = plan.process(basic_df) + assert "age" not in result.columns + + def test_pipe_leaves_protocol_unchanged(self, basic_df: pl.DataFrame) -> None: + """Test that pipe itself registers no step.""" + + def noop(plan: TransformPlan) -> TransformPlan: + return plan + + plan = TransformPlan().col_drop("age").pipe(noop) + _, protocol = plan.process(basic_df) + assert len(protocol) == 1 + + +class TestExtend: + """Tests for extend() and plan addition.""" + + def test_extend_appends_steps(self, basic_df: pl.DataFrame) -> None: + """Test that extend appends another plan's steps.""" + block = TransformPlan().col_drop("salary") + plan = TransformPlan().col_drop("age").extend(block) + result, _ = plan.process(basic_df) + assert "age" not in result.columns + assert "salary" not in result.columns + assert len(plan) == 2 + + def test_extend_deep_copies_params(self) -> None: + """Test that a reused block does not share mutable state.""" + block = TransformPlan().col_fill_null("name", value="X") + + first = TransformPlan().extend(block) + second = TransformPlan().extend(block) + + first._operations[0][1]["value"] = "CHANGED" + + assert second._operations[0][1]["value"] == "X" + assert block._operations[0][1]["value"] == "X" + + def test_extend_is_serializable(self, basic_df: pl.DataFrame) -> None: + """Test that extended steps survive a JSON round trip.""" + block = TransformPlan().col_drop("salary") + plan = TransformPlan().col_drop("age").extend(block) + restored = TransformPlan.from_json(plan.to_json()) + assert len(restored) == 2 + + expected, _ = plan.process(basic_df) + actual, _ = restored.process(basic_df) + assert actual.equals(expected) + + def test_add_leaves_operands_untouched(self) -> None: + """Test that plan_a + plan_b builds a new plan.""" + left = TransformPlan().col_drop("age") + right = TransformPlan().col_drop("salary") + combined = left + right + + assert len(left) == 1 + assert len(right) == 1 + assert len(combined) == 2 + + def test_extend_inherits_target_section(self) -> None: + """Test that a block without sections adopts the target's section.""" + block = TransformPlan().col_drop("salary") + plan = TransformPlan().section("Cleanup").extend(block) + assert plan.to_dict()["steps"][0]["section"] == "Cleanup" + + def test_extend_keeps_block_section(self) -> None: + """Test that a block's own section wins over the target's.""" + block = TransformPlan().section("Block").col_drop("salary") + plan = TransformPlan().section("Target").extend(block) + assert plan.to_dict()["steps"][0]["section"] == "Block" + + +class TestSectionAndReason: + """Tests for section() and because() metadata.""" + + def test_section_labels_following_steps(self) -> None: + """Test that steps inherit the current section.""" + plan = ( + TransformPlan() + .section("Filters") + .col_drop("age") + .col_drop("salary") + .section("Types") + .col_cast("id", "Float64") + ) + sections = [s.get("section") for s in plan.to_dict()["steps"]] + assert sections == ["Filters", "Filters", "Types"] + + def test_section_none_ends_the_section(self) -> None: + """Test that section(None) stops labelling.""" + plan = TransformPlan().section("A").col_drop("age").section(None).col_drop("id") + steps = plan.to_dict()["steps"] + assert steps[0]["section"] == "A" + assert "section" not in steps[1] + + def test_because_annotates_the_last_step(self) -> None: + """Test that because attaches a reason.""" + plan = TransformPlan().col_drop("age").because("not needed downstream") + assert plan.to_dict()["steps"][0]["reason"] == "not needed downstream" + + def test_because_annotates_every_step_of_a_sequence(self) -> None: + """Test that a sequence call gets the reason on all its steps.""" + plan = TransformPlan().col_drop(["age", "salary"]).because("unused") + reasons = [s.get("reason") for s in plan.to_dict()["steps"]] + assert reasons == ["unused", "unused"] + + def test_because_without_operation_raises(self) -> None: + """Test that because before any operation is an error.""" + with pytest.raises(ValueError, match="must follow an operation"): + TransformPlan().because("nothing to explain") + + def test_metadata_survives_round_trip(self) -> None: + """Test that section and reason survive a JSON round trip.""" + plan = TransformPlan().section("Cleanup").col_drop("age").because("not needed") + restored = TransformPlan.from_json(plan.to_json()) + step = restored.to_dict()["steps"][0] + assert step["section"] == "Cleanup" + assert step["reason"] == "not needed" + + def test_plain_plan_serializes_without_metadata_keys(self) -> None: + """Test that plans without metadata are unchanged in JSON.""" + step = TransformPlan().col_drop("age").to_dict()["steps"][0] + assert set(step) == {"operation", "params"} + + class TestToPython: """Tests for to_python() method.""" @@ -233,6 +386,21 @@ def test_to_python_custom_variable(self) -> None: code = plan.to_python(variable_name="my_plan") assert "my_plan = (" in code + def test_to_python_is_executable_with_cast(self) -> None: + """Test that generated code compiles when a dtype is involved.""" + plan = TransformPlan().col_cast("age", pl.Float64) + code = plan.to_python() + assert 'dtype="Float64"' in code + compile(code, "", "exec") + + def test_to_python_renders_section_and_reason(self) -> None: + """Test that metadata is reproduced in generated code.""" + plan = TransformPlan().section("Cleanup").col_drop("age").because("not needed") + code = plan.to_python() + assert '.section("Cleanup")' in code + assert '.because("not needed")' in code + compile(code, "", "exec") + def test_to_python_with_filter(self, basic_df: pl.DataFrame) -> None: """Test Python code generation with filter.""" plan = TransformPlan().rows_filter(Col("age") >= 30) diff --git a/tests/test_duckdb.py b/tests/test_duckdb.py index 90b5b1c..6d41607 100644 --- a/tests/test_duckdb.py +++ b/tests/test_duckdb.py @@ -10,6 +10,7 @@ from typing import Any import duckdb +import polars as pl import pytest from transformplan import Col, TransformPlan @@ -213,6 +214,52 @@ def test_col_cast_int_to_string( assert vals[0] == "1" +class TestColCastDtypeParity: + """Casting must produce the same logical type across backends.""" + + def test_polars_dtype_maps_to_double( + self, backend: DuckDBBackend, basic_rel: duckdb.DuckDBPyRelation + ) -> None: + """A Polars dtype must not silently degrade to VARCHAR.""" + result, _ = ( + TransformPlan() + .col_cast("age", pl.Float64) + .process(basic_rel, backend=backend) + ) + assert result.types[result.columns.index("age")] == "DOUBLE" + + def test_canonical_name_maps_to_double( + self, backend: DuckDBBackend, basic_rel: duckdb.DuckDBPyRelation + ) -> None: + """A canonical dtype name resolves to the DuckDB type.""" + result, _ = ( + TransformPlan() + .col_cast("age", "Float64") + .process(basic_rel, backend=backend) + ) + assert result.types[result.columns.index("age")] == "DOUBLE" + + def test_int64_maps_to_bigint( + self, backend: DuckDBBackend, basic_rel: duckdb.DuckDBPyRelation + ) -> None: + """Integer widths are mapped explicitly, not by fallback.""" + result, _ = ( + TransformPlan().col_cast("age", "Int64").process(basic_rel, backend=backend) + ) + assert result.types[result.columns.index("age")] == "BIGINT" + + def test_both_backends_agree( + self, backend: DuckDBBackend, basic_rel: duckdb.DuckDBPyRelation + ) -> None: + """The same plan yields the same values on both backends.""" + plan = TransformPlan().col_cast("age", "Float64") + + duck_result, _ = plan.process(basic_rel, backend=backend) + polars_result, _ = plan.process(basic_rel.pl()) + + assert _col_values(duck_result, "age") == polars_result["age"].to_list() + + class TestColReorder: def test_col_reorder( self, backend: DuckDBBackend, basic_rel: duckdb.DuckDBPyRelation diff --git a/tests/test_map.py b/tests/test_map.py index ebe57bc..98c9d44 100644 --- a/tests/test_map.py +++ b/tests/test_map.py @@ -399,3 +399,26 @@ def test_map_discretize_right_false_multiple_bins(self) -> None: assert result["grade"][2] == "C" assert result["grade"][3] == "B" assert result["grade"][4] == "A" + + +class TestMapSequences: + """Tests for map operations accepting a sequence of columns.""" + + def test_map_values_sequence(self) -> None: + """Test applying one mapping to several columns.""" + df = pl.DataFrame({"a": ["x", "y"], "b": ["y", "x"]}) + plan = TransformPlan().map_values(["a", "b"], {"x": "X", "y": "Y"}) + result, _ = plan.process(df) + + assert result["a"].to_list() == ["X", "Y"] + assert result["b"].to_list() == ["Y", "X"] + assert len(plan) == 2 + + def test_map_null_to_value_sequence(self) -> None: + """Test replacing nulls in several columns at once.""" + df = pl.DataFrame({"a": ["x", None], "b": [None, "y"]}) + plan = TransformPlan().map_null_to_value(["a", "b"], "MISSING") + result, _ = plan.process(df) + + assert result["a"].to_list() == ["x", "MISSING"] + assert result["b"].to_list() == ["MISSING", "y"] diff --git a/tests/test_math.py b/tests/test_math.py index ab94263..14de7e3 100644 --- a/tests/test_math.py +++ b/tests/test_math.py @@ -629,3 +629,26 @@ def test_serialization_roundtrip(self, numeric_df: pl.DataFrame) -> None: result1, _ = plan.process(numeric_df) result2, _ = restored.process(numeric_df) assert result1["diff"].to_list() == result2["diff"].to_list() + + +class TestMathSequences: + """Tests for math operations accepting a sequence of columns.""" + + def test_math_multiply_sequence(self, numeric_df: pl.DataFrame) -> None: + """Test multiplying several columns by the same factor.""" + cols = [c for c in numeric_df.columns if numeric_df[c].dtype.is_numeric()][:2] + expected = [[v * 2 for v in numeric_df[c].to_list()] for c in cols] + + plan = TransformPlan().math_multiply(cols, 2) + result, _ = plan.process(numeric_df) + + for col, exp in zip(cols, expected, strict=True): + assert result[col].to_list() == exp + assert len(plan) == 2 + + def test_math_round_sequence(self, numeric_df: pl.DataFrame) -> None: + """Test rounding several columns at once.""" + cols = [c for c in numeric_df.columns if numeric_df[c].dtype.is_numeric()][:2] + plan = TransformPlan().math_round(cols, 1) + result, _ = plan.process(numeric_df) + assert result.height == numeric_df.height diff --git a/tests/test_protocol.py b/tests/test_protocol.py index c7d18c5..656b546 100644 --- a/tests/test_protocol.py +++ b/tests/test_protocol.py @@ -268,6 +268,99 @@ def test_print( assert "TRANSFORM PROTOCOL" in captured.out +class TestProtocolSectionsAndReasons: + """Tests for section grouping and reasons in the protocol.""" + + def test_summary_shows_section_balance(self, basic_df: pl.DataFrame) -> None: + """Test that a per-section balance is rendered.""" + plan = ( + TransformPlan() + .section("Cleanup") + .col_drop("age") + .col_drop("salary") + .section("Types") + .col_cast("id", "Float64") + ) + _, protocol = plan.process(basic_df) + summary = protocol.summary() + + assert "SECTIONS" in summary + assert "Cleanup" in summary + assert "Types" in summary + assert "2 steps" in summary + + def test_summary_shows_reason(self, basic_df: pl.DataFrame) -> None: + """Test that a reason appears under its step.""" + plan = TransformPlan().col_drop("age").because("not needed downstream") + _, protocol = plan.process(basic_df) + assert "not needed downstream" in protocol.summary() + + def test_summary_unchanged_without_sections(self, basic_df: pl.DataFrame) -> None: + """Test that plans without sections render as before.""" + plan = TransformPlan().col_drop("age") + _, protocol = plan.process(basic_df) + assert "SECTIONS" not in protocol.summary() + + def test_dataframe_has_metadata_columns(self, basic_df: pl.DataFrame) -> None: + """Test that section and reason reach the DataFrame export.""" + plan = TransformPlan().section("A").col_drop("age").because("why") + _, protocol = plan.process(basic_df) + df = protocol.to_dataframe() + + assert "section" in df.columns + assert "reason" in df.columns + assert df.filter(pl.col("step") == 1)["section"][0] == "A" + assert df.filter(pl.col("step") == 1)["reason"][0] == "why" + + def test_csv_has_metadata_columns(self, basic_df: pl.DataFrame) -> None: + """Test that section and reason reach the CSV export.""" + plan = TransformPlan().section("A").col_drop("age").because("why") + _, protocol = plan.process(basic_df) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "protocol.csv" + protocol.to_csv(path) + content = path.read_text() + + assert "section" in content + assert "why" in content + + def test_dict_round_trip_keeps_metadata(self, basic_df: pl.DataFrame) -> None: + """Test that metadata survives a protocol round trip.""" + plan = TransformPlan().section("A").col_drop("age").because("why") + _, protocol = plan.process(basic_df) + + restored = Protocol.from_dict(protocol.to_dict()) + step = restored.to_dict()["steps"][0] + assert step["section"] == "A" + assert step["reason"] == "why" + + def test_from_dict_tolerates_missing_metadata(self) -> None: + """Test that protocols written before metadata existed still load.""" + legacy = { + "version": "1.0", + "created_at": "2024-01-01T00:00:00", + "metadata": {}, + "input": {"hash": "abc123", "shape": [5, 3]}, + "steps": [ + { + "step": 1, + "operation": "col_drop", + "params": {"column": "age"}, + "old_shape": [5, 3], + "new_shape": [5, 2], + "rows_changed": 0, + "cols_changed": 1, + "elapsed_seconds": 0.001, + "output_hash": "def456", + } + ], + } + restored = Protocol.from_dict(legacy) + assert len(restored) == 1 + assert restored.to_dict()["steps"][0]["section"] is None + + class TestProtocolStepTracking: """Tests for Protocol step tracking.""" diff --git a/tests/test_string.py b/tests/test_string.py index abd30eb..094cbfb 100644 --- a/tests/test_string.py +++ b/tests/test_string.py @@ -314,3 +314,27 @@ def test_unicode_string(self) -> None: result, _ = plan.process(df) assert result["text"][0] == "HÉLLO" assert result["text"][1] == "WÖRLD" + + +class TestStringSequences: + """Tests for string operations accepting a sequence of columns.""" + + def test_str_slice_sequence(self, string_df: pl.DataFrame) -> None: + """Test slicing two columns with identical parameters.""" + plan = TransformPlan().str_slice(["code", "first_name"], 0, 3) + result, _ = plan.process(string_df) + assert result["code"][0] == "PRD" + assert result["first_name"][0] == "Joh" + assert len(plan) == 2 + + def test_str_upper_sequence(self, string_df: pl.DataFrame) -> None: + """Test uppercasing several columns at once.""" + plan = TransformPlan().str_upper(["first_name", "last_name"]) + result, _ = plan.process(string_df) + assert result["first_name"][0] == "JOHN" + assert result["last_name"][0] == "DOE" + + def test_str_single_column_unchanged(self, string_df: pl.DataFrame) -> None: + """Test that a single name still registers exactly one step.""" + plan = TransformPlan().str_upper("first_name") + assert len(plan) == 1 diff --git a/transformplan/backends/base.py b/transformplan/backends/base.py index 0a6eff0..0ee2f5a 100644 --- a/transformplan/backends/base.py +++ b/transformplan/backends/base.py @@ -126,6 +126,22 @@ def type_name(self, dtype: Any) -> str: """Return a human-readable name for a dtype.""" ... + def resolve_dtype(self, dtype: Any) -> Any: + """Resolve a canonical dtype name to this backend's type. + + Plans store dtypes as canonical name strings so they stay serializable + (see ``transformplan.dtypes``). Backends override this to map a name to + their own type system. The default passes the value through unchanged, + so backends predating this method keep working. + + Args: + dtype: Canonical dtype name, or a native dtype to pass through. + + Returns: + The backend's dtype. + """ + return dtype + # ========================================================================= # Column operations (14) # ========================================================================= diff --git a/transformplan/backends/duckdb.py b/transformplan/backends/duckdb.py index 9c289d0..4081a6b 100644 --- a/transformplan/backends/duckdb.py +++ b/transformplan/backends/duckdb.py @@ -11,7 +11,6 @@ from __future__ import annotations -import datetime import hashlib import math import secrets @@ -29,6 +28,7 @@ Numeric, RankMethod, ) +from transformplan.dtypes import to_duckdb_type from transformplan.filters import Filter from transformplan.sql_utils import sql_format_value as _v from transformplan.sql_utils import sql_quote_identifier as _q @@ -37,27 +37,9 @@ # SQL helpers # ============================================================================= -_DTYPE_TO_DUCKDB: dict[type, str] = { - int: "BIGINT", - float: "DOUBLE", - str: "VARCHAR", - bool: "BOOLEAN", - datetime.datetime: "TIMESTAMP", - datetime.date: "DATE", -} - _VALID_AGG_FUNCTIONS = {"first", "sum", "mean", "median", "min", "max", "count"} -def _dtype_to_duckdb(dtype: type) -> str: - """Map Python type to DuckDB SQL type name. - - Returns: - DuckDB type string. - """ - return _DTYPE_TO_DUCKDB.get(dtype, "VARCHAR") - - def _sub(rel: duckdb.DuckDBPyRelation) -> str: """Wrap a relation's query as a subquery. @@ -216,10 +198,13 @@ def col_rename( col_list = ", ".join(cols) return self._con.sql(f"SELECT {col_list} FROM {_sub(data)}") + def resolve_dtype(self, dtype: Any) -> str: + return to_duckdb_type(dtype) + def col_cast( - self, data: duckdb.DuckDBPyRelation, column: str, dtype: type + self, data: duckdb.DuckDBPyRelation, column: str, dtype: str | type ) -> duckdb.DuckDBPyRelation: - ddb_type = _dtype_to_duckdb(dtype) + ddb_type = self.resolve_dtype(dtype) cols = [ f"CAST({_q(c)} AS {ddb_type}) AS {_q(c)}" if c == column else _q(c) for c in data.columns diff --git a/transformplan/backends/polars.py b/transformplan/backends/polars.py index 9b7e864..4113764 100644 --- a/transformplan/backends/polars.py +++ b/transformplan/backends/polars.py @@ -27,6 +27,7 @@ Numeric, RankMethod, ) +from transformplan.dtypes import to_polars_dtype from transformplan.filters import Filter from transformplan.protocol import frame_hash @@ -126,8 +127,13 @@ def col_rename( ) -> pl.DataFrame: return data.rename({column: new_name}) - def col_cast(self, data: pl.DataFrame, column: str, dtype: type) -> pl.DataFrame: - return data.with_columns(pl.col(column).cast(dtype)) + def resolve_dtype(self, dtype: Any) -> Any: + return to_polars_dtype(dtype) + + def col_cast( + self, data: pl.DataFrame, column: str, dtype: str | type + ) -> pl.DataFrame: + return data.with_columns(pl.col(column).cast(self.resolve_dtype(dtype))) def col_reorder(self, data: pl.DataFrame, columns: list[str]) -> pl.DataFrame: return data.select(columns) diff --git a/transformplan/core.py b/transformplan/core.py index 0d6d4f4..e5e1eb0 100644 --- a/transformplan/core.py +++ b/transformplan/core.py @@ -14,10 +14,11 @@ from __future__ import annotations +import copy import json import time from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Callable from typing import Protocol as TypingProtocol import polars as pl @@ -41,7 +42,9 @@ if TYPE_CHECKING: from collections.abc import Generator, Mapping - from typing_extensions import Self + from typing_extensions import Concatenate, ParamSpec, Self + + P = ParamSpec("P") from transformplan.backends.base import Backend @@ -55,6 +58,12 @@ def _register( params: dict[str, Any], ) -> Self: ... + def _register_each( + self, + op_name: str, + params_list: list[dict[str, Any]], + ) -> Self: ... + class TransformPlanBase: """Base class providing operation registration and execution.""" @@ -64,6 +73,15 @@ class TransformPlanBase: def __init__(self) -> None: """Initialize an empty TransformPlanBase.""" self._operations: list[tuple[str, dict[str, Any]]] = [] + # Step metadata (section, reason) is kept parallel to _operations rather + # than inside params, because params are dispatched to the backend as + # keyword arguments. Same length as _operations, one dict per step. + self._step_meta: list[dict[str, Any]] = [] + self._current_section: str | None = None + # Number of steps registered by the most recent operation call. A single + # call may register several steps when given a sequence of columns, and + # because() must annotate all of them. + self._last_registered: int = 0 @staticmethod def _resolve_backend(backend: Backend | None = None) -> Backend: @@ -88,8 +106,148 @@ def _register( Self for method chaining. """ self._operations.append((op_name, params)) + meta: dict[str, Any] = {} + if self._current_section is not None: + meta["section"] = self._current_section + self._step_meta.append(meta) + self._last_registered = 1 + return self + + def _register_each( + self, + op_name: str, + params_list: list[dict[str, Any]], + ) -> Self: + """Register one operation per parameter set as a single logical call. + + Used by operations that accept a sequence of columns. because() then + annotates every step the call produced, not just the last one. + + Returns: + Self for method chaining. + """ + for params in params_list: + self._register(op_name, params) + self._last_registered = len(params_list) + return self + + def pipe( + self, + fn: Callable[Concatenate[Self, P], Self], + *args: P.args, + **kwargs: P.kwargs, + ) -> Self: + """Apply a function that takes and returns a plan. + + Lets repeated blocks live in their own function without breaking the + method chain. The function registers ordinary steps, so the protocol is + unchanged. + + Args: + fn: Callable receiving this plan plus any extra arguments. + *args: Positional arguments forwarded to fn. + **kwargs: Keyword arguments forwarded to fn. + + Returns: + Whatever fn returns, normally the plan itself. + + Example: + >>> def decimal_hour(plan, source, target): + ... return plan.dt_format(source, "%H", target) + >>> plan = TransformPlan().pipe(decimal_hour, "START", "HOUR") + """ + return fn(self, *args, **kwargs) + + def section(self, name: str | None) -> Self: + """Start a named section that groups the following steps. + + Every step registered afterwards carries this label until the next + section() call. The protocol groups steps by section and reports a + per-section balance. Pass None to end the current section. + + Args: + name: Section label, or None to end the current section. + + Returns: + Self for method chaining. + + Example: + >>> plan = TransformPlan().section("Row filters") + >>> plan = plan.rows_drop_nulls("SHIPPED_AT") + """ + self._current_section = name return self + def because(self, reason: str) -> Self: + """Attach a reason to the step(s) just registered. + + The protocol records that rows were removed; the reason records why. + When the preceding call registered several steps (a sequence of + columns), the reason is attached to all of them. + + Args: + reason: Why this step is part of the pipeline. + + Returns: + Self for method chaining. + + Raises: + ValueError: If no operation has been registered yet. + + Example: + >>> plan = TransformPlan().rows_drop_nulls("DRG_CODE") + >>> plan = plan.because("PEPP cases carry no DRG") + """ + if not self._step_meta or self._last_registered == 0: + msg = "because() must follow an operation." + raise ValueError(msg) + + for meta in self._step_meta[-self._last_registered :]: + meta["reason"] = reason + return self + + def extend(self, other: TransformPlanBase) -> Self: + """Append another plan's steps to this plan. + + Lets a reusable building block be defined as its own plan and inserted + where needed. Steps are deep-copied, so the block can be reused in + several plans without them sharing mutable parameters. + + Args: + other: Plan whose steps are appended. + + Returns: + Self for method chaining. + + Example: + >>> BLOCK = TransformPlan().str_strip("NAME").str_upper("NAME") + >>> plan = TransformPlan().col_drop("temp").extend(BLOCK) + """ + for (op_name, params), meta in zip( + other._operations, + other._step_meta, + strict=False, + ): + self._operations.append((op_name, copy.deepcopy(params))) + new_meta = copy.deepcopy(meta) + # A block's own sections win; otherwise the target plan's section applies. + if "section" not in new_meta and self._current_section is not None: + new_meta["section"] = self._current_section + self._step_meta.append(new_meta) + + self._last_registered = len(other._operations) + return self + + def __add__(self, other: TransformPlanBase) -> Self: + """Combine two plans into a new plan. + + Returns: + New plan with the steps of both, leaving both operands unchanged. + """ + combined = copy.deepcopy(self) + combined.extend(other) + return combined + def process( self, data: Any, # noqa: ANN401 @@ -140,7 +298,8 @@ def process( } protocol.set_metadata(references=ref_meta) - for op_name, params in self._operations: + for index, (op_name, params) in enumerate(self._operations): + meta = self._step_meta[index] if index < len(self._step_meta) else {} old_shape = resolved.get_shape(data) start = time.perf_counter() @@ -163,6 +322,8 @@ def process( new_shape=resolved.get_shape(data), elapsed=elapsed, output_hash=resolved.compute_hash(data), + section=meta.get("section"), + reason=meta.get("reason"), ) return data, protocol @@ -269,13 +430,18 @@ def to_dict(self) -> dict[str, Any]: Dictionary representation of the pipeline. """ steps = [] - for op_name, params in self._operations: - steps.append( - { - "operation": op_name, - "params": params, - } - ) + for index, (op_name, params) in enumerate(self._operations): + step: dict[str, Any] = { + "operation": op_name, + "params": params, + } + meta = self._step_meta[index] if index < len(self._step_meta) else {} + # Only written when set, so plans without metadata serialize as before. + if meta.get("section") is not None: + step["section"] = meta["section"] + if meta.get("reason") is not None: + step["reason"] = meta["reason"] + steps.append(step) return { "version": self.VERSION, @@ -315,6 +481,14 @@ def from_dict(cls, data: dict[str, Any]) -> Self: msg = f"Invalid parameters for operation '{op_name}': {e}" raise ValueError(msg) from e + # Metadata keys are absent in plans written before they existed. + if plan._step_meta: + meta = plan._step_meta[-1] + if step.get("section") is not None: + meta["section"] = step["section"] + if step.get("reason") is not None: + meta["reason"] = step["reason"] + return plan def to_json(self, path: str | Path | None = None, indent: int = 2) -> str: @@ -379,10 +553,23 @@ def to_python(self, variable_name: str = "plan") -> str: lines = ["from transformplan import TransformPlan, Col", ""] lines.extend((f"{variable_name} = (", " TransformPlan()")) - for op_name, params in self._operations: + current_section: str | None = None + for index, (op_name, params) in enumerate(self._operations): + meta = self._step_meta[index] if index < len(self._step_meta) else {} + + section = meta.get("section") + if section != current_section: + arg = f'"{section}"' if section is not None else "None" + lines.append(f" .section({arg})") + current_section = section + param_str = self._format_params_as_python(params) lines.append(f" .{op_name}({param_str})") + reason = meta.get("reason") + if reason is not None: + lines.append(f' .because("{reason}")') + lines.append(")") return "\n".join(lines) diff --git a/transformplan/dtypes.py b/transformplan/dtypes.py new file mode 100644 index 0000000..707363f --- /dev/null +++ b/transformplan/dtypes.py @@ -0,0 +1,177 @@ +"""Canonical dtype names for serializable transformation plans. + +A plan must survive a round trip through JSON. Polars dtype objects do not: +``json.dumps`` rejects them, and ``to_python()`` renders them as bare names that +no import statement provides. This module defines a canonical string name per +dtype so that plans store a name, not an object. + +The names follow Polars spelling ("Float64", "String", "Boolean"), and each +backend resolves them to its own type system via ``Backend.resolve_dtype()``. + +Functions: + normalize_dtype: Convert a dtype specification to its canonical name. + canonical_dtype_names: List all supported canonical names. + +Example: + >>> normalize_dtype(pl.Float64) + 'Float64' + >>> normalize_dtype(float) + 'Float64' + >>> normalize_dtype("Utf8") + 'String' +""" + +from __future__ import annotations + +import datetime +from typing import Any + +import polars as pl + +# Canonical name -> Polars dtype. The canonical spelling is the Polars one. +_CANONICAL_TO_POLARS: dict[str, Any] = { + "Int8": pl.Int8, + "Int16": pl.Int16, + "Int32": pl.Int32, + "Int64": pl.Int64, + "UInt8": pl.UInt8, + "UInt16": pl.UInt16, + "UInt32": pl.UInt32, + "UInt64": pl.UInt64, + "Float32": pl.Float32, + "Float64": pl.Float64, + "Boolean": pl.Boolean, + "String": pl.String, + "Binary": pl.Binary, + "Date": pl.Date, + "Datetime": pl.Datetime, + "Time": pl.Time, + "Duration": pl.Duration, +} + +# Canonical name -> DuckDB SQL type. +_CANONICAL_TO_DUCKDB: dict[str, str] = { + "Int8": "TINYINT", + "Int16": "SMALLINT", + "Int32": "INTEGER", + "Int64": "BIGINT", + "UInt8": "UTINYINT", + "UInt16": "USMALLINT", + "UInt32": "UINTEGER", + "UInt64": "UBIGINT", + "Float32": "FLOAT", + "Float64": "DOUBLE", + "Boolean": "BOOLEAN", + "String": "VARCHAR", + "Binary": "BLOB", + "Date": "DATE", + "Datetime": "TIMESTAMP", + "Time": "TIME", + "Duration": "INTERVAL", +} + +# Builtin Python types accepted as dtype shorthands. +_PYTHON_TO_CANONICAL: list[tuple[type, str]] = [ + (bool, "Boolean"), # before int: bool is a subclass of int + (int, "Int64"), + (float, "Float64"), + (str, "String"), + (bytes, "Binary"), + (datetime.datetime, "Datetime"), + (datetime.date, "Date"), + (datetime.time, "Time"), + (datetime.timedelta, "Duration"), +] + +# Accepted spellings that are not the canonical one. +_ALIASES: dict[str, str] = { + "Utf8": "String", + "Bool": "Boolean", + "Str": "String", +} + + +def canonical_dtype_names() -> list[str]: + """List the canonical dtype names. + + Returns: + Sorted list of supported names. + """ + return sorted(_CANONICAL_TO_POLARS) + + +def normalize_dtype(dtype: Any) -> Any: # noqa: ANN401 + """Convert a dtype specification to its canonical name. + + Accepts a canonical name, a known alias, a Polars dtype class, or a builtin + Python type. Anything not recognized is returned unchanged, so that exotic + or parameterized dtypes (``pl.Datetime("us")``, ``pl.Categorical``, + ``pl.List``) keep working exactly as before -- such plans simply remain + non-serializable, as they were prior to this function existing. + + Args: + dtype: Dtype specification. + + Returns: + Canonical name string, or the input unchanged if not recognized. + + Raises: + ValueError: If a string is given that is neither canonical nor a known + alias. Strings are unambiguously meant as names, so a typo is + reported when the plan is built rather than when it runs. + """ + if isinstance(dtype, str): + resolved = _ALIASES.get(dtype, dtype) + if resolved not in _CANONICAL_TO_POLARS: + msg = ( + f"Unknown dtype name: {dtype!r}. " + f"Supported names: {', '.join(canonical_dtype_names())}" + ) + raise ValueError(msg) + return resolved + + # Builtin Python types. Identity comparison avoids triggering the custom + # __eq__ that Polars dtypes define. + for py_type, name in _PYTHON_TO_CANONICAL: + if dtype is py_type: + return name + + # Polars dtype classes. Parameterized instances (pl.Datetime("us")) have no + # __name__ and fall through unchanged, keeping their parameters intact. + name = getattr(dtype, "__name__", None) + if isinstance(name, str): + canonical = _ALIASES.get(name, name) + if _CANONICAL_TO_POLARS.get(canonical) is dtype: + return canonical + + return dtype + + +def to_polars_dtype(dtype: Any) -> Any: # noqa: ANN401 + """Resolve a dtype specification to a Polars dtype. + + Args: + dtype: Canonical name, alias, Python type, or Polars dtype. + + Returns: + Polars dtype, or the input unchanged if it is not a known name. + """ + normalized = normalize_dtype(dtype) + if isinstance(normalized, str): + return _CANONICAL_TO_POLARS[normalized] + return normalized + + +def to_duckdb_type(dtype: Any) -> str: # noqa: ANN401 + """Resolve a dtype specification to a DuckDB SQL type name. + + Args: + dtype: Canonical name, alias, Python type, or Polars dtype. + + Returns: + DuckDB SQL type string. Unrecognized dtypes fall back to VARCHAR. + """ + normalized = normalize_dtype(dtype) + if isinstance(normalized, str): + return _CANONICAL_TO_DUCKDB[normalized] + return "VARCHAR" diff --git a/transformplan/ops/_common.py b/transformplan/ops/_common.py new file mode 100644 index 0000000..15b7efe --- /dev/null +++ b/transformplan/ops/_common.py @@ -0,0 +1,25 @@ +"""Shared helpers for operation mixins. + +Functions: + as_columns: Normalize a column argument to a list of column names. +""" + +from __future__ import annotations + +from typing import Sequence + + +def as_columns(column: str | Sequence[str]) -> list[str]: + """Normalize a column argument to a list of names. + + Operations that work in place on a column accept either a single name or a + sequence of names. A sequence registers one step per column, which keeps the + protocol step-granular and leaves backends and validators untouched. + + Args: + column: A single column name or a sequence of names. + + Returns: + List of column names. + """ + return [column] if isinstance(column, str) else list(column) diff --git a/transformplan/ops/column.py b/transformplan/ops/column.py index 209c351..2690280 100644 --- a/transformplan/ops/column.py +++ b/transformplan/ops/column.py @@ -30,6 +30,9 @@ from typing import TYPE_CHECKING, Any, Literal, Sequence +from transformplan.dtypes import normalize_dtype +from transformplan.ops._common import as_columns + if TYPE_CHECKING: from typing import Any @@ -49,13 +52,24 @@ def _register( params: dict[str, Any], ) -> Self: ... - def col_drop(self, column: str) -> Self: - """Drop a column from the DataFrame. + def _register_each( + self, + op_name: str, + params_list: list[dict[str, Any]], + ) -> Self: ... + + def col_drop(self, column: str | Sequence[str]) -> Self: + """Drop one or more columns from the DataFrame. + + Args: + column: Column name, or a sequence of names to drop. Returns: Self for method chaining. """ - return self._register("col_drop", {"column": column}) + return self._register_each( + "col_drop", [{"column": col} for col in as_columns(column)] + ) def col_rename(self, column: str, new_name: str) -> Self: """Rename a column. @@ -65,13 +79,26 @@ def col_rename(self, column: str, new_name: str) -> Self: """ return self._register("col_rename", {"column": column, "new_name": new_name}) - def col_cast(self, column: str, dtype: type) -> Self: - """Cast a column to a different dtype. + def col_cast(self, column: str | Sequence[str], dtype: str | type) -> Self: + """Cast one or more columns to a different dtype. + + The dtype is stored under its canonical name (e.g. ``"Float64"``) so the + plan stays serializable. Polars dtypes and builtin Python types are + accepted and normalized; each backend resolves the name to its own type. + + Args: + column: Column name, or a sequence of names to cast. + dtype: Target dtype as a canonical name (``"Float64"``), a Polars + dtype (``pl.Float64``), or a Python type (``float``). Returns: Self for method chaining. """ - return self._register("col_cast", {"column": column, "dtype": dtype}) + normalized = normalize_dtype(dtype) + return self._register_each( + "col_cast", + [{"column": col, "dtype": normalized} for col in as_columns(column)], + ) def col_reorder(self, columns: Sequence[str]) -> Self: """Reorder columns. Unlisted columns are dropped. @@ -91,14 +118,14 @@ def col_duplicate(self, column: str, new_name: str) -> Self: def col_fill_null( self, - column: str, + column: str | Sequence[str], value: Any = None, # noqa: ANN401 strategy: FillNullStrategy | None = None, ) -> Self: """Fill null values in a column. Args: - column: Column to fill. + column: Column name, or a sequence of names to fill. value: Value to fill nulls with (if strategy is None). strategy: Fill strategy - 'forward', 'backward', 'mean', 'min', 'max', 'zero', 'one'. @@ -106,9 +133,12 @@ def col_fill_null( Returns: Self for method chaining. """ - return self._register( + return self._register_each( "col_fill_null", - {"column": column, "value": value, "strategy": strategy}, + [ + {"column": col, "value": value, "strategy": strategy} + for col in as_columns(column) + ], ) def col_drop_null(self, columns: str | Sequence[str] | None = None) -> Self: @@ -124,13 +154,18 @@ def col_drop_null(self, columns: str | Sequence[str] | None = None) -> Self: columns = [columns] return self._register("col_drop_null", {"columns": columns}) - def col_drop_zero(self, column: str) -> Self: - """Drop rows where the specified column is zero. + def col_drop_zero(self, column: str | Sequence[str]) -> Self: + """Drop rows where the specified column(s) are zero. + + Args: + column: Column name, or a sequence of names to check. Returns: Self for method chaining. """ - return self._register("col_drop_zero", {"column": column}) + return self._register_each( + "col_drop_zero", [{"column": col} for col in as_columns(column)] + ) def col_add( self, diff --git a/transformplan/ops/map.py b/transformplan/ops/map.py index 6aa8a93..07f2b84 100644 --- a/transformplan/ops/map.py +++ b/transformplan/ops/map.py @@ -34,6 +34,8 @@ from typing import TYPE_CHECKING, Any, Literal, Sequence +from transformplan.ops._common import as_columns + if TYPE_CHECKING: from typing_extensions import Self @@ -49,9 +51,15 @@ def _register( params: dict[str, Any], ) -> Self: ... + def _register_each( + self, + op_name: str, + params_list: list[dict[str, Any]], + ) -> Self: ... + def map_values( self, - column: str, + column: str | Sequence[str], mapping: dict[Any, Any], default: Any = None, # noqa: ANN401 *, @@ -60,7 +68,7 @@ def map_values( """Map values in a column using a dictionary. Args: - column: Column to transform. + column: Column name, or a sequence of names, to transform. mapping: Dictionary mapping old values to new values. default: Default value for unmapped values (if keep_unmapped=False). keep_unmapped: If True, keep original value when not in mapping. @@ -68,14 +76,17 @@ def map_values( Returns: Self for method chaining. """ - return self._register( + return self._register_each( "map_values", - { - "column": column, - "mapping": mapping, - "default": default, - "keep_unmapped": keep_unmapped, - }, + [ + { + "column": col, + "mapping": mapping, + "default": default, + "keep_unmapped": keep_unmapped, + } + for col in as_columns(column) + ], ) def map_discretize( @@ -110,29 +121,38 @@ def map_discretize( }, ) - def map_bool_to_int(self, column: str) -> Self: + def map_bool_to_int(self, column: str | Sequence[str]) -> Self: """Convert a boolean column to integer (True=1, False=0). Returns: Self for method chaining. """ - return self._register("map_bool_to_int", {"column": column}) + return self._register_each( + "map_bool_to_int", + [{"column": col} for col in as_columns(column)], + ) - def map_null_to_value(self, column: str, value: Any) -> Self: # noqa: ANN401 + def map_null_to_value(self, column: str | Sequence[str], value: Any) -> Self: # noqa: ANN401 """Replace null values with a specific value. Returns: Self for method chaining. """ - return self._register("map_null_to_value", {"column": column, "value": value}) + return self._register_each( + "map_null_to_value", + [{"column": col, "value": value} for col in as_columns(column)], + ) - def map_value_to_null(self, column: str, value: Any) -> Self: # noqa: ANN401 + def map_value_to_null(self, column: str | Sequence[str], value: Any) -> Self: # noqa: ANN401 """Replace a specific value with null. Returns: Self for method chaining. """ - return self._register("map_value_to_null", {"column": column, "value": value}) + return self._register_each( + "map_value_to_null", + [{"column": col, "value": value} for col in as_columns(column)], + ) def map_case( self, diff --git a/transformplan/ops/math.py b/transformplan/ops/math.py index a6fb226..2779419 100644 --- a/transformplan/ops/math.py +++ b/transformplan/ops/math.py @@ -47,7 +47,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Literal, Union +from typing import TYPE_CHECKING, Literal, Sequence, Union + +from transformplan.ops._common import as_columns if TYPE_CHECKING: from typing import Any @@ -70,41 +72,59 @@ def _register( params: dict[str, Any], ) -> Self: ... - def math_add(self, column: str, value: Numeric) -> Self: + def _register_each( + self, + op_name: str, + params_list: list[dict[str, Any]], + ) -> Self: ... + + def math_add(self, column: str | Sequence[str], value: Numeric) -> Self: """Add a scalar value to a column. Returns: Self for method chaining. """ - return self._register("math_add", {"column": column, "value": value}) + return self._register_each( + "math_add", + [{"column": col, "value": value} for col in as_columns(column)], + ) - def math_subtract(self, column: str, value: Numeric) -> Self: + def math_subtract(self, column: str | Sequence[str], value: Numeric) -> Self: """Subtract a scalar value from a column. Returns: Self for method chaining. """ - return self._register("math_subtract", {"column": column, "value": value}) + return self._register_each( + "math_subtract", + [{"column": col, "value": value} for col in as_columns(column)], + ) - def math_multiply(self, column: str, value: Numeric) -> Self: + def math_multiply(self, column: str | Sequence[str], value: Numeric) -> Self: """Multiply a column by a scalar value. Returns: Self for method chaining. """ - return self._register("math_multiply", {"column": column, "value": value}) + return self._register_each( + "math_multiply", + [{"column": col, "value": value} for col in as_columns(column)], + ) - def math_divide(self, column: str, value: Numeric) -> Self: + def math_divide(self, column: str | Sequence[str], value: Numeric) -> Self: """Divide a column by a scalar value. Returns: Self for method chaining. """ - return self._register("math_divide", {"column": column, "value": value}) + return self._register_each( + "math_divide", + [{"column": col, "value": value} for col in as_columns(column)], + ) def math_clamp( self, - column: str, + column: str | Sequence[str], lower: Numeric | None = None, upper: Numeric | None = None, ) -> Self: @@ -113,8 +133,12 @@ def math_clamp( Returns: Self for method chaining. """ - return self._register( - "math_clamp", {"column": column, "lower": lower, "upper": upper} + return self._register_each( + "math_clamp", + [ + {"column": col, "lower": lower, "upper": upper} + for col in as_columns(column) + ], ) def math_add_columns(self, column_a: str, column_b: str, new_column: str) -> Self: @@ -167,41 +191,49 @@ def math_divide_columns( {"column_a": column_a, "column_b": column_b, "new_column": new_column}, ) - def math_set_min(self, column: str, min_value: Numeric) -> Self: + def math_set_min(self, column: str | Sequence[str], min_value: Numeric) -> Self: """Set a minimum value for a column (values below are raised to min). Returns: Self for method chaining. """ - return self._register( - "math_set_min", {"column": column, "min_value": min_value} + return self._register_each( + "math_set_min", + [{"column": col, "min_value": min_value} for col in as_columns(column)], ) - def math_set_max(self, column: str, max_value: Numeric) -> Self: + def math_set_max(self, column: str | Sequence[str], max_value: Numeric) -> Self: """Set a maximum value for a column (values above are lowered to max). Returns: Self for method chaining. """ - return self._register( - "math_set_max", {"column": column, "max_value": max_value} + return self._register_each( + "math_set_max", + [{"column": col, "max_value": max_value} for col in as_columns(column)], ) - def math_abs(self, column: str) -> Self: + def math_abs(self, column: str | Sequence[str]) -> Self: """Take absolute value of a column. Returns: Self for method chaining. """ - return self._register("math_abs", {"column": column}) + return self._register_each( + "math_abs", + [{"column": col} for col in as_columns(column)], + ) - def math_round(self, column: str, decimals: int = 0) -> Self: + def math_round(self, column: str | Sequence[str], decimals: int = 0) -> Self: """Round a column to specified decimal places. Returns: Self for method chaining. """ - return self._register("math_round", {"column": column, "decimals": decimals}) + return self._register_each( + "math_round", + [{"column": col, "decimals": decimals} for col in as_columns(column)], + ) def math_percent_of( self, diff --git a/transformplan/ops/string.py b/transformplan/ops/string.py index c90dab2..507dca4 100644 --- a/transformplan/ops/string.py +++ b/transformplan/ops/string.py @@ -28,7 +28,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Sequence + +from transformplan.ops._common import as_columns if TYPE_CHECKING: from typing import Any @@ -47,9 +49,15 @@ def _register( params: dict[str, Any], ) -> Self: ... + def _register_each( + self, + op_name: str, + params_list: list[dict[str, Any]], + ) -> Self: ... + def str_replace( self, - column: str, + column: str | Sequence[str], pattern: str, replacement: str, *, @@ -58,7 +66,7 @@ def str_replace( """Replace occurrences of a pattern in a string column. Args: - column: Column to modify. + column: Column name, or a sequence of names, to modify. pattern: Pattern to search for. replacement: String to replace with. literal: If True, treat pattern as literal string. If False, treat as regex. @@ -66,50 +74,62 @@ def str_replace( Returns: Self for method chaining. """ - return self._register( + return self._register_each( "str_replace", - { - "column": column, - "pattern": pattern, - "replacement": replacement, - "literal": literal, - }, + [ + { + "column": col, + "pattern": pattern, + "replacement": replacement, + "literal": literal, + } + for col in as_columns(column) + ], ) def str_slice( self, - column: str, + column: str | Sequence[str], offset: int, length: int | None = None, ) -> Self: """Extract a substring from a string column. Args: - column: Column to modify. + column: Column name, or a sequence of names, to modify. offset: Start position (0-indexed, negative counts from end). length: Number of characters to extract (None = to end). Returns: Self for method chaining. """ - return self._register( - "str_slice", {"column": column, "offset": offset, "length": length} + return self._register_each( + "str_slice", + [ + {"column": col, "offset": offset, "length": length} + for col in as_columns(column) + ], ) - def str_truncate(self, column: str, max_length: int, suffix: str = "...") -> Self: + def str_truncate( + self, column: str | Sequence[str], max_length: int, suffix: str = "..." + ) -> Self: """Truncate strings to a maximum length with optional suffix. Args: - column: Column to modify. + column: Column name, or a sequence of names, to modify. max_length: Maximum length of the string (including suffix). suffix: Suffix to append to truncated strings. Returns: Self for method chaining. """ - return self._register( + return self._register_each( "str_truncate", - {"column": column, "max_length": max_length, "suffix": suffix}, + [ + {"column": col, "max_length": max_length, "suffix": suffix} + for col in as_columns(column) + ], ) def str_split( @@ -141,37 +161,46 @@ def str_split( }, ) - def str_lower(self, column: str) -> Self: + def str_lower(self, column: str | Sequence[str]) -> Self: """Convert string column to lowercase. Returns: Self for method chaining. """ - return self._register("str_lower", {"column": column}) + return self._register_each( + "str_lower", + [{"column": col} for col in as_columns(column)], + ) - def str_upper(self, column: str) -> Self: + def str_upper(self, column: str | Sequence[str]) -> Self: """Convert string column to uppercase. Returns: Self for method chaining. """ - return self._register("str_upper", {"column": column}) + return self._register_each( + "str_upper", + [{"column": col} for col in as_columns(column)], + ) - def str_strip(self, column: str, chars: str | None = None) -> Self: + def str_strip(self, column: str | Sequence[str], chars: str | None = None) -> Self: """Strip leading and trailing characters from a string column. Args: - column: Column to modify. + column: Column name, or a sequence of names, to modify. chars: Characters to strip (None = whitespace). Returns: Self for method chaining. """ - return self._register("str_strip", {"column": column, "chars": chars}) + return self._register_each( + "str_strip", + [{"column": col, "chars": chars} for col in as_columns(column)], + ) def str_pad( self, - column: str, + column: str | Sequence[str], length: int, fill_char: str = " ", side: str = "left", @@ -179,7 +208,7 @@ def str_pad( """Pad a string column to a specified length. Args: - column: Column to modify. + column: Column name, or a sequence of names, to modify. length: Target length. fill_char: Character to pad with. side: 'left' or 'right'. @@ -187,9 +216,12 @@ def str_pad( Returns: Self for method chaining. """ - return self._register( + return self._register_each( "str_pad", - {"column": column, "length": length, "fill_char": fill_char, "side": side}, + [ + {"column": col, "length": length, "fill_char": fill_char, "side": side} + for col in as_columns(column) + ], ) def str_concat( diff --git a/transformplan/protocol.py b/transformplan/protocol.py index 8c2ce95..d10fdaf 100644 --- a/transformplan/protocol.py +++ b/transformplan/protocol.py @@ -96,6 +96,8 @@ def add_step( new_shape: tuple[int, int], elapsed: float, output_hash: str, + section: str | None = None, + reason: str | None = None, ) -> None: """Record a transformation step in the protocol. @@ -106,6 +108,8 @@ def add_step( new_shape: Shape after operation (rows, cols). elapsed: Time taken in seconds. output_hash: Hash of the output DataFrame. + section: Name of the section this step belongs to, if any. + reason: Why this step is part of the pipeline, if given. """ self._steps.append( { @@ -118,6 +122,8 @@ def add_step( "cols_changed": old_shape[1] - new_shape[1], "elapsed_seconds": round(elapsed, 4), "output_hash": output_hash, + "section": section, + "reason": reason, } ) @@ -171,6 +177,8 @@ def to_dataframe(self) -> pl.DataFrame: "cols_changed": 0, "elapsed_seconds": 0.0, "output_hash": self._input_hash, + "section": None, + "reason": None, } ) @@ -202,6 +210,8 @@ def to_csv(self, path: str | Path) -> None: "cols_changed": 0, "elapsed_seconds": 0.0, "output_hash": self._input_hash, + "section": None, + "reason": None, } ) @@ -216,6 +226,8 @@ def to_csv(self, path: str | Path) -> None: "cols_changed": step["cols_changed"], "elapsed_seconds": step["elapsed_seconds"], "output_hash": step["output_hash"], + "section": step.get("section"), + "reason": step.get("reason"), } for step in self._steps ) @@ -247,6 +259,8 @@ def to_dict(self) -> dict[str, Any]: "cols_changed": s["cols_changed"], "elapsed_seconds": s["elapsed_seconds"], "output_hash": s["output_hash"], + "section": s.get("section"), + "reason": s.get("reason"), } for s in self._steps ], @@ -280,6 +294,9 @@ def from_dict(cls, data: dict[str, Any]) -> Protocol: "cols_changed": step["cols_changed"], "elapsed_seconds": step["elapsed_seconds"], "output_hash": step["output_hash"], + # Older protocol files predate these keys. + "section": step.get("section"), + "reason": step.get("reason"), } ) @@ -374,9 +391,16 @@ def summary(self, *, show_params: bool = True) -> str: # noqa: C901 # Total time total_time = sum(s["elapsed_seconds"] for s in self._steps) + lines.append(f"Total time: {total_time:.4f}s") + + # Per-section balance, only when sections are in use + section_lines = self._section_summary() + if section_lines: + lines.extend(["-" * 70, "", "SECTIONS", ""]) + lines.extend(section_lines) + lines.extend( [ - f"Total time: {total_time:.4f}s", "-" * 70, "", f"{'#':<4} {'Operation':<20} {'Rows':<12} {'Cols':<12} {'Time':<10} {'Hash':<16}", @@ -393,6 +417,7 @@ def summary(self, *, show_params: bool = True) -> str: # noqa: C901 # Operation rows no_effect_steps = [] + current_section: str | None = None for step in self._steps: step_num = str(step["step"]) op = step["operation"] @@ -427,6 +452,13 @@ def summary(self, *, show_params: bool = True) -> str: # noqa: C901 # Add marker for no-effect steps marker = " ○" if no_effect else "" + # Section heading whenever the section changes + section = step.get("section") + if section != current_section: + if section is not None: + lines.append(f"[{section}]") + current_section = section + lines.append( f"{step_num:<4} {op:<20} {row_str:<12} {col_str:<12} {time_str:<10} {hash_str:<16}{marker}" ) @@ -436,6 +468,10 @@ def summary(self, *, show_params: bool = True) -> str: # noqa: C901 params_str = self._format_params(step["params"]) lines.append(f" └─ {params_str}") + # Reason + if step.get("reason"): + lines.append(f" ↳ {step['reason']}") + lines.append("=" * 70) # Add note about no-effect steps @@ -446,6 +482,39 @@ def summary(self, *, show_params: bool = True) -> str: # noqa: C901 return "\n".join(lines) + def _section_summary(self) -> list[str]: + """Build the per-section balance lines. + + Returns: + Formatted lines, or an empty list if no step carries a section. + """ + if not any(s.get("section") for s in self._steps): + return [] + + groups: list[tuple[str | None, list[dict[str, Any]]]] = [] + for step in self._steps: + section = step.get("section") + if groups and groups[-1][0] == section: + groups[-1][1].append(step) + else: + groups.append((section, [step])) + + lines = [] + for section, steps in groups: + label = section if section is not None else "(no section)" + rows_before = steps[0]["old_shape"][0] + rows_after = steps[-1]["new_shape"][0] + cols_delta = steps[-1]["new_shape"][1] - steps[0]["old_shape"][1] + + detail = f"{len(steps)} step{'s' if len(steps) != 1 else ''}" + if cols_delta: + detail += f", {cols_delta:+d} cols" + + lines.append( + f" {label:<24} {rows_before:>8} → {rows_after:>8} rows ({detail})" + ) + return lines + def _format_params(self, params: dict[str, Any], max_length: int = 60) -> str: """Format params dict as a readable string. diff --git a/transformplan/validation.py b/transformplan/validation.py index f78a421..c7d633c 100644 --- a/transformplan/validation.py +++ b/transformplan/validation.py @@ -442,6 +442,17 @@ def __init__(self, schema: dict[str, Any], backend: Backend | None = None) -> No backend = PolarsBackend() self._backend = backend + def resolve_dtype(self, dtype: Any) -> Any: # noqa: ANN401 + """Resolve a canonical dtype name to the backend's type. + + Args: + dtype: Canonical dtype name or a native dtype. + + Returns: + The backend's dtype. + """ + return self._backend.resolve_dtype(dtype) + @property def columns(self) -> set[str]: """Get set of column names. @@ -704,7 +715,9 @@ def _validate_col_cast( tracker: SchemaTracker, params: dict[str, Any], result: ValidationResult, step: int ) -> None: column = params["column"] - target_dtype = params["dtype"] + # Plans store dtypes as canonical names; resolve to the backend's type so + # downstream validators see the same dtype objects as the rest of the schema. + target_dtype = tracker.resolve_dtype(params["dtype"]) if _check_column_exists(tracker, column, result, step, "col_cast"): tracker.set_dtype(column, target_dtype) From f504f4018822d6862ac2772c97dbdc7bffaa4c65 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philipp=20Gro=C3=9Fer?= Date: Thu, 27 Aug 2026 18:43:08 +0200 Subject: [PATCH 2/2] readme update --- README.md | 112 +++++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 89 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index c925559..d91b43e 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,8 @@ - **Audit trails**: Generate complete audit protocols with deterministic DataFrame hashing - **Multi-backend support**: Polars (default) and DuckDB backends with a pluggable Backend ABC - **Serializable pipelines**: Save and load transformation plans as JSON +- **Structured plans**: Group long pipelines into named sections and compose them from reusable blocks +- **Documented decisions**: Record *why* a step exists, not just what it did ## Quick Example @@ -25,20 +27,19 @@ from transformplan import TransformPlan, Col # Build readable pipelines with 89 chainable operations plan = ( TransformPlan() - # Standardize column names + .section("Standardize") .col_rename(column="PatientID", new_name="patient_id") .col_rename(column="DOB", new_name="date_of_birth") .str_strip(column="patient_id") + .col_drop(["tmp_a", "tmp_b"]).because("scratch columns from the export") - # Calculate derived values - .dt_age_years(column="date_of_birth", new_column="age") - .math_clamp(column="age", min_value=0, max_value=120) + .section("Derive") + .dt_age_years(birth_column="date_of_birth", new_column="age") + .col_cast("raw_weight", "Float64") + .map_discretize(column="age", bins=[18, 40, 65], labels=["minor", "young", "adult", "senior"], new_column="age_group") - # Categorize patients age - .map_discretize(column="age", bins=[18, 40, 65], labels=["young", "adult", "senior"], new_column="age_group") - - # Filter and clean - .rows_filter(Col("age") >= 18) + .section("Filter") + .rows_filter(Col("age") >= 18).because("cohort is adults only") .rows_drop_nulls(columns=["patient_id", "age"]) .col_drop(column="date_of_birth") ) @@ -64,27 +65,87 @@ protocol.print(show_params=False) ====================================================================== TRANSFORM PROTOCOL ====================================================================== -Input: 1000 rows × 5 cols [a4f8b2c1] -Output: 847 rows × 5 cols [e7d3f9a2] -Total time: 0.0247s +Input: 1000 rows x 5 cols [a8bfc98263e8aa4c] +Output: 841 rows x 4 cols [6f62d5e6d677cb81] +Total time: 0.0011s +---------------------------------------------------------------------- + +SECTIONS + + Standardize 1000 → 1000 rows (5 steps, -2 cols) + Derive 1000 → 1000 rows (3 steps, +2 cols) + Filter 1000 → 841 rows (3 steps, -1 cols) ---------------------------------------------------------------------- # Operation Rows Cols Time Hash ---------------------------------------------------------------------- -0 input 1000 5 - a4f8b2c1 -1 col_rename 1000 5 0.0012s b2e4a7f3 -2 col_rename 1000 5 0.0008s c9d1e5b8 -3 str_strip 1000 5 0.0013s c9d1e5b8 ○ -4 dt_age_years 1000 6 (+1) 0.0041s d4f2c8a1 -5 math_clamp 1000 6 0.0015s e1b7d3f9 -6 map_discretize 1000 7 (+1) 0.0028s f8a4c2e6 -7 rows_filter 858 (-142) 7 0.0037s a2e9f4b7 -8 rows_drop_nulls 847 (-11) 7 0.0019s b5c1d8e3 -9 col_drop 847 6 (-1) 0.0006s e7d3f9a2 +0 input 1000 5 - a8bfc98263e8aa4c +[Standardize] +1 col_rename 1000 5 0.0000s 4a3d3ba5644cf847 +2 col_rename 1000 5 0.0000s e535f4ddd1e78927 +3 str_strip 1000 5 0.0001s 14cf5b02000d13a4 +4 col_drop 1000 4 (-1) 0.0001s ccaddd0bc65974d7 + ↳ scratch columns from the export +5 col_drop 1000 3 (-1) 0.0000s 08a006cf6351723c + ↳ scratch columns from the export +[Derive] +6 dt_age_years 1000 4 (+1) 0.0001s 2e3c2e80b77f7e10 +7 col_cast 1000 4 0.0001s e7e31d60ee332aa6 +8 map_discretize 1000 5 (+1) 0.0002s f3c8c0998d8ee820 +[Filter] +9 rows_filter 851 (-149) 5 0.0002s 97f54c043b457d0c + ↳ cohort is adults only +10 rows_drop_nulls 841 (-10) 5 0.0002s 7bb8de08021dfe20 +11 col_drop 841 4 (-1) 0.0001s 6f62d5e6d677cb81 ====================================================================== -○ = no effect (steps 3 did not change data) ``` +Sections give a plan with a hundred steps a readable balance instead of one flat +list, and `because()` records the domain knowledge that otherwise lives only in a +code comment — the question an audit actually asks. + +### Composing Long Pipelines + +Repeated blocks can live in their own function or their own plan, without breaking +the chain: + +```python +# A function that takes and returns a plan +def decimal_hour(plan, source, target): + return ( + plan.dt_format(source, "%H", target) + .col_cast(target, "Float64") + ) + +# A reusable block, defined as a plan of its own +CLEAN_IDS = TransformPlan().str_strip("patient_id").str_upper("patient_id") + +plan = ( + TransformPlan() + .pipe(decimal_hour, "admitted_at", "admit_hour") + .pipe(decimal_hour, "discharged_at", "discharge_hour") + .extend(CLEAN_IDS) # or: plan_a + plan_b +) +``` + +`pipe()` registers no step of its own, so the protocol is unchanged. `extend()` +deep-copies the block, so the same block can be reused across plans without them +sharing state — and unlike `pipe()`, its steps are serialized with the plan. + +Operations that work in place also take a sequence of columns, so repetition +collapses into one call: + +```python +plan = ( + TransformPlan() + .col_drop(["fiscal_year", "import_batch", "row_checksum"]) + .str_slice(["primary_code", "secondary_code"], 0, 3) + .col_cast(["weight", "height"], "Float64") +) +``` + +Each column still becomes its own protocol step. + ### DuckDB Backend Run the same pipelines on DuckDB for SQL-based execution and native large-file handling: @@ -119,6 +180,11 @@ result, protocol = plan.process(rel, backend=DuckDBBackend(con)) | **dt\_** | Datetime operations | `dt_year`, `dt_month`, `dt_parse`, `dt_age_years`, `dt_diff_days` | | **map\_** | Value mapping & encoding | `map_values`, `map_discretize`, `map_onehot`, `map_ordinal` | +24 in-place operations accept either a single column or a sequence of columns. +Operations that name an output column (`new_column`) take one column at a time. + +See the [roadmap](docs/roadmap.md) for planned additions. + ## Installation ```bash