awx-zipline-ai 0.2.0__py3-none-any.whl

ai/chronon/resources/gcp/README.md

@@ -0,0 +1,174 @@
+
+# 🧠 Zipline AI: Sample Chronon Project
+
+This repository demonstrates how to author and run [Chronon](https://chronon.ai) pipelines, including GroupBy and Join definitions, using GCP (BigQuery + Iceberg) as the storage backend.
+
+Chronon is a unified platform for **feature engineering**, enabling **online and offline consistency**, **real-time feature generation**, and **historical backfills** from a single codebase.
+
+---
+
+## 📦 Project Structure
+
+```bash
+.
+├── group_bys/           # GroupBy definitions (feature aggregations)
+├── joins/               # Join definitions (how sources and GroupBys are combined)
+├── sources/             # Chronon Source definitions (event tables)
+├── compiled/            # Generated configs and outputs
+├── teams.py             # Chronon Team configurations
+└── README.md
+```
+
+---
+
+## 🚀 Quick Start
+
+### 🛠️ Requirements
+
+To get started, make sure you have the following set up:
+
+- ✅ **Python** 3.11 or higher
+- ✅ **Zipline CLI** — Only for **upgrades or downgrades**, install via:
+  ```bash
+  ./zipline-cli-install.sh
+- ✅ gcloud CLI — authenticated and configured with the correct GCP project
+- ✅ Google Cloud credentials — either:
+  - Application Default Credentials (ADC)
+  - A service account with access to BigQuery and GCS
+- ✅Add this to your shell config (e.g., .bashrc, .zshrc):
+
+```bash
+# From the same directory as this README
+export PYTHONPATH="$(pwd):$PYTHONPATH"
+```
+
+---
+## Requirements
+
+Teams define metadata, Spark config, and environment variables.
+
+In [teams.py](teams.py), fill in the appropriate values in the TODO section.
+
+Make sure to replace placeholders like `<project-id>` and `<gcs-prefix>` with real values.
+
+### Partition format and column
+Chronon expects tables to be date partitioned. Please specify the partition format and the column in teams.py here:
+
+```python
+            "spark.chronon.partition.format": "<date-format>", # ex: "yyyy-MM-dd",
+            "spark.chronon.partition.column": "<partition-column-name>", # ex: "ds",
+```
+
+---
+
+## 🧪 Compiling
+
+To generate the user configs from the Python chronon objects to be used in the CLI, run:
+
+```bash
+zipline compile
+```
+
+This will create a `compiled` directory.
+
+---
+
+## 🧪 Running a GroupBy or Join Backfill
+
+Run a GroupBy backfill from the CLI:
+
+```bash
+zipline run \
+--mode backfill \
+--conf compiled/group_bys/<TEAM_NAME>/<GROUPBY_NAME>
+```
+
+Run a Join backfill from the CLI:
+
+```bash
+zipline run \
+--mode backfill \
+--conf compiled/joins/<TEAM_NAME>/<JOIN_NAME>
+```
+
+Results are written to the configured BigQuery + Iceberg tables under the `outputNamespace` (e.g. `default.group_by_v1` or `default.v1`).
+
+---
+
+## 🧪 Running a GroupBy upload (GBU) job.
+
+```bash
+zipline run \
+--mode upload \
+--conf compiled/group_bys/<TEAM_NAME>/<GROUP_BY_NAME> \
+--ds <DATE>
+```
+
+Results are written to the configured BigQuery + Iceberg tables under the `outputNamespace` (e.g. `default.group_by_v1` or `default.v1`).
+
+---
+
+## 🧪 Upload the GBU values to online KV store.
+
+```bash
+zipline run \
+--mode upload-to-kv \
+--conf compiled/group_bys/<TEAM_NAME>/<GROUP_BY_NAME> \
+--ds <DATE>
+```
+
+---
+
+## 🧪 Upload the metadata of Chronon GroupBy or Join to online KV store for serving.
+
+GroupBy metadata upload:
+```bash
+zipline run \
+--mode metadata-upload \
+--conf compiled/group_bys/<TEAM_NAME>/<GROUP_BY_NAME>
+```
+
+Join metadata upload:
+```bash
+zipline run \
+--mode metadata-upload \
+--conf compiled/joins/<TEAM_NAME>/<JOIN_NAME>
+```
+
+---
+
+## 🧪 Fetch feature values from Chronon GroupBy or Join.
+
+**Note:** This is only for debugging purposes. Not for production use.
+
+Fetching from a GroupBy:
+```bash
+zipline run \
+--mode fetch \
+--conf compiled/group_bys/<TEAM_NAME>/<GROUP_BY_NAME> \
+--name <GROUP_BY_NAME> \
+-k '{"<ENTITY_KEY>": "<VALUE>"}'
+```
+
+Fetching from a Join:
+```bash
+zipline run \
+--mode fetch \
+--conf compiled/joins/<TEAM_NAME>/<JOIN_NAME> \
+--name <JOIN_NAME> \
+-k '{"<ENTITY_KEY>": "<VALUE>"}'
+```
+
+---
+
+## 📚 Resources
+
+- [Chronon Docs](https://chronon.ai)
+- [GitHub](https://github.com/airbnb/chronon)
+- [Community Slack](https://join.slack.com/t/chrononworkspace/shared_invite/zt-33zbnzwac-ghPZXpYNZJsArXZ5WdBy9g)
+
+---
+
+## 👋 About
+
+This project is a reference scaffold for building scalable feature pipelines using Chronon on GCP. It provides end-to-end visibility from source to production features.

ai/chronon/resources/gcp/group_bys/test/data.py

@@ -0,0 +1,34 @@
+
+from sources.test.data import source_v1
+
+from ai.chronon.group_by import Aggregation, GroupBy, Operation, TimeUnit, Window
+
+window_sizes = [Window(length=day, time_unit=TimeUnit.DAYS) for day in [3, 14, 30]] # Define some window sizes to use below
+
+group_by_v1 = GroupBy(
+    backfill_start_date="2023-11-01",
+    sources=[source_v1],
+    keys=["user_id"], # We are aggregating by user
+    online=True,
+    aggregations=[Aggregation(
+        input_column="purchase_price",
+        operation=Operation.SUM,
+        windows=window_sizes
+    ), # The sum of purchases prices in various windows
+        Aggregation(
+            input_column="purchase_price",
+            operation=Operation.COUNT,
+            windows=window_sizes
+        ), # The count of purchases in various windows
+        Aggregation(
+            input_column="purchase_price",
+            operation=Operation.AVERAGE,
+            windows=window_sizes
+        ), # The average purchases by user in various windows
+        Aggregation(
+            input_column="purchase_price",
+            operation=Operation.LAST_K(10),
+        ),
+    ],
+    version=0,
+)

ai/chronon/resources/gcp/joins/test/data.py

@@ -0,0 +1,30 @@
+from group_bys.test.data import group_by_v1
+
+from ai.chronon.api.ttypes import EventSource, Source
+from ai.chronon.join import Join, JoinPart
+from ai.chronon.query import Query, selects
+
+"""
+This is the "left side" of the join that will comprise our training set. It is responsible for providing the primary keys
+and timestamps for which features will be computed.
+"""
+source = Source(
+    events=EventSource(
+        table="data.checkouts",
+        query=Query(
+            selects=selects(
+                "user_id"
+            ),  # The primary key used to join various GroupBys together
+            time_column="ts",
+        ),  # The event time used to compute feature values as-of
+    )
+)
+
+v1 = Join(
+    left=source,
+    right_parts=[
+        JoinPart(group_by=group_by_v1)
+    ],
+    row_ids="user_id",
+    version=0,
+)

ai/chronon/resources/gcp/sources/test/data.py

@@ -0,0 +1,23 @@
+from ai.chronon.api.ttypes import EventSource, Source
+from ai.chronon.query import Query, selects
+
+"""
+Example: Defining a Chronon Source from a Batch Table
+
+This example demonstrates how to configure a Chronon `Source` from a BigQuery or Hive table,
+with a clear event time column and selected fields for downstream feature computation.
+"""
+
+# Define the EventSource using the batch table and query
+# Wrap the EventSource in a Source object
+
+source_v1 = Source(
+    events=EventSource(
+        table="data.purchases", # This points to the log table in the warehouse with historical purchase events, updated in batch daily
+        topic=None, # See the 'returns' GroupBy for an example that has a streaming source configured. In this case, this would be the streaming source topic that can be listened to for realtime events
+        query=Query(
+            selects=selects("user_id","purchase_price"), # Select the fields we care about
+            time_column="ts") # The event time
+    ))
+
+# The `source_v1` object can now be used in a Chronon join or pipeline definition

ai/chronon/resources/gcp/teams.py

@@ -0,0 +1,70 @@
+from ai.chronon.api.ttypes import Team
+from ai.chronon.repo.constants import RunMode
+from ai.chronon.types import ConfigProperties, EnvironmentVariables
+
+default = Team(
+    description="Default team",
+    email="<responsible-team-email>",
+    outputNamespace="default",
+    conf=ConfigProperties(
+        common={
+            "spark.chronon.table.format_provider.class": "ai.chronon.integrations.cloud_gcp.GcpFormatProvider",
+            "spark.chronon.table_write.format": "iceberg",
+
+            "spark.sql.defaultCatalog": "bigquery_catalog",
+
+            "spark.sql.catalog.bigquery_catalog": "ai.chronon.integrations.cloud_gcp.DelegatingBigQueryMetastoreCatalog",
+            "spark.sql.catalog.bigquery_catalog.catalog-impl": "org.apache.iceberg.gcp.bigquery.BigQueryMetastoreCatalog",
+            "spark.sql.catalog.bigquery_catalog.io-impl": "org.apache.iceberg.io.ResolvingFileIO",
+
+            "spark.sql.defaultUrlStreamHandlerFactory.enabled": "false",
+            "spark.kryo.registrator": "ai.chronon.integrations.cloud_gcp.ChrononIcebergKryoRegistrator",
+
+            "spark.chronon.coalesce.factor": "10",
+            "spark.default.parallelism": "10",
+            "spark.sql.shuffle.partitions": "10",
+
+            # TODO: Please fill in the following values
+            "spark.sql.catalog.bigquery_catalog.warehouse": "gs://zipline-warehouse-<customer_id>/data/tables/",
+            "spark.sql.catalog.bigquery_catalog.gcp_location": "<region>",
+            "spark.sql.catalog.bigquery_catalog.gcp_project": "<project-id>",
+            "spark.chronon.partition.format": "<date-format>", # ex: "yyyy-MM-dd",
+            "spark.chronon.partition.column": "<partition-column-name>", # ex: "ds",
+        },
+    ),
+    env=EnvironmentVariables(
+        common={
+            "JOB_MODE": "local[*]",
+            "CHRONON_ONLINE_CLASS": "[ONLINE-TODO]your.online.class",
+            "CHRONON_ONLINE_ARGS": "[ONLINE-TODO]args prefixed with -Z become constructor map for your implementation of ai.chronon.online.Api, -Zkv-host=<YOUR_HOST> -Zkv-port=<YOUR_PORT>",
+
+            # TODO: Please fill in the following values
+            "CUSTOMER_ID": "<customer_id>",
+            "GCP_PROJECT_ID": "<project-id>",
+            "GCP_REGION": "<region>",
+            "GCP_DATAPROC_CLUSTER_NAME": "<dataproc-cluster-name>",
+            "GCP_BIGTABLE_INSTANCE_ID": "<bigtable-instance-id>",
+            "ARTIFACT_PREFIX": "<customer-artifact-bucket>",
+            "CLOUD_PROVIDER": "<gcp | aws>"
+        },
+    ),
+)
+
+
+test = Team(
+    outputNamespace="data",
+    env=EnvironmentVariables(
+        common={},
+        modeEnvironments={
+            RunMode.BACKFILL: {},
+            RunMode.UPLOAD: {}
+        }
+    ),
+)
+
+team_conf = Team(
+    outputNamespace="test",
+    env=EnvironmentVariables(
+        common={},
+    ),
+)

ai/chronon/resources/gcp/zipline-cli-install.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+
+function print_usage() {
+    echo "Usage: $0 [OPTIONS]"
+    echo "Options:"
+    echo "  --artifact_prefix <gcs_bucket>      Specify the gcs bucket to upload artifacts to e.g. \"gs://ck-zipline-artifacts\""
+    echo "  --version <version>                 Specify the version you want to run"
+    echo "  -h, --help  Show this help message"
+}
+
+if [ $# -ne 4 ]; then
+    print_usage
+    exit 1
+fi
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --artifact_prefix)
+            if [[ -z $2 ]]; then
+                echo "Error: --artifact_prefix requires a value"
+                print_usage
+                exit 1
+            fi
+            ARTIFACT_PREFIX="$2"
+            shift 2
+            ;;
+        -h|--help)
+            print_usage
+            exit 0
+            ;;
+        --version)
+            if [[ -z $2 ]]; then
+                echo "Error: --version requires a value"
+                print_usage
+                exit 1
+            fi
+            VERSION="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown option: $1"
+            print_usage
+            exit 1
+            ;;
+    esac
+done
+
+gcloud storage cp "${ARTIFACT_PREFIX%/}/release/$VERSION/wheels/zipline_ai-$VERSION-py3-none-any.whl" .
+
+trap 'rm -f ./zipline_ai-$VERSION-py3-none-any.whl' EXIT
+
+pip3 uninstall zipline-ai
+
+pip3 install ./zipline_ai-$VERSION-py3-none-any.whl

ai/chronon/source.py

@@ -0,0 +1,88 @@
+"""
+Wrappers to directly create Source objects.
+"""
+
+import ai.chronon.api.ttypes as ttypes
+
+
+def EventSource(
+    table: str,
+    query: ttypes.Query,
+    topic: str = None,
+    is_cumulative: bool = None,
+) -> ttypes.Source:
+    """
+    Event Sources represent data that gets generated over-time.
+    Typically, but not necessarily, logged to message buses like kafka, kinesis or google pub/sub.
+    fct tables are also event source worthy.
+
+    Attributes:
+
+     - table: Table currently needs to be a 'ds' (date string - yyyy-MM-dd) partitioned hive table.
+              Table names can contain subpartition specs, example db.table/system=mobile/currency=USD
+     - topic: Topic is a kafka table. The table contains all the events historically came through this topic.
+     - query: The logic used to scan both the table and the topic. Contains row level transformations
+              and filtering expressed as Spark SQL statements.
+     - isCumulative: If each new hive partition contains not just the current day's events but the entire set
+                     of events since the begininng. The key property is that the events are not mutated
+                     across partitions.
+
+    """
+    return ttypes.Source(
+        events=ttypes.EventSource(
+            table=table, topic=topic, query=query, isCumulative=is_cumulative
+        )
+    )
+
+
+def EntitySource(
+    snapshot_table: str,
+    query: ttypes.Query,
+    mutation_table: str = None,
+    mutation_topic: str = None,
+) -> ttypes.Source:
+    """
+    Entity Sources represent data that gets mutated over-time - at row-level. This is a group of three data elements.
+    snapshotTable, mutationTable and mutationTopic. mutationTable and mutationTopic are only necessary if we are trying
+    to create realtime or point-in-time aggregations over these sources. Entity sources usually map 1:1 with a database
+    tables in your OLTP store that typically serves live application traffic. When mutation data is absent they map 1:1
+    to `dim` tables in star schema.
+
+    Attributes:
+     - snapshotTable: Snapshot table currently needs to be a 'ds' (date string - yyyy-MM-dd) partitioned hive table.
+     - mutationTable: Topic is a kafka table. The table contains
+                      all the events that historically came through this topic.
+                      We need all the fields present in the snapshot table, PLUS two additional fields,
+                      `mutation_time` - milliseconds since epoch of type Long that represents the time of the mutation
+                      `is_before` - a boolean flag that represents whether
+                                    this row contains values before or after the mutation.
+     - mutationTopic: The logic used to scan both the table and the topic. Contains row level transformations
+                      and filtering expressed as Spark SQL statements.
+     - query: If each new hive partition contains not just the current day's events but the entire set
+              of events since the begininng. The key property is that the events are not mutated across partitions.
+    """
+    return ttypes.Source(
+        entities=ttypes.EntitySource(
+            snapshotTable=snapshot_table,
+            mutationTable=mutation_table,
+            mutationTopic=mutation_topic,
+            query=query,
+        )
+    )
+
+
+def JoinSource(join: ttypes.Join, query: ttypes.Query) -> ttypes.Source:
+    """
+    The output of a join can be used as a source for `GroupBy`.
+    Useful for expressing complex computation in chronon.
+
+    Offline this simply means that we will compute the necessary date ranges of the join
+    before we start computing the `GroupBy`.
+
+    Online we will:
+    1. enrich the stream/topic of `join.left` with all the columns defined by the join
+    2. apply the selects & wheres defined in the `query`
+    3. perform aggregations defined in the *downstream* `GroupBy`
+    4. write the result to the kv store.
+    """
+    return ttypes.Source(joinSource=ttypes.JoinSource(join=join, query=query))

ai/chronon/staging_query.py

@@ -0,0 +1,185 @@
+
+import inspect
+import json
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+import ai.chronon.airflow_helpers as airflow_helpers
+import ai.chronon.api.common.ttypes as common
+import ai.chronon.api.ttypes as ttypes
+from ai.chronon.constants import AIRFLOW_DEPENDENCIES_KEY
+
+
+# Wrapper for EngineType
+class EngineType:
+    SPARK = ttypes.EngineType.SPARK
+    BIGQUERY = ttypes.EngineType.BIGQUERY
+
+@dataclass
+class TableDependency:
+    table: str
+    partition_column: Optional[str] = None
+    partition_format: Optional[str] = None
+    additional_partitions: Optional[List[str]] = None
+    offset: Optional[int] = None
+
+    def to_thrift(self):
+        if self.offset is None:
+            raise ValueError(f"Dependency offset for table {self.table} must be specified.")
+        offset_window = common.Window(length = self.offset, timeUnit= common.TimeUnit.DAYS)
+        return common.TableDependency(
+            tableInfo=common.TableInfo(
+                table=self.table, 
+                partitionColumn=self.partition_column,
+                partitionFormat=self.partition_format,
+                partitionInterval=common.Window(1, common.TimeUnit.DAYS)
+            ),
+            startOffset=offset_window,
+            endOffset=offset_window,
+            startCutOff=None,
+            endCutOff=None
+        )
+
+def StagingQuery(
+    name: str,
+    query: str,
+    version: int,
+    output_namespace: Optional[str] = None,
+    start_partition: Optional[str] = None,
+    table_properties: Optional[Dict[str, str]] = None,
+    setups: Optional[List[str]] = None,
+    engine_type: Optional[EngineType] = None,
+    dependencies: Optional[List[Union[TableDependency, Dict]]] = None,
+    tags: Optional[Dict[str, str]] = None,
+    # execution params
+    offline_schedule: str = "@daily",
+    conf: Optional[common.ConfigProperties] = None,
+    env_vars: Optional[common.EnvironmentVariables] = None,
+    cluster_conf: common.ClusterConfigProperties = None,
+    step_days: Optional[int] = None,
+    recompute_days: Optional[int] = None,
+) -> ttypes.StagingQuery:
+    """
+    Creates a StagingQuery object for executing arbitrary SQL queries with templated date parameters.
+
+    :param query:
+        Arbitrary spark query that should be written with template parameters:
+        - `{{ start_date }}`: Initial run uses start_partition, future runs use latest partition + 1 day
+        - `{{ end_date }}`: The end partition of the computing range
+        - `{{ latest_date }}`: End partition independent of the computing range (for cumulative sources)
+        - `{{ max_date(table=namespace.my_table) }}`: Max partition available for a given table
+        These parameters can be modified with offset and bounds:
+        - `{{ start_date(offset=-10, lower_bound='2023-01-01', upper_bound='2024-01-01') }}`
+    :type query: str
+    :param start_partition:
+        On the first run, `{{ start_date }}` will be set to this user provided start date,
+        future incremental runs will set it to the latest existing partition + 1 day.
+    :type start_partition: str
+    :param setups:
+        Spark SQL setup statements. Used typically to register UDFs.
+    :type setups: List[str]
+    :type partition_column: str
+    :param engine_type:
+        By default, spark is the compute engine. You can specify an override (eg. bigquery, etc.)
+        Use the EngineType class constants: EngineType.SPARK, EngineType.BIGQUERY, etc.
+    :type engine_type: int
+    :param tags:
+        Additional metadata that does not directly affect computation, but is useful for management.
+    :type tags: Dict[str, str]
+    :param offline_schedule:
+        The offline schedule interval for batch jobs. Format examples:
+        '@hourly': '0 * * * *',
+        '@daily': '0 0 * * *',
+        '@weekly': '0 0 * * 0',
+        '@monthly': '0 0 1 * *',
+        '@yearly': '0 0 1 1 *'
+    :type offline_schedule: str
+    :param conf:
+        Configuration properties for the StagingQuery.
+    :type conf: common.ConfigProperties
+    :param env_vars:
+        Environment variables for the StagingQuery.
+    :type env_vars: common.EnvironmentVariables
+    :param cluster_conf:
+        Cluster configuration properties for the join.
+    :param step_days:
+        The maximum number of days to process at once
+    :type step_days: int
+    :param dependencies:
+        List of dependencies for the StagingQuery. Each dependency can be either a TableDependency object
+        or a dictionary with 'name' and 'spec' keys.
+    :type dependencies: List[Union[TableDependency, Dict]]
+    :param recompute_days:
+        Used by orchestrator to determine how many days are recomputed on each incremental scheduled run. Should be
+        set when the source data is changed in-place (i.e. existing partitions overwritten with new data each day up to
+        X days later) or when you want partially mature aggregations (i.e. a 7 day window, but start computing it from
+        day 1, and refresh it for the next 6 days)
+    :type recompute_days: int
+    :return:
+        A StagingQuery object
+    """
+    # Get caller's filename to assign team
+    team = inspect.stack()[1].filename.split("/")[-2]
+
+    # Create execution info
+    exec_info = common.ExecutionInfo(
+        scheduleCron=offline_schedule,
+        conf=conf,
+        env=env_vars,
+        stepDays=step_days,
+        clusterConf=cluster_conf
+    )
+
+    airflow_dependencies = []
+
+    if dependencies:
+        for d in dependencies:
+            if isinstance(d, TableDependency):
+                # Create an Airflow dependency object for the table
+                airflow_dependency = airflow_helpers.create_airflow_dependency(
+                    d.table,
+                    d.partition_column,
+                    d.additional_partitions,
+                    d.offset,
+                )
+                airflow_dependencies.append(airflow_dependency)
+            elif isinstance(d, dict):
+                # If it's already a dictionary, just append it
+                airflow_dependencies.append(d)
+            else:
+                raise ValueError(
+                    "Dependencies must be either TableDependency instances or dictionaries."
+                )
+
+    custom_json = json.dumps({AIRFLOW_DEPENDENCIES_KEY: airflow_dependencies})
+
+    # Create metadata
+    meta_data = ttypes.MetaData(
+        name=name,
+        outputNamespace=output_namespace,
+        team=team,
+        executionInfo=exec_info,
+        tags=tags,
+        customJson=custom_json,
+        tableProperties=table_properties,
+        version=str(version)
+    )
+
+    thrift_deps = []
+    if dependencies and len(dependencies) > 0:
+        for d in dependencies:
+            if d and isinstance(d, TableDependency):
+                thrift_deps.append(d.to_thrift())
+
+    # Create and return the StagingQuery object with camelCase parameter names
+    staging_query = ttypes.StagingQuery(
+        metaData=meta_data,
+        query=query,
+        startPartition=start_partition,
+        setups=setups,
+        engineType=engine_type,
+        tableDependencies=thrift_deps,
+        recomputeDays=recompute_days,
+    )
+
+    return staging_query