Inspecting Pipeline Runs

Pipeline.run() returns a PipelineResult. Most tests already have the Step returned by add_step, so they can read its output with result.get_output(step). If you only have the step name, pass the expected runtime type: result.get_output("step_name", ValueType).

Inspect what a writer received

For an end-to-end test, check the DataFrame passed to the final writer. Use get_one_input("write_orders", pl.LazyFrame) when the writer has one DataFrame input. The example also reads the upstream step's output:

from dataclasses import dataclass
from typing import Annotated

import polars as pl

from data_joinery import Pipeline, Project, transform


@dataclass
class Order:
    order_id: int


@transform
def read_orders() -> Annotated[pl.LazyFrame, Project(Order)]:
    return pl.LazyFrame({"order_id": [1, 2]})


@transform
def write_orders(orders: Annotated[pl.LazyFrame, Project(Order)]) -> None:
    pass


pipeline = Pipeline()
read_step = pipeline.add_step(read_orders)
writer_step = pipeline.add_step(write_orders)
read_step >> writer_step

result = pipeline.run()
print("Read output:", result.get_output(read_step).collect().to_dicts())
print(
    "Passed to writer:",
    result.get_one_input("write_orders", pl.LazyFrame).collect().to_dicts(),
)

Outputs:

Read output: [{'order_id': 1}, {'order_id': 2}]
Passed to writer: [{'order_id': 1}, {'order_id': 2}]

Writer inputs are recorded before coercion

get_input() and get_one_input() return the DataFrame bound to the writer before its input contract projects columns or casts values. The writer function may therefore see different columns or types. To check the values after coercion, capture them inside a replacement writer.

get_one_input raises LookupError if no input or more than one input matches the requested runtime type. It considers all bound inputs, including context values and defaults. When several match, use get_input(writer_step, "orders", pl.LazyFrame) to name the parameter.

In a test, compare the recorded input with an expected DataFrame using the backend's frame equality helper.

Assert the writer input in a test

This test compares the complete DataFrame passed to the writer:

from dataclasses import dataclass
from typing import Annotated

import polars as pl
from polars.testing import assert_frame_equal

from data_joinery import Pipeline, Project, transform


@dataclass
class Order:
    order_id: int


@transform
def read_orders() -> Annotated[pl.LazyFrame, Project(Order)]:
    return pl.LazyFrame({"order_id": [1, 2]})


@transform
def write_orders(orders: Annotated[pl.LazyFrame, Project(Order)]) -> None:
    pass


def test_writer_input():
    pipeline = Pipeline()
    read = pipeline.add_step(read_orders)
    write = pipeline.add_step(write_orders)
    read >> write

    result = pipeline.run()
    actual = result.get_one_input(write, pl.LazyFrame)
    expected = pl.LazyFrame({"order_id": [1, 2]})
    assert_frame_equal(actual.collect(), expected.collect())

Outputs:

============================= test session starts ==============================
platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/runner/work/data-joinery/data-joinery
configfile: pyproject.toml
plugins: platformdirs-4.12.1, layout-enforcer-0.1.0
collected 1 item

docs_src/learn/pipelines/inspecting_pipeline_runs/test_writer_input_polars.py . [100%]

============================== 1 passed ===============================

Read an intermediate DataFrame output

For a DataFrame step, use the same method. A saved Step gives the result its declared output type. When looking up a step by name, pass the backend's DataFrame class:

from dataclasses import dataclass
from typing import Annotated

import polars as pl

from data_joinery import Pipeline, Project, transform


@dataclass
class Order:
    amount: int


@dataclass
class OrderTotal:
    total: int


@transform
def read_orders() -> Annotated[pl.LazyFrame, Project(Order)]:
    return pl.LazyFrame({"amount": [10, 20]})


@transform
def total_orders(
    orders: Annotated[pl.LazyFrame, Project(Order)],
) -> Annotated[pl.LazyFrame, Project(OrderTotal)]:
    return orders.select(pl.col("amount").sum().alias("total"))


pipeline = Pipeline()
read_step = pipeline.add_step(read_orders)
total_step = pipeline.add_step(total_orders)
read_step >> total_step

result = pipeline.run()
print("Orders:", result.get_output(read_step).collect().to_dicts())
print("Total:", result.get_output("total_orders", pl.LazyFrame).collect().to_dicts())

Outputs:

Orders: [{'amount': 10}, {'amount': 20}]
Total: [{'total': 30}]

For PySpark, pass pyspark.sql.DataFrame in place of pl.LazyFrame. A named lookup raises TypeError if the output has a different runtime type.

The Pipeline result reference documents the methods and their errors.