Testing

Polars DataFrame Testing in Python and Rust: A Complete Guide

HelpMeTest

HelpMeTest

— 5 min read
Polars DataFrame Testing in Python and Rust: A Complete Guide

Polars has become a serious alternative to pandas for data processing work. Its Rust core delivers performance that pandas cannot match, and its expressive lazy API makes complex transformations readable. But testing Polars code requires different patterns than testing ordinary Python or Rust functions. DataFrames are not simple values — they have schemas, row counts, column types, and floating-point precision concerns. This guide covers practical testing patterns for Polars in both Python and Rust.

Why DataFrame Testing Is Different

When you test a function that returns an integer, you compare two integers. When you test a function that returns a DataFrame, you need to verify:

  • The schema (column names and types)
  • The row count
  • The values in specific columns
  • The order of rows (if order matters)
  • Null handling
  • Floating-point approximate equality

Polars provides built-in tools for all of these, but you need to know where to look.

Python: Testing with pytest and Polars Assertions

Setting Up

Install the testing dependencies:

pip install polars pytest

Polars ships with polars.testing, a module specifically designed for DataFrame assertions. Import it in your test files:

import polars as pl
from polars.testing import assert_frame_equal, assert_series_equal

Basic DataFrame Assertion

Suppose you have a data cleaning function:

# src/transforms.py
import polars as pl

def normalize_names(df: pl.DataFrame) -> pl.DataFrame:
    return df.with_columns([
        pl.col("first_name").str.strip_chars().str.to_titlecase(),
        pl.col("last_name").str.strip_chars().str.to_titlecase(),
    ])

The test:

# tests/test_transforms.py
import polars as pl
from polars.testing import assert_frame_equal
from src.transforms import normalize_names

def test_normalize_names_strips_whitespace():
    input_df = pl.DataFrame({
        "first_name": ["  alice", "BOB  ", " carol "],
        "last_name": ["smith", "JONES", " davis "],
    })
    expected = pl.DataFrame({
        "first_name": ["Alice", "Bob", "Carol"],
        "last_name": ["Smith", "Jones", "Davis"],
    })
    result = normalize_names(input_df)
    assert_frame_equal(result, expected)

assert_frame_equal checks schema, shape, and values. On failure it produces a diff that shows exactly which rows and columns diverged.

Testing with Approximate Equality

Floating-point results from aggregations or mathematical transformations should use check_exact=False:

def test_revenue_per_user():
    input_df = pl.DataFrame({
        "user_id": [1, 1, 2, 2],
        "revenue": [10.5, 20.0, 5.25, 15.75],
    })
    result = (
        input_df
        .group_by("user_id")
        .agg(pl.col("revenue").sum().alias("total_revenue"))
        .sort("user_id")
    )
    expected = pl.DataFrame({
        "user_id": [1, 2],
        "total_revenue": [30.5, 21.0],
    })
    assert_frame_equal(result, expected, check_exact=False, atol=1e-6)

The atol parameter sets the absolute tolerance for float comparison.

Schema Validation Tests

Test that your transformations produce the expected schema before testing values:

def test_output_schema():
    df = pl.DataFrame({"value": [1, 2, 3]})
    result = my_transform(df)
    expected_schema = {
        "value": pl.Int64,
        "value_squared": pl.Int64,
        "value_log": pl.Float64,
    }
    assert result.schema == expected_schema

Schema failures surface integration bugs early — for example, when an upstream data source changes a column type from Int32 to Int64.

Testing Lazy Frames

Polars lazy evaluation defers work until you call .collect(). Test lazy pipelines the same way, just collect before asserting:

def test_lazy_pipeline():
    lf = pl.LazyFrame({
        "category": ["A", "B", "A", "B"],
        "amount": [100, 200, 150, 250],
    })
    result = (
        lf
        .filter(pl.col("amount") > 100)
        .group_by("category")
        .agg(pl.col("amount").sum())
        .sort("category")
        .collect()
    )
    expected = pl.DataFrame({
        "category": ["A", "B"],
        "amount": [150, 450],
    })
    assert_frame_equal(result, expected)

Testing Null Handling

Null handling is a common source of bugs. Write explicit tests for null behavior:

def test_fill_nulls():
    df = pl.DataFrame({
        "score": [1.0, None, 3.0, None, 5.0],
    })
    result = df.with_columns(
        pl.col("score").fill_null(strategy="mean").alias("score_filled")
    )
    # mean of [1.0, 3.0, 5.0] = 3.0
    assert result["score_filled"].null_count() == 0
    assert result["score_filled"][1] == 3.0
    assert result["score_filled"][3] == 3.0

Rust: Testing Polars DataFrames

The Rust API is lower-level but the patterns are similar. Polars for Rust uses polars-core and provides its own testing utilities.

Setting Up

In Cargo.toml:

[dependencies]
polars = { version = "0.39", features = ["lazy"] }

[dev-dependencies]
polars = { version = "0.39", features = ["lazy", "testing"] }

Basic Rust Test

#[cfg(test)]
mod tests {
    use polars::prelude::*;

    fn make_sales_df() -> DataFrame {
        df![
            "product" => ["A", "B", "A", "C"],
            "units" => [10i32, 5, 8, 3],
            "price" => [2.5f64, 4.0, 2.5, 10.0],
        ]
        .unwrap()
    }

    #[test]
    fn test_total_revenue() {
        let df = make_sales_df();
        let result = df
            .lazy()
            .with_column((col("units").cast(DataType::Float64) * col("price"))
                .alias("revenue"))
            .collect()
            .unwrap();

        let revenue = result.column("revenue").unwrap();
        let expected = Series::new("revenue", &[25.0f64, 20.0, 20.0, 30.0]);
        assert!(revenue.equals(&expected));
    }

    #[test]
    fn test_schema_has_revenue_column() {
        let df = make_sales_df();
        let result = df
            .lazy()
            .with_column((col("units").cast(DataType::Float64) * col("price"))
                .alias("revenue"))
            .collect()
            .unwrap();

        assert!(result.schema().contains("revenue"));
        assert_eq!(
            result.schema().get("revenue").unwrap(),
            &DataType::Float64
        );
    }
}

Comparing DataFrames in Rust

For full DataFrame equality, iterate over columns:

fn assert_dataframes_equal(actual: &DataFrame, expected: &DataFrame) {
    assert_eq!(actual.shape(), expected.shape(), "Shape mismatch");
    for name in expected.get_column_names() {
        let actual_col = actual.column(name).expect("Column missing in actual");
        let expected_col = expected.column(name).unwrap();
        assert!(
            actual_col.equals(expected_col),
            "Column '{}' differs: actual={:?}, expected={:?}",
            name, actual_col, expected_col
        );
    }
}

Use this helper across your test suite instead of repeating the assertion logic.

CI Integration

Both Python and Rust Polars tests integrate cleanly into GitHub Actions:

name: Polars Tests

on: [push, pull_request]

jobs:
  python-tests:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: "3.12"
      - run: pip install polars pytest
      - run: pytest tests/

  rust-tests:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: dtolnay/rust-toolchain@stable
      - uses: Swatinem/rust-cache@v2
      - run: cargo test

For Rust, Swatinem/rust-cache is essential — it caches the compiled dependencies including the Polars codebase, cutting CI time from 10+ minutes to under 2.

Testing Data Pipelines End-to-End

Unit tests on individual transformations are not enough for a real data pipeline. You also need integration tests that run the pipeline on a representative dataset and verify the shape and statistical properties of the output:

def test_pipeline_output_shape(sample_dataset_path):
    df = pl.read_parquet(sample_dataset_path)
    result = run_full_pipeline(df)
    # Shape assertions
    assert result.shape[1] == 12, "Expected 12 output columns"
    assert result.shape[0] > 0, "Pipeline produced no rows"
    # Null rate assertions
    assert result["customer_id"].null_count() == 0
    assert result["revenue"].null_count() / len(result) < 0.01

For testing the deployed version of your pipeline — especially when it runs as an API or scheduled job — HelpMeTest can run automated checks against your endpoints on a schedule. The Pro plan at $100/month covers continuous monitoring for data pipelines that expose HTTP interfaces, so regressions get caught before they affect downstream consumers.

Common Pitfalls

Column ordering. assert_frame_equal checks column order by default. Use check_column_order=False if your transformation may reorder columns as a side effect.

Categorical vs String. Polars uses a Categorical type for low-cardinality string columns. A column that comes in as Utf8 (String) may be cast to Categorical by a join or group_by. Verify types explicitly in schema tests.

Integer overflow. Polars defaults to Int64 in Python but Int32 in many scan operations. Arithmetic on mixed types can silently upcast or fail. Always specify dtype in DataFrame constructors used for testing.

Group-by order. Group-by results are unordered in Polars. Always call .sort() before comparing with assert_frame_equal unless you use check_row_order=False.

Summary

Testing Polars code is more involved than testing pure functions, but Polars gives you the tools to do it right. Use assert_frame_equal for Python and column-level equals for Rust. Write dedicated schema tests. Test null handling explicitly. Sort before comparing group-by results. Add CI from the beginning so your data transformations are as well-tested as your application code.

Read more