CHIP: Feature Unit Testing Framework for Chronon (Phase 1)

Table of Contents

1. Problem Statement
2. Value Proposition
3. Requirements
4. Architecture
5. Bazel Integration
6. Detailed Design
7. User Experience
8. Verification
9. Future Phases

Problem Statement

Currently, Chronon users can define GroupBys, but there's no way to unit test them locally before deploying to production:

1. No local testing - Users must run full Spark jobs on production infrastructure to validate feature logic
2. No testing API - No Python functions to execute a GroupBy against mock data and verify results
3. No comparison utilities - No DataFrame diffing tools to compare actual vs expected output
4. Long feedback loops - Bugs are caught late in the deployment pipeline, often after expensive backfills

Current Workflow

Today, validating a GroupBy requires:

1. Define GroupBy in Python
2. Compile to Thrift JSON
3. Submit to production Spark cluster via run.py --mode backfill
4. Wait for job to complete (minutes to hours)
5. Manually query output table and inspect results
6. If wrong, repeat from step 1

There's no way to run a quick local test with mock data before committing.

Value Proposition

This CHIP introduces a comprehensive pytest-based testing framework that:

Enables Local Testing

Run GroupBy tests locally in Spark local mode without production infrastructure:

pytest tests/test_my_features.py -v

Supports All Accuracy Modes

Test both TEMPORAL and SNAPSHOT GroupBys with the same simple API:

result = test_group_by(my_group_by, input_df, query_df, spark)

Provides Clear Assertions

Get actionable error messages when tests fail:

AssertionError: GroupBy output mismatch
  Mismatched rows: 2
  | key  | actual_sum | expected_sum |
  | m1   | 350        | 300          |  # Diff: 50

Integrates with Any CI

Standard pytest fixtures work with GitHub Actions, CircleCI, Jenkins, etc.

Enables Agentic Development

Fast local feedback loops make this framework ideal for AI-assisted feature development. Agents can:

• Define a GroupBy, run tests, and iterate in seconds (not hours)
• Validate aggregation logic without production infrastructure
• Self-correct based on test failures with clear diff output

This lays the groundwork for more autonomous feature development workflows in Chronon.

Requirements

Goals

1. Support testing GroupBys with both TEMPORAL and SNAPSHOT accuracy
2. Provide pytest fixtures for easy SparkSession setup
3. Include DataFrame comparison utilities with clear diff output
4. Work with Spark local mode in CI environments
5. Bazel-native chronon_pyspark_test macro for hermetic test execution
6. Dual-mode support: works with both Bazel and sbt builds
7. Users provide their own mock DataFrames (no data generators)
8. Comprehensive documentation and examples

Non-Goals (Phase 1)

1. Join testing (future CHIP)
2. Derivations testing
3. Mock data generation utilities
4. Streaming/Flink testing
5. Online serving testing

Architecture

Module Structure

api/py/ai/chronon/repo/testing/
├── __init__.py           # Main API: test_group_by(), compare_dataframes()
├── assertions.py         # assert_group_by_output_equals(), etc.
├── comparison.py         # ComparisonResult class
├── fixtures.py           # pytest fixtures: chronon_spark
└── _bridge.py            # Internal Py4J bridge helpers

Python-JVM Bridge

Python Test                         Scala Execution
──────────────                      ───────────────
test_group_by()
    │
    ├─► thrift_simple_json(gb)  ────► PySparkUtils.parseGroupBy(json)
    │                                          │
    ├─► Detect accuracy mode                   │
    │                                          ▼
    │   TEMPORAL ─────────────────► GroupBy.temporalEvents(queryDf)
    │   SNAPSHOT ─────────────────► GroupBy.snapshotEvents(partitionRange)
    │                                          │
    ◄───────────────────────────────── DataFrame result

Accuracy Mode Execution Paths

Bazel Integration

PySpark tests require Chronon's Scala JARs on the classpath. This section defines a Bazel-native solution for hermetic, reproducible test execution.

The chronon_pyspark_test Macro

A custom Bazel macro that wraps py_test with automatic JAR dependency injection:

# tools/build_rules/python/chronon_pyspark_test.bzl

load("@rules_python//python:defs.bzl", "py_test")
load("@pypi//:requirements.bzl", "requirement")

def chronon_pyspark_test(name, srcs, deps = [], data = [], **kwargs):
    """
    Bazel macro for running PySpark tests with Chronon JARs.

    Automatically includes the Chronon spark-assembly JAR and configures
    the test environment for JAR discovery.

    Example:
        chronon_pyspark_test(
            name = "test_merchant_features",
            srcs = ["test_merchant_features.py"],
        )
    """

    # The spark uber JAR contains all Chronon Scala dependencies
    chronon_jar = "//spark:spark-assembly_deploy.jar"

    py_test(
        name = name,
        srcs = [
            "//tools/build_rules/python:pytest_wrapper.py",
        ] + srcs,
        main = "//tools/build_rules/python:pytest_wrapper.py",
        args = [
            "--capture=no",
            # Pass JAR path to pytest for fixture discovery
            "--chronon-jar=$(rootpath {})".format(chronon_jar),
        ] + ["$(location :%s)" % x for x in srcs],
        python_version = "PY3",
        deps = deps + [
            "//api/py:api_py",
            requirement("pytest"),
            requirement("pyspark"),
        ],
        data = data + [chronon_jar],
        **kwargs
    )

JAR Discovery

The fixture locates the JAR via --chronon-jar pytest arg (set by the macro) or falls back to sbt's target/ directories:

@pytest.fixture(scope="session")
def chronon_spark(request):
    jar_path = request.config.getoption("--chronon-jar", default=None)
    classpath = jar_path if jar_path else _discover_jars_from_sbt()

    spark = (
        SparkSession.builder
        .config("spark.driver.extraClassPath", classpath)
        .config("spark.executor.extraClassPath", classpath)
        .master("local[2]")
        .getOrCreate()
    )
    yield spark
    spark.stop()

Alternative: Use Bazel's runfiles library to locate JARs by label instead of pytest args.

Example BUILD.bazel Usage

# my_team/features/BUILD.bazel

load("//tools/build_rules/python:chronon_pyspark_test.bzl", "chronon_pyspark_test")

chronon_pyspark_test(
    name = "test_merchant_features",
    srcs = ["test_merchant_features.py"],
    data = [
        "test_data/purchases.json",  # Optional test fixtures
    ],
)

Run with:

bazel test //my_team/features:test_merchant_features

Detailed Design

Public API

test_group_by Signature

def test_group_by(
    group_by: GroupBy,
    input_df: DataFrame,
    query_df: Optional[DataFrame] = None,      # Required for TEMPORAL
    spark: SparkSession = None,
    mutations_df: Optional[DataFrame] = None,  # For TEMPORAL entities
    partition_range: Optional[tuple[str, str]] = None,  # Required for SNAPSHOT
) -> DataFrame

Automatically selects execution path based on GroupBy accuracy mode.

User Experience

Example Test

# tests/test_merchant_features.py

from ai.chronon.repo.testing import test_group_by, assert_group_by_output_equals
from ai.chronon.group_by import GroupBy, Aggregation, Operation, Window, TimeUnit, Accuracy

def test_sum_aggregation(chronon_spark):
    # Mock input data
    input_df = chronon_spark.createDataFrame([
        ("m1", 100, 1640995200000, "2022-01-01"),
        ("m1", 50,  1640998800000, "2022-01-01"),
        ("m2", 75,  1640995200000, "2022-01-01"),
    ], ["merchant_id", "amount", "ts", "ds"])

    query_df = chronon_spark.createDataFrame([
        ("m1", 1641006000000, "2022-01-01"),
        ("m2", 1641006000000, "2022-01-01"),
    ], ["merchant_id", "ts", "ds"])

    # Define and execute GroupBy
    gb = GroupBy(
        keys=["merchant_id"],
        aggregations=[Aggregation("amount", Operation.SUM, [Window(1, TimeUnit.DAYS)])],
        accuracy=Accuracy.TEMPORAL,
    )
    result = test_group_by(gb, input_df, query_df, chronon_spark)

    # Verify
    expected = chronon_spark.createDataFrame([
        ("m1", 150), ("m2", 75)
    ], ["merchant_id", "amount_sum_1d"])

    assert_group_by_output_equals(result, expected, keys=["merchant_id"])

For SNAPSHOT accuracy, use partition_range=("2022-01-01", "2022-01-02") instead of query_df.

CI Integration

# .github/workflows/test-features.yml (Bazel)
steps:
  - uses: actions/checkout@v3
  - uses: bazelbuild/setup-bazelisk@v2
  - run: bazel test //my_team/features:test_merchant_features

Bazel automatically builds the JAR (if not cached) and injects it via --chronon-jar.

sbt alternative: Build JARs with sbt '++ 2.12.12 spark/assembly', then run pytest tests/ -v.

Verification

# Bazel 
bazel test //api/py:chronon_pyspark_test

# sbt alternative
sbt '++ 2.12.12 spark/assembly' && pytest api/py/test/ -v

Future Phases

Conclusion

This CHIP introduces a testing framework that enables Chronon users to validate their GroupBy definitions locally before production deployment. By supporting all accuracy modes and providing clear assertion helpers, the framework significantly reduces the feedback loop for feature development while integrating seamlessly with existing CI/CD workflows.