Testing Polars Applications: Unit Tests to Integration
Data engineering pipelines built with Polars are fast, expressive, and increasingly popular — but speed means nothing if your transformations silently produce wrong results. Testing Polars applications ensures that your data transformations, aggregations, and joins behave as expected across code changes, schema drift, and growing datasets. This tutorial walks you through a complete testing strategy, from isolated unit tests on expressions to full integration tests that validate end-to-end pipelines.
Why Testing Polars Applications Matters
Polars encourages a functional, declarative style of data manipulation. While this is cleaner than imperative pandas code, it introduces its own failure modes:
- Silent schema changes: a column rename upstream can break a downstream expression without raising an error until runtime.
- Type coercion surprises: Polars is strict, but operations like joins on mismatched dtypes can still produce unexpected results.
- Lazy optimization edge cases: the lazy engine reorders and optimizes operations, which can expose bugs that don't appear in eager mode.
- Scaling behavior: code that works on a 100-row sample may fail on partitioned data or produce different results due to floating point aggregation order.
A robust test suite catches these issues early, documents expected behavior, and gives you confidence to refactor aggressively.
Structuring Your Test Strategy
A good testing strategy for Polars applications follows the classic pyramid:
- Unit tests: validate individual expressions, custom functions, and small transformations in isolation.
- Component tests: validate a logical group of transformations, such as a single cleaning stage.
- Integration tests: validate the full pipeline end-to-end, including I/O and schema contracts.
- Property-based tests: validate invariants that should hold for any valid input.
Setting Up the Project
Assume the following project layout:
polars_app/
├── polars_app/
│ ├── __init__.py
│ ├── transforms.py
│ └── pipeline.py
├── tests/
│ ├── __init__.py
│ ├── conftest.py
│ ├── test_transforms.py
│ ├── test_pipeline.py
│ └── test_properties.py
├── pyproject.toml
└── data/
├── sample_input.csv
└── expected_output.csv
Install the required dependencies:
pip install polars pytest pytest-cov hypothesis
Unit Testing Polars Expressions
Unit tests focus on the smallest testable pieces of logic. In a Polars application, that usually means individual transformation functions that accept and return DataFrames or LazyFrames.
The Transformation Module
Here is a sample transforms.py module containing a few reusable transformation functions:
import polars as pl
def normalize_amounts(df: pl.DataFrame) -> pl.DataFrame:
return df.with_columns(
(pl.col("amount") - pl.col("amount").mean())
/ pl.col("amount").std()
.alias("amount_zscore")
)
def filter_active_users(df: pl.DataFrame) -> pl.DataFrame:
return df.filter(pl.col("status") == "active")
def aggregate_by_user(df: pl.DataFrame) -> pl.DataFrame:
return (
df.group_by("user_id")
.agg(
pl.col("amount").sum().alias("total_amount"),
pl.col("amount").count().alias("transaction_count"),
pl.col("amount").mean().alias("avg_amount"),
)
.sort("user_id")
)
Writing Unit Tests
Each function gets its own test that constructs a tiny, deterministic input DataFrame and asserts on the output. Use pl.testing helpers and Polars' built-in equality for clean assertions.
import polars as pl
from polars_app.transforms import (
normalize_amounts,
filter_active_users,
aggregate_by_user,
)
def test_filter_active_users_keeps_only_active():
df = pl.DataFrame({
"user_id": [1, 2, 3],
"status": ["active", "inactive", "active"],
"amount": [10.0, 20.0, 30.0],
})
result = filter_active_users(df)
assert result.height == 2
assert result["user_id"].to_list() == [1, 3]
def test_aggregate_by_user_sums_correctly():
df = pl.DataFrame({
"user_id": [1, 1, 2, 2, 2],
"amount": [10.0, 20.0, 5.0, 5.0, 10.0],
})
result = aggregate_by_user(df)
assert result.height == 2
assert result["total_amount"].to_list() == [30.0, 20.0]
assert result["transaction_count"].to_list() == [2, 3]
def test_normalize_amounts_produces_zero_mean():
df = pl.DataFrame({"amount": [1.0, 2.0, 3.0, 4.0]})
result = normalize_amounts(df)
zscores = result["amount_zscore"].to_list()
assert round(sum(zscores) / len(zscores), 10) == 0.0
Using Fixtures for Reusable Data
To avoid repeating DataFrame construction across tests, define fixtures in conftest.py:
import polars as pl
import pytest
@pytest.fixture
def sample_transactions():
return pl.DataFrame({
"user_id": [1, 1, 2, 3, 3, 3],
"status": ["active", "active", "inactive", "active", "active", "active"],
"amount": [100.0, 200.0, 50.0, 10.0, 20.0, 30.0],
"timestamp": [
"2024-01-01",
"2024-01-02",
"2024-01-03",
"2024-01-04",
"2024-01-05",
"2024-01-06",
],
}).with_columns(pl.col("timestamp").str.to_date())
Now tests can request the fixture by name:
def test_filter_active_users_with_fixture(sample_transactions):
result = filter_active_users(sample_transactions)
assert "inactive" not in result["status"].to_list()
Testing LazyFrames
Polars' lazy API defers execution and applies query optimization. To test lazy pipelines, you should test the optimized plan and the materialized result separately.
import polars as pl
from polars_app.pipeline import build_lazy_pipeline
def test_lazy_pipeline_optimizes_predicates():
source = pl.scan_csv("data/sample_input.csv")
lazy = build_lazy_pipeline(source)
plan = lazy.explain(optimized=True)
# The optimizer should push the filter below the projection
assert "FILTER" in plan or "SELECTION" in plan
def test_lazy_pipeline_collects_expected_schema():
source = pl.scan_csv("data/sample_input.csv")
lazy = build_lazy_pipeline(source)
schema = lazy.collect_schema()
assert "total_amount" in schema
assert schema["total_amount"] == pl.Float64
Calling collect_schema() is a cheap way to verify that your lazy pipeline produces the expected columns and types without executing the full query.
Component Tests for Multi-Step Stages
Component tests validate a coherent stage of the pipeline, such as a cleaning stage that combines several transformations. Suppose your pipeline.py contains:
import polars as pl
from polars_app.transforms import filter_active_users, aggregate_by_user
def clean_stage(df: pl.DataFrame) -> pl.DataFrame:
return (
df.lazy()
.pipe(filter_active_users)
.pipe(aggregate_by_user)
.collect()
)
def build_lazy_pipeline(source: pl.LazyFrame) -> pl.LazyFrame:
return (
source
.filter(pl.col("status") == "active")
.group_by("user_id")
.agg(
pl.col("amount").sum().alias("total_amount"),
pl.col("amount").count().alias("transaction_count"),
)
.sort("user_id")
)
A component test verifies the combined behavior:
import polars as pl
from polars_app.pipeline import clean_stage
def test_clean_stage_produces_aggregated_output(sample_transactions):
result = clean_stage(sample_transactions)
assert result.height == 2 # users 1 and 3 are active
assert result.columns == ["user_id", "total_amount", "transaction_count"]
assert result["total_amount"].sum() == 360.0
Integration Testing the Full Pipeline
Integration tests run the entire pipeline against realistic inputs and compare outputs to known-good results. The most maintainable approach is to store fixtures as files and compare against expected output files.
Snapshot-Based Integration Test
import polars as pl
from polars_app.pipeline import build_lazy_pipeline
def test_pipeline_end_to_end_matches_expected():
source = pl.scan_csv("data/sample_input.csv")
result = build_lazy_pipeline(source).collect()
expected = pl.read_csv("data/expected_output.csv")
assert result.columns == expected.columns
assert result.schema == expected.schema
assert result.equals(expected)
The equals() method performs a strict comparison including dtypes and row order. If row order should not matter, sort both frames first or use frame_equal with appropriate options.
Testing I/O Boundaries
Integration tests should also exercise the I/O layer. If your pipeline reads from CSV and writes to Parquet, test both directions:
import polars as pl
from pathlib import Path
from polars_app.pipeline import run_pipeline
def test_pipeline_writes_parquet(tmp_path):
output_path = tmp_path / "result.parquet"
run_pipeline(
input_path="data/sample_input.csv",
output_path=output_path,
)
assert output_path.exists()
written = pl.read_parquet(output_path)
assert written.height > 0
assert "total_amount" in written.columns
Using pytest's tmp_path fixture keeps your test workspace clean and avoids polluting the repository with generated files.
Property-Based Testing with Hypothesis
Property-based testing is especially powerful for data transformations because you can assert invariants that must hold for any valid input. The hypothesis library integrates cleanly with Polars through custom strategies.
Defining a DataFrame Strategy
import polars as pl
from hypothesis import given, strategies as st, settings
from polars_app.transforms import aggregate_by_user
@st.composite
def transactions(draw):
n = draw(st.integers(min_value=1, max_value=50))
user_ids = draw(
st.lists(st.integers(min_value=1, max_value=10), min_size=n, max_size=n)
)
amounts = draw(
st.lists(
st.floats(min_value=0.01, max_value=1000.0, allow_nan=False),
min_size=n,
max_size=n,
)
)
return pl.DataFrame({"user_id": user_ids, "amount": amounts})
@given(transactions())
@settings(max_examples=50)
def test_aggregate_total_equals_sum_of_input(df):
result = aggregate_by_user(df)
input_total = df["amount"].sum()
output_total = result["total_amount"].sum()
assert abs(output_total - input_total) < 1e-6
This test verifies a conservation invariant: the sum of all amounts must be preserved through aggregation. Such properties catch subtle bugs like accidental filtering or double counting.
Testing Schema Invariants
@given(transactions())
def test_aggregate_always_returns_expected_columns(df):
result = aggregate_by_user(df)
assert set(result.columns) == {"user_id", "total_amount", "transaction_count", "avg_amount"}
assert result.schema["user_id"] == pl.Int64
assert result.schema["total_amount"] == pl.Float64
Best Practices for Testing Polars Applications
- Test schemas explicitly: assert on
df.schemaorcollect_schema()to catch silent type changes early. - Prefer small, deterministic fixtures: tiny DataFrames make failures easy to diagnose and tests fast to run.
- Test lazy and eager paths separately: the optimizer can change behavior, so verify both the plan and the collected result.
- Use
tmp_pathfor file I/O tests: never write test outputs into the repository tree. - Sort before comparing: group_by operations do not guarantee row order, so sort both frames on a stable key before equality checks.
- Use
approxfor floats: floating point aggregation order can vary, so compare with tolerance rather than exact equality. - Snapshot expected outputs: store known-good outputs as files and regenerate them deliberately, not automatically on every run.
- Cover edge cases: empty DataFrames, single-row frames, nulls, and duplicate keys often expose hidden bugs.
- Property-test invariants: conservation laws, monotonicity, and idempotence are perfect candidates for Hypothesis.
Testing Edge Cases
import polars as pl
from polars_app.transforms import aggregate_by_user
def test_aggregate_handles_empty_dataframe():
df = pl.DataFrame({"user_id": pl.Series([], dtype=pl.Int64),
"amount": pl.Series([], dtype=pl.Float64)})
result = aggregate_by_user(df)
assert result.height == 0
assert result.columns == ["user_id", "total_amount", "transaction_count", "avg_amount"]
def test_aggregate_handles_nulls():
df = pl.DataFrame({
"user_id": [1, 1, None],
"amount": [10.0, None, 30.0],
})
result = aggregate_by_user(df)
# Null user_id should be grouped separately; null amounts ignored by sum
assert result.height >= 1
user_1 = result.filter(pl.col("user_id") == 1)
assert user_1["total_amount"].item() == 10.0
Measuring and Maintaining Coverage
Use pytest-cov to track how much of your transformation code is exercised by tests:
pytest --cov=polars_app --cov-report=term-missing
Aim for high coverage on transformation logic, but do not obsess over I/O wrappers or configuration loaders. Focus coverage on the expressions and aggregations where bugs hide.
Conclusion
Testing Polars applications requires a layered approach that respects the unique characteristics of declarative data transformations. Unit tests pin down individual expressions, component tests validate coherent stages, integration tests confirm end-to-end correctness against known-good outputs, and property-based tests enforce invariants that hold across the entire input space. By combining deterministic fixtures, explicit schema assertions, lazy plan inspection, and Hypothesis strategies, you build a safety net that lets you refactor aggressively and scale confidently. The upfront investment in a thorough test suite pays dividends every time your data sources evolve, your schemas shift, or your pipeline grows in complexity — ensuring that speed of execution is matched by reliability of results.