zipline-pyspark 0.1.1

Notebook execution utilities for Zipline AI

zipline-pyspark

Execution utilities to help with iteration and UX when working with Zipline in notebook environments.

Status: Work in progress.

Overview

zipline-pyspark provides thin wrappers around Zipline's core entities — GroupBy, Join, and StagingQuery — that are designed to run interactively inside a Jupyter (or compatible) notebook against a live PySpark session. The goal is to make it easy to prototype, debug, and iterate on Zipline feature definitions without leaving the notebook.

Installation

%pip install zipline-pyspark

Usage

from ai.chronon.pyspark.jupyter import JupyterGroupBy, JupyterJoin, JupyterStagingQuery

# Run a GroupBy over a date range, chunked by day
result_df = JupyterGroupBy(my_group_by, spark).run(
    start_date="2026-04-01",
    end_date="2026-04-07",
    step_days=1,
)

# Run a Join
result_df = JupyterJoin(my_join, spark).run(
    start_date="2026-04-01",
    end_date="2026-04-07",
)

# Run a StagingQuery
result_df = JupyterStagingQuery(my_staging_query, spark).run(
    end_date="2026-04-07",
    start_date="2026-04-01",
    step_days=1,
)

Setup

Prerequisites

• Python 3.11+
• PySpark 3.5+
• A running Spark session (local or remote)
• zipline-ai package (installed automatically as a dependency)

Databricks limitations

The Databricks notebook executors invoke Chronon's Scala batch driver through the Spark driver JVM. Databricks serverless compute does not support direct access to the underlying driver JVM via sparkContext, so these executors are not compatible with serverless notebook compute. Use classic Databricks compute with Dedicated access mode, and attach the Chronon batch JAR to the cluster.

Development install

git clone https://github.com/zipline-ai/zipline-pyspark.git
cd zipline-pyspark
pip install -e ".[dev]"
pre-commit install

Running tests

pytest tests/ -v

Coverage is enforced automatically (configured in pyproject.toml). To skip the coverage check locally:

pytest tests/ -v --no-cov

Release Flow

Each PR must carry exactly one semantic version label: Semver-Major, Semver-Minor, or Semver-Patch. Release Drafter uses those labels to keep a draft GitHub release up to date and determine the next release version.

Publishing that GitHub release builds the zipline-pyspark wheel with the release tag as the package version and publishes it to PyPI through trusted publishing. To smoke-test the same wheel build locally:

scripts/build_release_wheel.sh v0.2.0

Roadmap

1. Implement jupyter and databricks backends

• jupyter.py — complete JupyterGroupBy and JupyterJoin to match the existing JupyterStagingQuery pattern: date-range chunking, setup-statement execution, and union of per-step DataFrames.
• databricks.py — a parallel set of classes (DatabricksGroupBy, DatabricksJoin, DatabricksStagingQuery) that adapt the same interface for Databricks notebooks: dbutils-aware progress display, widget-based date inputs, and Databricks displayHTML / display integration instead of plain DataFrame returns.
• Shared base layer — extract common chunking and template-rendering logic into an internal _base.py so both backends stay in sync without code duplication.

2. Distribution

• Publish to PyPI under zipline-pyspark and document the %pip install zipline-pyspark notebook workflow.

Credits

This project is the Zipline counterpart of the original implementation contributed to the open-source Chronon project by Airbnb. See the original pull request for reference: airbnb/chronon#981.

License

See LICENSE.