Merge branch 'main' into nikhil/orchestrator_stubs

Author: varant-zlai

WORKSPACE

@@ -102,6 +102,21 @@ scalafmt_default_config()
 
 scalafmt_repositories()
 
+# For jar jar to help with shading
+load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
+
+git_repository(
+    name = "com_github_johnynek_bazel_jar_jar",
+    commit = "352e66efa42434154ff2c0406ffd395efcbec92c", # Latest commit SHA as of 2024/11/05
+    remote = "https://github.com/johnynek/bazel_jar_jar.git",
+)
+
+load(
+    "@com_github_johnynek_bazel_jar_jar//:jar_jar.bzl",
+    "jar_jar_repositories",
+)
+jar_jar_repositories()
+
 # For Protobuf support
 http_archive(
     name = "rules_proto",

aggregator/BUILD.bazel

@@ -16,7 +16,6 @@ scala_library(
         maven_artifact("org.apache.datasketches:datasketches-java"),
         maven_artifact("org.apache.commons:commons-lang3"),
         maven_artifact("org.slf4j:slf4j-api"),
-        maven_artifact("org.apache.logging.log4j:log4j-slf4j-impl"),
         maven_artifact_with_suffix("org.scala-lang.modules:scala-collection-compat"),
     ],
 )
@@ -27,7 +26,6 @@ test_deps = [
     "//api:thrift_java",
     # Libraries
     maven_artifact("org.slf4j:slf4j-api"),
-    maven_artifact("org.apache.logging.log4j:log4j-slf4j-impl"),
     maven_artifact("com.google.code.gson:gson"),
     maven_artifact("org.apache.datasketches:datasketches-memory"),
     maven_artifact("org.apache.datasketches:datasketches-java"),

api/BUILD.bazel

@@ -12,7 +12,6 @@ java_library(
     deps = [
         maven_artifact("javax.annotation:javax.annotation.api"),
         maven_artifact("org.slf4j:slf4j-api"),
-        maven_artifact("org.apache.logging.log4j:log4j-slf4j-impl"),
         maven_artifact("org.apache.commons:commons-lang3"),
         maven_artifact("com.google.code.gson:gson"),
     ],
@@ -35,7 +34,6 @@ scala_library(
         maven_artifact("com.fasterxml.jackson.core:jackson-core"),
         maven_artifact("com.fasterxml.jackson.core:jackson-databind"),
         maven_artifact("org.slf4j:slf4j-api"),
-        maven_artifact("org.apache.logging.log4j:log4j-slf4j-impl"),
         maven_artifact("org.apache.commons:commons-lang3"),
         maven_artifact("com.google.code.gson:gson"),
         maven_artifact_with_suffix("org.scala-lang.modules:scala-collection-compat"),
@@ -51,7 +49,6 @@ test_deps = [
     maven_artifact("com.fasterxml.jackson.core:jackson-core"),
     maven_artifact("com.fasterxml.jackson.core:jackson-databind"),
     maven_artifact("org.slf4j:slf4j-api"),
-    maven_artifact("org.apache.logging.log4j:log4j-slf4j-impl"),
     # Testing
     maven_artifact_with_suffix("org.scala-lang.modules:scala-parser-combinators"),
     maven_artifact_with_suffix("org.scala-lang.modules:scala-collection-compat"),

api/py/ai/chronon/cli/compile/parse_teams.py

@@ -123,6 +123,7 @@ def _merge_mode_maps(
                 result.backfill = _merge_maps(result.common, result.backfill)
                 result.upload = _merge_maps(result.common, result.upload)
                 result.streaming = _merge_maps(result.common, result.streaming)
+                result.serving = _merge_maps(result.common, result.serving)
                 result.common = None
             continue
 
@@ -135,5 +136,6 @@ def _merge_mode_maps(
         result.streaming = _merge_maps(
             result.streaming, mode_map.common, mode_map.streaming
         )
+        result.serving = _merge_maps(result.serving, mode_map.common, mode_map.serving)
 
     return result

api/py/ai/chronon/repo/aws.py

@@ -31,14 +31,15 @@
 class AwsRunner(Runner):
     def __init__(self, args):
         aws_jar_path = AwsRunner.download_zipline_aws_jar(
-            ZIPLINE_DIRECTORY, get_customer_id(), ZIPLINE_AWS_JAR_DEFAULT
+            ZIPLINE_DIRECTORY, get_customer_id(), args["version"], ZIPLINE_AWS_JAR_DEFAULT
         )
         service_jar_path = AwsRunner.download_zipline_aws_jar(
-            ZIPLINE_DIRECTORY, get_customer_id(), ZIPLINE_AWS_SERVICE_JAR
+            ZIPLINE_DIRECTORY, get_customer_id(), args["version"], ZIPLINE_AWS_SERVICE_JAR
         )
         jar_path = (
             f"{service_jar_path}:{aws_jar_path}" if args['mode'] == "fetch" else aws_jar_path
         )
+        self.version = args.get("version", "latest")
 
         super().__init__(args, os.path.expanduser(jar_path))
 
@@ -58,10 +59,10 @@ def upload_s3_file(
             raise RuntimeError(f"Failed to upload {source_file_name}: {str(e)}")
 
     @staticmethod
-    def download_zipline_aws_jar(destination_dir: str, customer_id: str, jar_name: str):
+    def download_zipline_aws_jar(destination_dir: str, customer_id: str, version: str, jar_name: str):
         s3_client = boto3.client("s3")
         destination_path = f"{destination_dir}/{jar_name}"
-        source_key_name = f"jars/{jar_name}"
+        source_key_name = f"release/{version}/jars/{jar_name}"
         bucket_name = f"zipline-artifacts-{customer_id}"
 
         are_identical = (
@@ -163,7 +164,7 @@ def generate_emr_submitter_args(
         # include jar uri. should also already be in the bucket
         jar_uri = (
             f"{zipline_artifacts_bucket_prefix}-{get_customer_id()}"
-            + f"/jars/{ZIPLINE_AWS_JAR_DEFAULT}"
+            + f"/release/{self.version}/jars/{ZIPLINE_AWS_JAR_DEFAULT}"
         )
 
         final_args = "{user_args} --jar-uri={jar_uri} --job-type={job_type} --main-class={main_class}"
@@ -193,7 +194,7 @@ def generate_emr_submitter_args(
                     job_type=job_type.value,
                     main_class=main_class,
                 )
-                + f" --additional-conf-path=additional-confs.yaml --files={s3_file_args}"
+                + f" --additional-conf-path={EMR_MOUNT_FILE_PREFIX}additional-confs.yaml --files={s3_file_args}"
             )
         else:
             raise ValueError(f"Invalid job type: {job_type}")

api/py/ai/chronon/repo/default_runner.py

@@ -24,6 +24,8 @@ def __init__(self, args, jar_path):
         self.kafka_bootstrap = args.get("kafka_bootstrap")
         self.mock_source = args.get("mock_source")
         self.savepoint_uri = args.get("savepoint_uri")
+        self.validate = args.get("validate")
+        self.validate_rows = args.get("validate_rows")
 
         valid_jar = args["online_jar"] and os.path.exists(args["online_jar"])
 

api/py/ai/chronon/repo/gcp.py

@@ -31,10 +31,10 @@
 class GcpRunner(Runner):
     def __init__(self, args):
         gcp_jar_path = GcpRunner.download_zipline_dataproc_jar(
-            ZIPLINE_DIRECTORY, get_customer_id(), ZIPLINE_GCP_JAR_DEFAULT
+            ZIPLINE_DIRECTORY, get_customer_id(), args["version"], ZIPLINE_GCP_JAR_DEFAULT
         )
         service_jar_path = GcpRunner.download_zipline_dataproc_jar(
-            ZIPLINE_DIRECTORY, get_customer_id(), ZIPLINE_GCP_SERVICE_JAR
+            ZIPLINE_DIRECTORY, get_customer_id(), args["version"], ZIPLINE_GCP_SERVICE_JAR
         )
         jar_path = (
             f"{service_jar_path}:{gcp_jar_path}" if args["mode"] == "fetch" else gcp_jar_path
@@ -156,11 +156,11 @@ def compare_gcs_and_local_file_hashes(
 
     @staticmethod
     def download_zipline_dataproc_jar(
-        destination_dir: str, customer_id: str, jar_name: str
+        destination_dir: str, customer_id: str, version: str, jar_name: str
     ):
         bucket_name = f"zipline-artifacts-{customer_id}"
 
-        source_blob_name = f"jars/{jar_name}"
+        source_blob_name = f"release/{version}/jars/{jar_name}"
         destination_path = f"{destination_dir}/{jar_name}"
 
         are_identical = (
@@ -186,6 +186,7 @@ def generate_dataproc_submitter_args(
         self,
         user_args: str,
         job_type: JobType = JobType.SPARK,
+        version: str = "latest",
         local_files_to_upload: List[str] = [],
     ):
         customer_warehouse_bucket_name = f"zipline-warehouse-{get_customer_id()}"
@@ -215,7 +216,7 @@ def generate_dataproc_submitter_args(
         # include jar uri. should also already be in the bucket
         jar_uri = (
             f"{zipline_artifacts_bucket_prefix}-{get_customer_id()}"
-            + f"/jars/{ZIPLINE_GCP_JAR_DEFAULT}"
+            + f"/release/{version}/jars/{ZIPLINE_GCP_JAR_DEFAULT}"
         )
 
         final_args = "{user_args} --jar-uri={jar_uri} --job-type={job_type} --main-class={main_class}"
@@ -224,7 +225,7 @@ def generate_dataproc_submitter_args(
             main_class = "ai.chronon.flink.FlinkJob"
             flink_jar_uri = (
                 f"{zipline_artifacts_bucket_prefix}-{get_customer_id()}"
-                + f"/jars/{ZIPLINE_GCP_FLINK_JAR_DEFAULT}"
+                + f"/release/{version}/jars/{ZIPLINE_GCP_FLINK_JAR_DEFAULT}"
             )
             return (
                 final_args.format(
@@ -258,9 +259,13 @@ def run_dataproc_flink_streaming(self):
             "-ZGCP_PROJECT_ID": GcpRunner.get_gcp_project_id(),
             "-ZGCP_BIGTABLE_INSTANCE_ID": GcpRunner.get_gcp_bigtable_instance_id(),
             "--savepoint-uri": self.savepoint_uri,
+            "--validate-rows": self.validate_rows,
         }
 
-        flag_args = {"--mock-source": self.mock_source}
+        flag_args = {
+            "--mock-source": self.mock_source,
+            "--validate": self.validate
+        }
         flag_args_str = " ".join(key for key, value in flag_args.items() if value)
 
         user_args_str = " ".join(

api/py/ai/chronon/repo/init.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python
+
+import shutil
+import os
+from importlib_resources import files
+
+import click
+
+
+@click.command(name="init")
+@click.option(
+    "--chronon_root",
+    envvar="CHRONON_ROOT",
+    help="Path to the root chronon folder",
+    default=os.path.join(os.getcwd(), "zipline"),
+    type=click.Path(file_okay=False, writable=True),
+)
+@click.pass_context
+def main(ctx, chronon_root):
+    try:
+        template_path = files("ai.chronon").joinpath("../../test/sample")
+        print(template_path)
+    except AttributeError:
+        click.echo("Error: Template directory not found in package.", err=True)
+        return
+
+    if not template_path.exists():
+        click.echo("Error: Template directory not found in package.", err=True)
+        return
+
+    target_path = os.path.abspath(chronon_root)
+
+    # Prevent accidental overwrites
+    if os.path.exists(target_path) and os.listdir(target_path):
+        click.confirm(f"Warning: {target_path} is not empty. Proceed?", abort=True)
+
+    click.echo(f"Generating scaffolding at {target_path}...")
+
+    try:
+        shutil.copytree(template_path, target_path, dirs_exist_ok=True)
+        click.echo("Project scaffolding created successfully! 🎉")
+        click.echo(
+            f""""Please copy the following command to your shell config:
+    `export PYTHONPATH={target_path}:$PYTHONPATH`"""
+        )
+    except Exception as e:
+        click.echo(f"Error: {e}", err=True)
+
+
+if __name__ == "__main__":
+    main()

api/py/ai/chronon/repo/run.py

@@ -53,7 +53,7 @@ def set_defaults(ctx):
     defaults = {
         "mode": "backfill",
         "dataproc": False,
-        "ds": today,
+        "ds": today,  # TODO: this breaks if the partition column is not the same as yyyy-MM-dd.
         "app_name": os.environ.get("APP_NAME"),
         "online_jar": os.environ.get("CHRONON_ONLINE_JAR"),
         "repo": chronon_repo_path,
@@ -118,7 +118,7 @@ def set_defaults(ctx):
     "--online-class",
     help="Class name of Online Impl. Used for streaming and metadata-upload mode.",
 )
-@click.option("--version", help="Chronon version to use.")
+@click.option("--version", default="latest", help="Chronon version to use.")
 @click.option(
     "--spark-version", default="2.4.0", help="Spark version to use for downloading jar."
 )
@@ -159,6 +159,10 @@ def set_defaults(ctx):
     help="Use a mocked data source instead of a real source for groupby-streaming Flink.",
 )
 @click.option("--savepoint-uri", help="Savepoint URI for Flink streaming job")
+@click.option("--validate", is_flag=True,
+              help="Validate the catalyst util Spark expression evaluation logic")
+@click.option("--validate-rows", default="10000",
+              help="Number of rows to  run the validation on")
 @click.pass_context
 def main(
     ctx,
@@ -190,6 +194,8 @@ def main(
     kafka_bootstrap,
     mock_source,
     savepoint_uri,
+    validate,
+    validate_rows
 ):
     unknown_args = ctx.args
     click.echo("Running with args: {}".format(ctx.params))

api/py/ai/chronon/repo/zipline.py

@@ -2,6 +2,7 @@
 
 from ai.chronon.repo.compile import extract_and_convert
 from ai.chronon.repo.run import main as run_main
+from ai.chronon.repo.init import main as init_main
 
 
 @click.group()
@@ -11,3 +12,4 @@ def zipline():
 
 zipline.add_command(extract_and_convert)
 zipline.add_command(run_main)
+zipline.add_command(init_main)

api/py/requirements/base.in

@@ -6,4 +6,5 @@ pyspark==3.5.4
 sqlglot
 crcmod==1.7
 glom
-boto3
+boto3
+importlib-resources==6.5.2

api/py/setup.py

@@ -1,4 +1,3 @@
-
 #     Copyright (C) 2023 The Chronon Authors.
 #
 #     Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,6 +15,7 @@
 import os
 import re
 from setuptools import find_packages, setup
+import glob
 
 current_dir = os.path.abspath(os.path.dirname(__file__))
 with open(os.path.join(current_dir, "README.md"), "r") as fh:
@@ -26,8 +26,10 @@
     basic_requirements = [line for line in infile]
 
 
-__version__ = "local"
+__version__ = "0.0.1"
 __branch__ = "main"
+
+
 def get_version():
     version_str = os.environ.get("VERSION", __version__)
     branch_str = os.environ.get("BRANCH", __branch__)
@@ -36,19 +38,20 @@ def get_version():
     # If the prefix is the branch name, then convert it as suffix after '+' to make it Python PEP440 complaint
     if version_str.startswith(branch_str + "-"):
         version_str = "{}+{}".format(
-            version_str.replace(branch_str + "-", ""),
-            branch_str
+            version_str.replace(branch_str + "-", ""), branch_str
         )
 
     # Replace multiple continuous '-' or '_' with a single period '.'.
     # In python version string, the label identifier that comes after '+', is all separated by periods '.'
-    version_str = re.sub(r'[-_]+', '.', version_str)
+    version_str = re.sub(r"[-_]+", ".", version_str)
 
     return version_str
 
+
+resources = [f for f in glob.glob('test/sample/**/*', recursive=True) if os.path.isfile(f)]
 setup(
     classifiers=[
-        "Programming Language :: Python :: 3.7"
+        "Programming Language :: Python :: 3.11"
     ],
     long_description=long_description,
     long_description_content_type="text/markdown",
@@ -58,16 +61,17 @@ def get_version():
         ]
     },
     description="Zipline python API library",
-    include_package_data=True,
     install_requires=basic_requirements,
     name="zipline-ai",
-    packages=find_packages(),
+    packages=find_packages(include=["ai*"]),
+    include_package_data=True,
+    package_data={"": resources},
     extras_require={
         # Extra requirement to have access to cli commands in python2 environments.
         "pip2compat": ["click<8"]
     },
-    python_requires=">=3.7",
+    python_requires=">=3.11",
     url=None,
     version=get_version(),
-    zip_safe=False
+    zip_safe=False,
 )