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

ai/chronon/cli/git_utils.py

@@ -0,0 +1,156 @@
+import subprocess
+import sys
+from pathlib import Path
+from typing import List, Optional
+
+from ai.chronon.cli.logger import get_logger
+
+logger = get_logger()
+
+
+def get_current_branch() -> str:
+
+    try:
+        subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL)
+
+        return (
+            subprocess.check_output(["git", "rev-parse", "--abbrev-ref", "HEAD"])
+            .decode("utf-8")
+            .strip()
+        )
+
+    except subprocess.CalledProcessError as e:
+
+        try:
+            head_file = Path(".git/HEAD").resolve()
+
+            if head_file.exists():
+                content = head_file.read_text().strip()
+
+                if content.startswith("ref: refs/heads/"):
+                    return content.split("/")[-1]
+
+        except Exception:
+            pass
+
+        print(
+            f"⛔ Error: {e.stderr.decode('utf-8') if e.stderr else 'Not a git repository or no commits'}",
+            file=sys.stderr,
+        )
+
+        raise
+
+
+def get_fork_point(base_branch: str = "main") -> str:
+    try:
+
+        return (
+            subprocess.check_output(["git", "merge-base", base_branch, "HEAD"])
+            .decode("utf-8")
+            .strip()
+        )
+
+    except subprocess.CalledProcessError as e:
+        print(
+            f"⛔ Error: {e.stderr.decode('utf-8') if e.stderr else f'Could not determine fork point from {base_branch}'}",
+            file=sys.stderr,
+        )
+        raise
+
+
+def get_file_content_at_commit(file_path: str, commit: str) -> Optional[str]:
+    try:
+        return subprocess.check_output(["git", "show", f"{commit}:{file_path}"]).decode(
+            "utf-8"
+        )
+    except subprocess.CalledProcessError:
+        return None
+
+
+def get_current_file_content(file_path: str) -> Optional[str]:
+    try:
+        return Path(file_path).read_text()
+    except Exception:
+        return None
+
+
+def get_changes_since_commit(path: str, commit: Optional[str] = None) -> List[str]:
+
+    path = Path(path).resolve()
+    if not path.exists():
+        print(f"⛔ Error: Path does not exist: {path}", file=sys.stderr)
+        raise ValueError(f"Path does not exist: {path}")
+
+    try:
+        subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL)
+        commit_range = f"{commit}..HEAD" if commit else "HEAD"
+
+        changes = (
+            subprocess.check_output(
+                ["git", "diff", "--name-only", commit_range, "--", str(path)]
+            )
+            .decode("utf-8")
+            .splitlines()
+        )
+
+    except subprocess.CalledProcessError:
+
+        changes = (
+            subprocess.check_output(["git", "diff", "--name-only", "--", str(path)])
+            .decode("utf-8")
+            .splitlines()
+        )
+
+    try:
+
+        untracked = (
+            subprocess.check_output(
+                ["git", "ls-files", "--others", "--exclude-standard", str(path)]
+            )
+            .decode("utf-8")
+            .splitlines()
+        )
+
+        changes.extend(untracked)
+
+    except subprocess.CalledProcessError as e:
+
+        print(
+            f"⛔ Error: {e.stderr.decode('utf-8') if e.stderr else 'Failed to get untracked files'}",
+            file=sys.stderr,
+        )
+
+        raise
+
+    logger.info(f"Changes since commit: {changes}")
+
+    return [change for change in changes if change.strip()]
+
+
+def get_changes_since_fork(path: str, base_branch: str = "main") -> List[str]:
+    try:
+        fork_point = get_fork_point(base_branch)
+        path = Path(path).resolve()
+
+        # Get all potential changes
+        changed_files = set(get_changes_since_commit(str(path), fork_point))
+
+        # Filter out files that are identical to fork point
+        real_changes = []
+        for file in changed_files:
+            fork_content = get_file_content_at_commit(file, fork_point)
+            current_content = get_current_file_content(file)
+
+            if fork_content != current_content:
+                real_changes.append(file)
+
+        logger.info(f"Changes since fork: {real_changes}")
+
+        return real_changes
+
+    except subprocess.CalledProcessError as e:
+        print(
+            f"⛔ Error: {e.stderr.decode('utf-8') if e.stderr else f'Failed to get changes since fork from {base_branch}'}",
+            file=sys.stderr,
+        )
+        raise

ai/chronon/cli/logger.py

@@ -0,0 +1,61 @@
+import logging
+import sys
+from datetime import datetime
+
+TIME_COLOR = "\033[36m"  # Cyan
+LEVEL_COLORS = {
+    logging.DEBUG: "\033[36m",  # Cyan
+    logging.INFO: "\033[32m",  # Green
+    logging.WARNING: "\033[33m",  # Yellow
+    logging.ERROR: "\033[31m",  # Red
+    logging.CRITICAL: "\033[41m",  # White on Red
+}
+FILE_COLOR = "\033[35m"  # Purple
+RESET = "\033[0m"
+
+
+class ColorFormatter(logging.Formatter):
+
+    def format(self, record):
+
+        time_str = datetime.fromtimestamp(record.created).strftime("%H:%M:%S")
+        level_color = LEVEL_COLORS.get(record.levelno)
+
+        return (
+            f"{TIME_COLOR}{time_str}{RESET} "
+            f"{level_color}{record.levelname}{RESET} "
+            f"{FILE_COLOR}{record.filename}:{record.lineno}{RESET} - "
+            f"{record.getMessage()}"
+        )
+
+
+def get_logger(log_level=logging.INFO):
+    logger = logging.getLogger(__name__)
+
+    # no need to reset if a handler already exists
+    if not logger.hasHandlers():
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(ColorFormatter())
+
+        logger.addHandler(handler)
+        logger.setLevel(log_level)
+
+    return logger
+
+
+def red(text):
+    return f"\033[1;91m{text}\033[0m"
+
+
+def green(text):
+    return f"\033[1;92m{text}\033[0m"
+
+
+def require(cond, message):
+    if not cond:
+        print(f"X: {message}")
+        sys.exit(1)
+
+
+def done(cond, message):
+    print(f"DONE: {message}")

ai/chronon/constants.py

@@ -0,0 +1,3 @@
+AIRFLOW_DEPENDENCIES_KEY = "airflowDependencies"
+AIRFLOW_LABEL_DEPENDENCIES_KEY = "airflowLabelDependencies"
+PARTITION_COLUMN_KEY = "spark.chronon.partition.column"

ai/chronon/eval/__init__.py

@@ -0,0 +1,122 @@
+from typing import Any, List
+
+from pyspark.sql import DataFrame, SparkSession
+
+import ai.chronon.api.ttypes as chronon
+from ai.chronon.eval.query_parsing import get_tables_from_query
+from ai.chronon.eval.sample_tables import sample_tables, sample_with_query
+from ai.chronon.eval.table_scan import (
+    TableScan,
+    clean_table_name,
+    table_scans_in_group_by,
+    table_scans_in_join,
+    table_scans_in_source,
+)
+
+
+def eval(obj: Any) -> List[DataFrame]:
+
+    if isinstance(obj, chronon.Source):
+        return _run_table_scans(table_scans_in_source(obj))
+
+    elif isinstance(obj, chronon.GroupBy):
+        return _run_table_scans(table_scans_in_group_by(obj))
+
+    elif isinstance(obj, chronon.Join):
+        return _run_table_scans(table_scans_in_join(obj))
+
+    elif isinstance(obj, chronon.StagingQuery):
+        return _sample_and_eval_query(_render_staging_query(obj))
+
+    elif isinstance(obj, str):
+        has_white_spaces = any(char.isspace() for char in obj)
+        if has_white_spaces:
+            return _sample_and_eval_query(obj)
+        else:
+            return _sample_and_eval_query(f"SELECT * FROM {obj} LIMIT 1000")
+
+    elif isinstance(obj, chronon.Model):
+        _run_table_scans(table_scans_in_source(obj.source))
+
+    else:
+        raise Exception(f"Unsupported object type for: {obj}")
+
+
+def _sample_and_eval_query(query: str) -> DataFrame:
+
+    table_names = get_tables_from_query(query)
+    sample_tables(table_names)
+
+    clean_query = query
+    for table_name in table_names:
+        clean_name = clean_table_name(table_name)
+        clean_query = clean_query.replace(table_name, clean_name)
+
+    return _run_query(clean_query)
+
+
+def _run_query(query: str) -> DataFrame:
+    spark = _get_spark()
+    return spark.sql(query)
+
+
+def _sample_table_scan(table_scan: TableScan) -> str:
+    table = table_scan.table
+    output_path = table_scan.output_path()
+    query = table_scan.raw_scan_query(local_table_view=False)
+    return sample_with_query(table, query, output_path)
+
+
+def _run_table_scans(table_scans: List[TableScan]) -> List[DataFrame]:
+    spark = _get_spark()
+    df_list = []
+
+    for table_scan in table_scans:
+        output_path = table_scan.output_path()
+
+        status = " (exists)" if output_path.exists() else ""
+        print(
+            f"table: {table_scan.table}\n"
+            f"view: {table_scan.view_name()}\n"
+            f"local_file: {output_path}{status}\n"
+        )
+
+    for table_scan in table_scans:
+
+        view_name = table_scan.view_name()
+        output_path = _sample_table_scan(table_scan)
+
+        print(f"Creating view {view_name} from parquet file {output_path}")
+        df = spark.read.parquet(str(output_path))
+        df.createOrReplaceTempView(view_name)
+
+        scan_query = table_scan.scan_query(local_table_view=True)
+        print(f"Scanning {table_scan.table} with query: \n{scan_query}\n")
+        df = spark.sql(scan_query)
+        df.show(5)
+        df_list.append(df)
+
+    return df_list
+
+
+_spark: SparkSession = None
+
+
+def _get_spark() -> SparkSession:
+    global _spark
+    if not _spark:
+        _spark = (
+            SparkSession.builder.appName("Chronon Evaluator")
+            .config("spark.driver.bindAddress", "127.0.0.1")
+            .config("spark.driver.host", "127.0.0.1")
+            .config("spark.sql.parquet.columnarReaderBatchSize", "16")
+            .config("spark.executor.memory", "4g")
+            .config("spark.driver.memory", "4g")
+            .config("spark.driver.maxResultSize", "2g")
+            .getOrCreate()
+        )
+    return _spark
+
+
+def _render_staging_query(staging_query: chronon.StagingQuery) -> str:
+    raise NotImplementedError("Staging query evals are not yet implemented")

ai/chronon/eval/query_parsing.py

@@ -0,0 +1,19 @@
+from typing import List
+
+
+def get_tables_from_query(sql_query) -> List[str]:
+    import sqlglot
+
+    # Parse the query
+    parsed = sqlglot.parse_one(sql_query, dialect="bigquery")
+
+    # Extract all table references
+    tables = parsed.find_all(sqlglot.exp.Table)
+
+    table_names = []
+    for table in tables:
+        name_parts = [part for part in [table.catalog, table.db, table.name] if part]
+        table_name = ".".join(name_parts)
+        table_names.append(table_name)
+
+    return table_names

ai/chronon/eval/sample_tables.py

@@ -0,0 +1,100 @@
+import os
+from typing import List
+
+from ai.chronon.eval.table_scan import local_warehouse
+
+
+def sample_with_query(table, query, output_path) -> str:
+    # if file exists, skip
+    if os.path.exists(output_path):
+        print(f"File {output_path} already exists. Skipping sampling.")
+        return output_path
+
+    raw_scan_query = query
+    print(f"Sampling {table} with query: {raw_scan_query}")
+
+    _sample_internal(raw_scan_query, output_path)
+    return output_path
+
+
+def sample_tables(table_names: List[str]) -> None:
+
+    for table in table_names:
+        query = f"SELECT * FROM {table} LIMIT 10000"
+        sample_with_query(table, query, local_warehouse / f"{table}.parquet")
+
+
+_sampling_engine = os.getenv("CHRONON_SAMPLING_ENGINE", "bigquery")
+
+
+def _sample_internal(query, output_path) -> str:
+    if _sampling_engine == "bigquery":
+        _sample_bigquery(query, output_path)
+    elif _sampling_engine == "trino":
+        _sample_trino(query, output_path)
+    else:
+        raise ValueError("Invalid sampling engine")
+
+
+def _sample_trino(query, output_path):
+    raise NotImplementedError("Trino sampling is not yet implemented")
+
+
+def _sample_bigquery(query, output_path):
+
+    from google.cloud import bigquery
+
+    project_id = os.getenv("GCP_PROJECT_ID")
+    assert project_id, "Please set the GCP_PROJECT_ID environment variable"
+
+    client = bigquery.Client(project=project_id)
+
+    results = client.query_and_wait(query)
+
+    df = results.to_dataframe()
+    df.to_parquet(output_path)
+
+
+def _sample_bigquery_fast(query, destination_path):
+    import os
+
+    import pyarrow.parquet as pq
+    from google.cloud import bigquery
+    from google.cloud.bigquery_storage import BigQueryReadClient
+    from google.cloud.bigquery_storage_v1.types import DataFormat, ReadSession
+
+    project_id = os.getenv("GCP_PROJECT_ID")
+    assert project_id, "Please set the GCP_PROJECT_ID environment variable"
+
+    client = bigquery.Client(project=project_id)
+    bqstorage_client = BigQueryReadClient()
+
+    # Create query job
+    query_job = client.query(query)
+    table_ref = query_job.destination
+
+    # Create read session
+    read_session = ReadSession()
+    read_session.table = table_ref.to_bqstorage()
+    read_session.data_format = DataFormat.ARROW
+
+    print("Fetching from BigQuery... (this might take a while)")
+
+    session = bqstorage_client.create_read_session(
+        parent=f"projects/{client.project}",
+        read_session=read_session,
+        max_stream_count=1,
+    )
+
+    print("Writing to local parquet file...")
+
+    # Read using Arrow
+    stream = bqstorage_client.read_rows(session.streams[0].name)
+    table = stream.to_arrow(read_session=session)
+
+    # Write to Parquet directly
+    pq.write_table(table, destination_path)
+
+    print(f"Wrote results to {destination_path}")
+
+    return destination_path

ai/chronon/eval/table_scan.py

@@ -0,0 +1,186 @@
+import hashlib
+import os
+import re
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import List, Tuple
+
+import ai.chronon.api.ttypes as chronon
+
+
+def clean_table_name(name: str) -> str:
+    return re.sub(r"[^a-zA-Z0-9_]", "_", name)
+
+
+local_warehouse = Path(os.getenv("CHRONON_ROOT", os.getcwd())) / "local_warehouse"
+limit = int(os.getenv("SAMPLE_LIMIT", "100"))
+# create local_warehouse if it doesn't exist
+local_warehouse.mkdir(parents=True, exist_ok=True)
+
+
+@dataclass
+class TableScan:
+    table: str
+    partition_col: str
+    partition_date: str
+    query: chronon.Query
+    is_mutations: bool = False
+
+    def output_path(self) -> str:
+        return Path(local_warehouse) / f"{self.view_name()}.parquet"
+
+    def view_name(self) -> str:
+        return clean_table_name(self.table) + "_" + self.where_id()
+
+    def table_name(self, local_table_view) -> str:
+        return self.view_name() if local_table_view else self.table
+
+    def where_id(self) -> str:
+        return "_" + hashlib.md5(self.where_block().encode()).hexdigest()[:3]
+
+    def where_block(self) -> str:
+        wheres = []
+        partition_scan = f"{self.partition_col} = '{self.partition_date}'"
+        wheres.append(partition_scan)
+
+        if self.query.wheres:
+            wheres.extend(self.query.wheres)
+
+        return " AND\n    ".join([f"({where})" for where in wheres])
+
+    def raw_scan_query(self, local_table_view: bool = True) -> str:
+        return f"""
+SELECT * FROM {self.table_name(local_table_view)}
+WHERE
+    {self.where_block()}
+LIMIT {limit}
+"""
+
+    def scan_query(self, local_table_view=True) -> str:
+        selects = []
+        base_selects = self.query.selects.copy()
+
+        if self.is_mutations:
+            base_selects["is_before"] = coalesce(self.query.reversalColumn, "is_before")
+            base_selects["mutation_ts"] = coalesce(
+                self.query.mutationTimeColumn, "mutation_ts"
+            )
+
+        if self.query.timeColumn:
+            base_selects["ts"] = coalesce(self.query.timeColumn, "ts")
+
+        for k, v in base_selects.items():
+            selects.append(f"{v} as {k}")
+        select_clauses = ",\n    ".join(selects)
+
+        return f"""
+SELECT
+    {select_clauses}
+FROM
+    {self.table_name(local_table_view)}
+WHERE
+    {self.where_block()}
+LIMIT
+    {limit}
+"""
+
+
+# TODO: use teams.py to get the default date column
+DEFAULT_DATE_COLUMN = "_date"
+DEFAULT_DATE_FORMAT = "%Y-%m-%d"
+
+two_days_ago = (datetime.now() - timedelta(days=2)).strftime(DEFAULT_DATE_FORMAT)
+
+_sample_date = os.getenv("SAMPLE_DATE", two_days_ago)
+
+
+def get_date(query: chronon.Query) -> Tuple[str, str]:
+    assert query and query.selects, "please specify source.query.selects"
+
+    partition_col = query.selects.get("ds", DEFAULT_DATE_COLUMN)
+    partition_date = coalesce(query.endPartition, _sample_date)
+
+    return (partition_col, partition_date)
+
+
+def coalesce(*args):
+    for arg in args:
+        if arg:
+            return arg
+
+
+def table_scans_in_source(source: chronon.Source) -> List[TableScan]:
+    result = []
+
+    if not source:
+        return result
+
+    if source.entities:
+        query: chronon.Query = source.entities.query
+        col, date = get_date(query)
+
+        snapshot = TableScan(source.entities.snapshotTable, col, date, query)
+        result.append(snapshot)
+
+        if source.entities.mutationTable:
+            mutations = TableScan(source.entities.mutationTable, col, date, query, True)
+            result.append(mutations)
+
+    if source.events:
+        query = source.events.query
+        col, date = get_date(query)
+        table = TableScan(source.events.table, col, date, query)
+        result.append(table)
+
+    if source.joinSource:
+        result.extend(table_scans_in_source(source.joinSource.join.left))
+
+    return result
+
+
+def table_scans_in_sources(sources: List[chronon.Source]) -> List[TableScan]:
+    result = []
+
+    for source in sources:
+        result.extend(table_scans_in_source(source))
+
+    return result
+
+
+def table_scans_in_group_by(gb: chronon.GroupBy) -> List[TableScan]:
+    if not gb:
+        return []
+
+    return table_scans_in_sources(gb.sources)
+
+
+def table_scans_in_join(join: chronon.Join) -> List[TableScan]:
+
+    result = []
+
+    if not join:
+        return result
+
+    result.extend(table_scans_in_source(join.left))
+
+    parts: List[chronon.JoinPart] = join.joinParts
+    if parts:
+        for part in parts:
+            result.extend(table_scans_in_group_by(part.groupBy))
+
+    bootstraps: List[chronon.BootstrapPart] = join.bootstrapParts
+    if bootstraps:
+        for bootstrap in bootstraps:
+            query = bootstrap.query
+            col, date = get_date(query)
+            bootstrap = TableScan(bootstrap.table, col, date, query)
+
+            result.append(bootstrap)
+
+    if join.labelParts:
+        labelParts: List[chronon.JoinPart] = join.labelParts.labels
+        for part in labelParts:
+            result.extend(table_scans_in_sources(part.groupBy))
+
+    return result

ai/chronon/fetcher/__init__.py

@@ -0,0 +1 @@
+__all__ = ['ttypes', 'constants']

ai/chronon/fetcher/constants.py

@@ -0,0 +1,15 @@
+#
+# Autogenerated by Thrift Compiler (0.22.0)
+#
+# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING
+#
+#  options string: py
+#
+
+from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException
+from thrift.protocol.TProtocol import TProtocolException
+from thrift.TRecursive import fix_spec
+from uuid import UUID
+
+import sys
+from .ttypes import *