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

ai/chronon/cli/compile/column_hashing.py

@@ -0,0 +1,313 @@
+import hashlib
+import re
+from collections import defaultdict
+from typing import Dict, List
+
+from ai.chronon.api.ttypes import Derivation, ExternalPart, GroupBy, Join, Source
+from ai.chronon.group_by import get_output_col_names
+
+
+# Returns a map of output column to semantic hash, including derivations
+def compute_group_by_columns_hashes(group_by: GroupBy, exclude_keys: bool = False) -> Dict[str, str]:
+    """
+    From the group_by object, get the final output columns after derivations.
+    """
+    # Get the output columns and their input expressions
+    output_to_input = get_pre_derived_group_by_columns(group_by)
+
+    # Get the base semantic fields that apply to all columns
+    base_semantics = []
+    for source in group_by.sources:
+        base_semantics.extend(_extract_source_semantic_info(source, group_by.keyColumns))
+
+    # Add the major version to semantic fields
+    group_by_minor_version_suffix = f"__{group_by.metaData.version}"
+    group_by_major_version = group_by.metaData.name
+    if group_by_major_version.endswith(group_by_minor_version_suffix):
+        group_by_major_version = group_by_major_version[:-len(group_by_minor_version_suffix)]
+    base_semantics.append(f"group_by_name:{group_by_major_version}")
+
+    # Compute the semantic hash for each output column
+    output_to_hash = {}
+    for output_col, input_expr in output_to_input.items():
+        semantic_components = base_semantics + [f"input_expr:{input_expr}"]
+        semantic_hash = _compute_semantic_hash(semantic_components)
+        output_to_hash[output_col] = semantic_hash
+
+    if exclude_keys:
+        output = {k: v for k, v in output_to_hash.items() if k not in group_by.keyColumns}
+    else:
+        output = output_to_hash
+    if group_by.derivations:
+        derived = build_derived_columns(output, group_by.derivations, base_semantics)
+        if not exclude_keys:
+            # We need to add keys back at this point
+            for key in group_by.keyColumns:
+                if key not in derived:
+                    derived[key] = output_to_input.get(key)
+        return derived
+    else:
+        return output
+
+
+def compute_join_column_hashes(join: Join) -> Dict[str, str]:
+    """
+    From the join object, get the final output columns -> semantic hash after derivations.
+    """
+    # Get the base semantics from the left side (table and key expression)
+    base_semantic_fields = []
+
+    output_columns = get_pre_derived_join_features(join) | get_pre_derived_source_keys(join.left)
+
+    if join.derivations:
+        return build_derived_columns(output_columns, join.derivations, base_semantic_fields)
+    else:
+        return output_columns
+
+
+def get_pre_derived_join_features(join: Join) -> Dict[str, str]:
+    return get_pre_derived_join_internal_features(join) | get_pre_derived_external_features(join)
+
+
+def get_pre_derived_external_features(join: Join) -> Dict[str, str]:
+    external_cols = []
+    if join.onlineExternalParts:
+        for external_part in join.onlineExternalParts:
+            original_external_columns = [
+                param.name for param in external_part.source.valueSchema.params
+            ]
+            prefix = get_external_part_full_name(external_part) + "_"
+            for col in original_external_columns:
+                external_cols.append(prefix + col)
+    # No meaningful semantic information on external columns, so we just return the column names as a self map
+    return {x: x for x in external_cols}
+
+
+def get_pre_derived_source_keys(source: Source) -> Dict[str, str]:
+    base_semantics = _extract_source_semantic_info(source)
+    source_keys_to_hashes = {}
+    for key, expression in extract_selects(source).items():
+        source_keys_to_hashes[key] = _compute_semantic_hash(base_semantics + [f"select:{key}={expression}"])
+    return source_keys_to_hashes
+
+
+def extract_selects(source: Source) -> Dict[str, str]:
+    if source.events:
+        return source.events.query.selects
+    elif source.entities:
+        return source.entities.query.selects
+    elif source.joinSource:
+        return source.joinSource.query.selects
+
+
+def get_pre_derived_join_internal_features(join: Join) -> Dict[str, str]:
+
+    # Get the base semantic fields from join left side (without key columns)
+    join_base_semantic_fields = _extract_source_semantic_info(join.left)
+
+    internal_features = {}
+    for jp in join.joinParts:
+        # Build key mapping semantics - include left side key expressions
+        if jp.keyMapping:
+            key_mapping_semantics = ["join_keys:" + ",".join(f"{k}:{v}" for k, v in sorted(jp.keyMapping.items()))]
+        else:
+            key_mapping_semantics = []
+
+        # Include left side key expressions that this join part uses
+        left_key_expressions = []
+        left_selects = extract_selects(join.left)
+        for gb_key in jp.groupBy.keyColumns:
+            # Get the mapped key name or use the original key name
+            left_key_name = gb_key
+            if jp.keyMapping:
+                # Find the left side key that maps to this groupby key
+                for left_key, mapped_key in jp.keyMapping.items():
+                    if mapped_key == gb_key:
+                        left_key_name = left_key
+                        break
+
+            # Add the left side expression for this key
+            left_key_expr = left_selects.get(left_key_name, left_key_name)
+            left_key_expressions.append(f"left_key:{left_key_name}={left_key_expr}")
+
+        # These semantics apply to all features in the joinPart
+        jp_base_semantics = key_mapping_semantics + left_key_expressions + join_base_semantic_fields
+
+        pre_derived_group_by_features = get_pre_derived_group_by_features(jp.groupBy, jp_base_semantics)
+
+        if jp.groupBy.derivations:
+            derived_group_by_features = build_derived_columns(
+                pre_derived_group_by_features, jp.groupBy.derivations, jp_base_semantics
+            )
+        else:
+            derived_group_by_features = pre_derived_group_by_features
+
+        for col, semantic_hash in derived_group_by_features.items():
+            prefix = jp.prefix + "_" if jp.prefix else ""
+            if join.useLongNames:
+                gb_prefix = prefix + jp.groupBy.metaData.name.replace(".", "_")
+            else:
+                key_str = "_".join(jp.groupBy.keyColumns)
+                gb_prefix = prefix + key_str
+            internal_features[gb_prefix + "_" + col] = semantic_hash
+    return internal_features
+
+
+def get_pre_derived_group_by_columns(group_by: GroupBy) -> Dict[str, str]:
+    output_columns_to_hashes = get_pre_derived_group_by_features(group_by)
+    for key_column in group_by.keyColumns:
+        key_expressions = []
+        for source in group_by.sources:
+            source_selects = extract_selects(source)  # Map[str, str]
+            key_expressions.append(source_selects.get(key_column, key_column))
+        output_columns_to_hashes[key_column] = _compute_semantic_hash(sorted(key_expressions))
+    return output_columns_to_hashes
+
+
+def get_pre_derived_group_by_features(group_by: GroupBy, additional_semantic_fields=None) -> Dict[str, str]:
+    # Get the base semantic fields that apply to all aggs
+    if additional_semantic_fields is None:
+        additional_semantic_fields = []
+    base_semantics = _get_base_group_by_semantic_fields(group_by)
+
+    output_columns = {}
+    # For group_bys with aggregations, aggregated columns
+    if group_by.aggregations:
+        for agg in group_by.aggregations:
+            input_expression_str = ",".join(get_input_expression_across_sources(group_by, agg.inputColumn))
+            for output_col_name in get_output_col_names(agg):
+                output_columns[output_col_name] = _compute_semantic_hash(base_semantics + [input_expression_str] + additional_semantic_fields)
+    # For group_bys without aggregations, selected fields from query
+    else:
+        combined_selects = defaultdict(set)
+
+        for source in group_by.sources:
+            source_selects = extract_selects(source)  # Map[str, str]
+            for key, val in source_selects.items():
+                combined_selects[key].add(val)
+
+        # Build a unified map of key to select expression from all sources
+        unified_selects = {key: ",".join(sorted(vals)) for key, vals in combined_selects.items()}
+
+        # now compute the hashes on base semantics + expression
+        selected_hashes = {key: _compute_semantic_hash(base_semantics + [val] + additional_semantic_fields) for key, val in unified_selects.items()}
+        output_columns.update(selected_hashes)
+    return output_columns
+
+
+def _get_base_group_by_semantic_fields(group_by: GroupBy) -> List[str]:
+    """
+    Extract base semantic fields from the group_by object.
+    This includes source table names, key columns, and any other relevant metadata.
+    """
+    base_semantics = []
+    for source in group_by.sources:
+        base_semantics.extend(_extract_source_semantic_info(source, group_by.keyColumns))
+
+    return sorted(base_semantics)
+
+
+def _extract_source_semantic_info(source: Source, key_columns: List[str] = None) -> List[str]:
+    """
+    Extract source information for semantic hashing.
+    Returns list of semantic components.
+
+    Args:
+        source: The source to extract info from
+        key_columns: List of key columns to include expressions for. If None, includes all selects.
+    """
+    components = []
+
+    if source.events:
+        table = source.events.table
+        mutationTable = ""
+        query = source.events.query
+        cumulative = str(source.events.isCumulative or "")
+    elif source.entities:
+        table = source.entities.snapshotTable
+        mutationTable = source.entities.mutationTable
+        query = source.entities.query
+        cumulative = ""
+    elif source.joinSource:
+        table = source.joinSource.join.metaData.name
+        mutationTable = ""
+        query = source.joinSource.query
+        cumulative = ""
+
+    components.append(f"table:{table}")
+    components.append(f"mutation_table:{mutationTable}")
+    components.append(f"cumulative:{cumulative}")
+    components.append(f"filters:{query.wheres or ''}")
+
+    selects = query.selects or {}
+    if key_columns:
+        # Only include expressions for the specified key columns
+        for key_col in sorted(key_columns):
+            expr = selects.get(key_col, key_col)
+            components.append(f"select:{key_col}={expr}")
+
+    # Add time column expression
+    if query.timeColumn:
+        time_expr = selects.get(query.timeColumn, query.timeColumn)
+        components.append(f"time_column:{query.timeColumn}={time_expr}")
+
+    return sorted(components)
+
+
+def _compute_semantic_hash(components: List[str]) -> str:
+    """
+    Compute semantic hash from a list of components.
+    Components should be ordered consistently to ensure reproducible hashes.
+    """
+    # Sort components to ensure consistent ordering
+    sorted_components = sorted(components)
+    hash_input = "|".join(sorted_components)
+    return hashlib.md5(hash_input.encode('utf-8')).hexdigest()
+
+
+def build_derived_columns(
+    base_columns_to_hashes: Dict[str, str], derivations: List[Derivation], additional_semantic_fields: List[str]
+) -> Dict[str, str]:
+    """
+    Build the derived columns from pre-derived columns and derivations.
+    """
+    # if derivations contain star, then all columns are included except the columns which are renamed
+    output_columns = {}
+    if derivations:
+        found = any(derivation.expression == "*" for derivation in derivations)
+        if found:
+            output_columns.update(base_columns_to_hashes)
+        for derivation in derivations:
+            if base_columns_to_hashes.get(derivation.expression):
+               # don't change the semantics if you're just passing a base column through derivations
+               output_columns[derivation.name] = base_columns_to_hashes[derivation.expression]
+            if derivation.name != "*":
+                # Identify base fields present within the derivation to include in the semantic hash
+                # We go long to short to avoid taking both a windowed feature and the unwindowed feature
+                # i.e. f_7d and f
+                derivation_expression = derivation.expression
+                base_col_semantic_fields = []
+                tokens = re.findall(r'\b\w+\b', derivation_expression)
+                for token in tokens:
+                    if token in base_columns_to_hashes:
+                        base_col_semantic_fields.append(base_columns_to_hashes[token])
+
+                output_columns[derivation.name] = _compute_semantic_hash(additional_semantic_fields + [f"derivation:{derivation.expression}"] + base_col_semantic_fields)
+    return output_columns
+
+
+def get_external_part_full_name(external_part: ExternalPart) -> str:
+    # The logic should be consistent with the full name logic defined
+    # in https://github.com/airbnb/chronon/blob/main/api/src/main/scala/ai/chronon/api/Extensions.scala#L677.
+    prefix = external_part.prefix + "_" if external_part.prefix else ""
+    name = external_part.source.metadata.name
+    sanitized_name = re.sub("[^a-zA-Z0-9_]", "_", name)
+    return "ext_" + prefix + sanitized_name
+
+
+def get_input_expression_across_sources(group_by: GroupBy, input_col: str):
+    expressions = []
+    for source in group_by.sources:
+        selects = extract_selects(source)
+        expressions.extend(selects.get(input_col, input_col))
+    return sorted(expressions)

ai/chronon/cli/compile/compile_context.py

@@ -0,0 +1,177 @@
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Type
+
+import ai.chronon.cli.compile.parse_teams as teams
+from ai.chronon.api.ttypes import GroupBy, Join, MetaData, Model, StagingQuery, Team
+from ai.chronon.cli.compile.conf_validator import ConfValidator
+from ai.chronon.cli.compile.display.compile_status import CompileStatus
+from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
+from ai.chronon.cli.compile.serializer import file2thrift
+from ai.chronon.cli.logger import get_logger, require
+from ai.chronon.orchestration.ttypes import ConfType
+
+logger = get_logger()
+
+
+@dataclass
+class ConfigInfo:
+    folder_name: str
+    cls: Type
+    config_type: Optional[ConfType]
+
+
+@dataclass
+class CompileContext:
+
+    def __init__(self):
+        self.chronon_root: str = os.getenv("CHRONON_ROOT", os.getcwd())
+        self.teams_dict: Dict[str, Team] = teams.load_teams(self.chronon_root)
+        self.compile_dir: str = "compiled"
+
+        self.config_infos: List[ConfigInfo] = [
+            ConfigInfo(folder_name="joins", cls=Join, config_type=ConfType.JOIN),
+            ConfigInfo(
+                folder_name="group_bys",
+                cls=GroupBy,
+                config_type=ConfType.GROUP_BY,
+            ),
+            ConfigInfo(
+                folder_name="staging_queries",
+                cls=StagingQuery,
+                config_type=ConfType.STAGING_QUERY,
+            ),
+            ConfigInfo(folder_name="models", cls=Model, config_type=ConfType.MODEL),
+            ConfigInfo(folder_name="teams_metadata", cls=MetaData, config_type=None), # only for team metadata
+        ]
+
+        self.compile_status = CompileStatus(use_live=False)
+
+        self.existing_confs: Dict[Type, Dict[str, Any]] = {}
+        for config_info in self.config_infos:
+            cls = config_info.cls
+            self.existing_confs[cls] = self._parse_existing_confs(cls)
+
+
+
+        self.validator: ConfValidator = ConfValidator(
+            input_root=self.chronon_root,
+            output_root=self.compile_dir,
+            existing_gbs=self.existing_confs[GroupBy],
+            existing_joins=self.existing_confs[Join],
+        )
+
+    def input_dir(self, cls: type) -> str:
+        """
+        - eg., input: group_by class
+        - eg., output: root/group_bys/
+        """
+        config_info = self.config_info_for_class(cls)
+        return os.path.join(self.chronon_root, config_info.folder_name)
+
+    def staging_output_dir(self, cls: type = None) -> str:
+        """
+        - eg., input: group_by class
+        - eg., output: root/compiled_staging/group_bys/
+        """
+        if cls is None:
+            return os.path.join(self.chronon_root, self.compile_dir + "_staging")
+        else:
+            config_info = self.config_info_for_class(cls)
+            return os.path.join(
+                self.chronon_root,
+                self.compile_dir + "_staging",
+                config_info.folder_name,
+            )
+
+    def output_dir(self, cls: type = None) -> str:
+        """
+        - eg., input: group_by class
+        - eg., output: root/compiled/group_bys/
+        """
+        if cls is None:
+            return os.path.join(self.chronon_root, self.compile_dir)
+        else:
+            config_info = self.config_info_for_class(cls)
+            return os.path.join(
+                self.chronon_root, self.compile_dir, config_info.folder_name
+            )
+
+    def staging_output_path(self, compiled_obj: CompiledObj):
+        """
+        - eg., input: group_by with name search.clicks.features.v1
+        - eg., output: root/compiled_staging/group_bys/search/clicks.features.v1
+        """
+
+        output_dir = self.staging_output_dir(compiled_obj.obj.__class__)  # compiled/joins
+
+        team, rest = compiled_obj.name.split(".", 1)  # search, clicks.features.v1
+
+        return os.path.join(
+            output_dir,
+            team,
+            rest,
+        )
+
+    def config_info_for_class(self, cls: type) -> ConfigInfo:
+        for info in self.config_infos:
+            if info.cls == cls:
+                return info
+
+        require(False, f"Class {cls} not found in CONFIG_INFOS")
+
+    def _parse_existing_confs(self, obj_class: type) -> Dict[str, object]:
+
+        result = {}
+
+        output_dir = self.output_dir(obj_class)
+
+        # Check if output_dir exists before walking
+        if not os.path.exists(output_dir):
+            return result
+
+        for sub_root, _sub_dirs, sub_files in os.walk(output_dir):
+
+            for f in sub_files:
+
+                if f.startswith("."):  # ignore hidden files - such as .DS_Store
+                    continue
+
+                full_path = os.path.join(sub_root, f)
+
+                try:
+                    obj = file2thrift(full_path, obj_class)
+
+                    if obj:
+                        if hasattr(obj, "metaData"):
+                            result[obj.metaData.name] = obj
+                            compiled_obj = CompiledObj(
+                                name=obj.metaData.name,
+                                obj=obj,
+                                file=obj.metaData.sourceFile,
+                                errors=None,
+                                obj_type=obj_class.__name__,
+                                tjson=open(full_path).read(),
+                            )
+                            self.compile_status.add_existing_object_update_display(compiled_obj)
+                        elif isinstance(obj, MetaData):
+                            team_metadata_name = '.'.join(full_path.split('/')[-2:]) # use the name of the file as team metadata won't have name
+                            result[team_metadata_name] = obj
+                            compiled_obj = CompiledObj(
+                                name=team_metadata_name,
+                                obj=obj,
+                                file=obj.sourceFile,
+                                errors=None,
+                                obj_type=obj_class.__name__,
+                                tjson=open(full_path).read(),
+                            )
+                            self.compile_status.add_existing_object_update_display(compiled_obj)
+                    else:
+                        logger.errors(
+                            f"Parsed object from {full_path} has no metaData attribute"
+                        )
+
+                except Exception as e:
+                    print(f"Failed to parse file {full_path}: {str(e)}", e)
+
+        return result

ai/chronon/cli/compile/compiler.py

@@ -0,0 +1,160 @@
+import os
+import shutil
+import traceback
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import ai.chronon.cli.compile.display.compiled_obj
+import ai.chronon.cli.compile.parse_configs as parser
+import ai.chronon.cli.logger as logger
+from ai.chronon.cli.compile import serializer
+from ai.chronon.cli.compile.compile_context import CompileContext, ConfigInfo
+from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
+from ai.chronon.cli.compile.display.console import console
+from ai.chronon.cli.compile.parse_teams import merge_team_execution_info
+from ai.chronon.orchestration.ttypes import ConfType
+from ai.chronon.types import MetaData
+
+logger = logger.get_logger()
+
+
+@dataclass
+class CompileResult:
+    config_info: ConfigInfo
+    obj_dict: Dict[str, Any]
+    error_dict: Dict[str, List[BaseException]]
+
+
+class Compiler:
+
+    def __init__(self, compile_context: CompileContext):
+        self.compile_context = compile_context
+
+    def compile(self) -> Dict[ConfType, CompileResult]:
+
+        config_infos = self.compile_context.config_infos
+
+        compile_results = {}
+
+        for config_info in config_infos:
+            configs = self._compile_class_configs(config_info)
+
+            compile_results[config_info.config_type] = configs
+        self._compile_team_metadata()
+
+        # check if staging_output_dir exists
+        staging_dir = self.compile_context.staging_output_dir()
+        if os.path.exists(staging_dir):
+            # replace staging_output_dir to output_dir
+            output_dir = self.compile_context.output_dir()
+            if os.path.exists(output_dir):
+                shutil.rmtree(output_dir)
+            shutil.move(staging_dir, output_dir)
+        else:
+            print(
+                f"Staging directory {staging_dir} does not exist. "
+                "Happens when every chronon config fails to compile or when no chronon configs exist."
+            )
+
+        # TODO: temporarily just print out the final results of the compile until live fix is implemented:
+        #  https://github.com/Textualize/rich/pull/3637
+        console.print(self.compile_context.compile_status.render())
+
+        return compile_results
+
+    def _compile_team_metadata(self):
+        """
+        Compile the team metadata and return the compiled object.
+        """
+        teams_dict = self.compile_context.teams_dict
+        for team in teams_dict:
+            m = MetaData()
+            merge_team_execution_info(m, teams_dict, team)
+
+            tjson = serializer.thrift_simple_json(m)
+            name = f"{team}.{team}_team_metadata"
+            result = CompiledObj(
+                name=name,
+                obj=m,
+                file=name,
+                errors=None,
+                obj_type=MetaData.__name__,
+                tjson=tjson,
+            )
+            self._write_object(result)
+            self.compile_context.compile_status.add_object_update_display(result, MetaData.__name__)
+
+        # Done writing team metadata, close the class
+        self.compile_context.compile_status.close_cls(MetaData.__name__)
+
+    def _compile_class_configs(self, config_info: ConfigInfo) -> CompileResult:
+
+        compile_result = CompileResult(
+            config_info=config_info, obj_dict={}, error_dict={}
+        )
+
+        input_dir = self.compile_context.input_dir(config_info.cls)
+
+        compiled_objects = parser.from_folder(
+            config_info.cls, input_dir, self.compile_context
+        )
+
+        objects, errors = self._write_objects_in_folder(compiled_objects)
+
+        if objects:
+            compile_result.obj_dict.update(objects)
+
+        if errors:
+            compile_result.error_dict.update(errors)
+
+        self.compile_context.compile_status.close_cls(config_info.cls.__name__)
+
+        return compile_result
+
+    def _write_objects_in_folder(
+        self,
+        compiled_objects: List[ai.chronon.cli.compile.display.compiled_obj.CompiledObj],
+    ) -> Tuple[Dict[str, Any], Dict[str, List[BaseException]]]:
+
+        error_dict = {}
+        object_dict = {}
+
+        for co in compiled_objects:
+
+            if co.obj:
+
+                if co.errors:
+                    error_dict[co.name] = co.errors
+
+                    for error in co.errors:
+                        self.compile_context.compile_status.print_live_console(
+                            f"Error processing conf {co.name}: {error}"
+                        )
+                        traceback.print_exception(
+                            type(error), error, error.__traceback__
+                        )
+
+                else:
+                    self._write_object(co)
+                    object_dict[co.name] = co.obj
+            else:
+                error_dict[co.file] = co.errors
+
+                self.compile_context.compile_status.print_live_console(
+                    f"Error processing file {co.file}: {co.errors}"
+                )
+                for error in co.errors:
+                    traceback.print_exception(type(error), error, error.__traceback__)
+
+        return object_dict, error_dict
+
+    def _write_object(self, compiled_obj: CompiledObj) -> Optional[List[BaseException]]:
+        output_path = self.compile_context.staging_output_path(compiled_obj)
+
+        folder = os.path.dirname(output_path)
+
+        if not os.path.exists(folder):
+            os.makedirs(folder)
+
+        with open(output_path, "w") as f:
+            f.write(compiled_obj.tjson)