churnkit 0.75.0a1__py3-none-any.whl

customer_retention/generators/spec_generator/__init__.py

@@ -0,0 +1,37 @@
+from .databricks_generator import DatabricksSpecGenerator
+from .generic_generator import GenericSpecGenerator
+from .mlflow_pipeline_generator import (
+    CleanAction,
+    MLflowConfig,
+    MLflowPipelineGenerator,
+    RecommendationParser,
+    TransformAction,
+)
+from .pipeline_spec import (
+    ColumnSpec,
+    FeatureSpec,
+    ModelSpec,
+    PipelineSpec,
+    QualityGateSpec,
+    SchemaSpec,
+    SourceSpec,
+    TransformSpec,
+)
+
+__all__ = [
+    "PipelineSpec",
+    "SourceSpec",
+    "SchemaSpec",
+    "ColumnSpec",
+    "TransformSpec",
+    "FeatureSpec",
+    "ModelSpec",
+    "QualityGateSpec",
+    "GenericSpecGenerator",
+    "DatabricksSpecGenerator",
+    "MLflowPipelineGenerator",
+    "MLflowConfig",
+    "RecommendationParser",
+    "CleanAction",
+    "TransformAction",
+]

customer_retention/generators/spec_generator/databricks_generator.py

@@ -0,0 +1,433 @@
+import json
+from pathlib import Path
+from typing import Any, Dict, List
+
+from .pipeline_spec import PipelineSpec
+
+
+class DatabricksSpecGenerator:
+    def __init__(self,
+                 catalog: str = "main",
+                 schema: str = "default",
+                 output_dir: str = "./databricks_artifacts"):
+        self.catalog = catalog
+        self.schema = schema
+        self.output_dir = Path(output_dir)
+
+    def generate_lakeflow_connect(self, spec: PipelineSpec) -> dict:
+        sources = []
+        for source in spec.sources:
+            sources.append({
+                "name": source.name,
+                "type": self._infer_connector_type(source.format),
+                "path": source.path,
+                "options": source.options
+            })
+
+        return {
+            "ingestion": {
+                "name": f"{spec.name}_ingestion",
+                "version": spec.version,
+                "target_catalog": self.catalog,
+                "target_schema": self.schema,
+                "sources": sources,
+                "target_tables": [
+                    {
+                        "name": f"{spec.name}_bronze",
+                        "format": "delta",
+                        "mode": "append"
+                    }
+                ],
+                "schedule": {
+                    "quartz_cron_expression": "0 0 * * * ?",
+                    "timezone_id": "UTC"
+                }
+            }
+        }
+
+    def _infer_connector_type(self, format: str) -> str:
+        format_map = {
+            "csv": "file",
+            "parquet": "file",
+            "json": "file",
+            "delta": "delta",
+            "jdbc": "jdbc",
+            "kafka": "kafka"
+        }
+        return format_map.get(format.lower(), "file")
+
+    def generate_dlt_pipeline(self, spec: PipelineSpec) -> str:
+        lines = [
+            "import dlt",
+            "from pyspark.sql import functions as F",
+            "",
+            f"# DLT Pipeline: {spec.name}",
+            f"# Version: {spec.version}",
+            "",
+        ]
+
+        lines.extend(self._generate_bronze_tables(spec))
+        lines.append("")
+        lines.extend(self._generate_silver_tables(spec))
+        lines.append("")
+        lines.extend(self._generate_gold_tables(spec))
+
+        return "\n".join(lines)
+
+    def _generate_bronze_tables(self, spec: PipelineSpec) -> List[str]:
+        lines = ["# Bronze Layer - Raw Data Ingestion", ""]
+
+        for source in spec.sources:
+            table_name = f"{spec.name}_bronze"
+            lines.extend([
+                '@dlt.table(',
+                f'    name="{table_name}",',
+                f'    comment="Raw data from {source.name}"',
+                ')',
+                '@dlt.expect_or_drop("valid_record", "1=1")',
+                f'def {table_name}():',
+                '    return (',
+                '        spark.read',
+                f'        .format("{source.format}")',
+                f'        .load("{source.path}")',
+                '    )',
+                ""
+            ])
+
+        return lines
+
+    def _generate_silver_tables(self, spec: PipelineSpec) -> List[str]:
+        lines = ["# Silver Layer - Cleaned and Standardized", ""]
+
+        table_name = f"{spec.name}_silver"
+        bronze_table = f"{spec.name}_bronze"
+
+        expectations = []
+        for gate in spec.quality_gates:
+            if gate.gate_type == "null_percentage":
+                expectations.append(
+                    f'@dlt.expect_or_warn("{gate.name}", "{gate.column} IS NOT NULL")'
+                )
+
+        lines.extend([
+            '@dlt.table(',
+            f'    name="{table_name}",',
+            '    comment="Cleaned and standardized data"',
+            ')',
+        ])
+        for exp in expectations:
+            lines.append(exp)
+
+        lines.extend([
+            f'def {table_name}():',
+            f'    df = dlt.read("{bronze_table}")',
+            ""
+        ])
+
+        if spec.silver_transforms:
+            for transform in spec.silver_transforms:
+                if transform.transform_type == "standard_scaling":
+                    col = transform.input_columns[0]
+                    out_col = transform.output_columns[0]
+                    lines.append(f'    # {transform.name}')
+                    lines.append(f'    df = df.withColumn("{out_col}", F.col("{col}"))')
+                elif transform.transform_type == "one_hot_encoding":
+                    col = transform.input_columns[0]
+                    lines.append(f'    # {transform.name}')
+                    lines.append('    # Note: One-hot encoding applied in feature engineering')
+
+            lines.append("")
+
+        lines.extend([
+            '    return df',
+            ""
+        ])
+
+        return lines
+
+    def _generate_gold_tables(self, spec: PipelineSpec) -> List[str]:
+        lines = ["# Gold Layer - Feature Engineering", ""]
+
+        table_name = f"{spec.name}_gold"
+        silver_table = f"{spec.name}_silver"
+
+        lines.extend([
+            '@dlt.table(',
+            f'    name="{table_name}",',
+            '    comment="Feature-engineered data ready for modeling"',
+            ')',
+            f'def {table_name}():',
+            f'    df = dlt.read("{silver_table}")',
+            ""
+        ])
+
+        if spec.feature_definitions:
+            for feature in spec.feature_definitions:
+                lines.append(f'    # Feature: {feature.name}')
+                if feature.computation == "days_since_today":
+                    col = feature.source_columns[0]
+                    lines.append(
+                        f'    df = df.withColumn("{feature.name}", '
+                        f'F.datediff(F.current_date(), F.col("{col}")))'
+                    )
+
+            lines.append("")
+
+        lines.extend([
+            '    return df',
+            ""
+        ])
+
+        return lines
+
+    def generate_workflow_jobs(self, spec: PipelineSpec) -> dict:
+        tasks = [
+            {
+                "task_key": "run_dlt_pipeline",
+                "pipeline_task": {
+                    "pipeline_id": f"{{{{pipelines.{spec.name}_dlt.id}}}}"
+                }
+            },
+            {
+                "task_key": "train_model",
+                "depends_on": [{"task_key": "run_dlt_pipeline"}],
+                "notebook_task": {
+                    "notebook_path": f"/Repos/{spec.name}/notebooks/train_model",
+                    "base_parameters": {
+                        "catalog": self.catalog,
+                        "schema": self.schema,
+                        "table": f"{spec.name}_gold"
+                    }
+                }
+            },
+            {
+                "task_key": "validate_model",
+                "depends_on": [{"task_key": "train_model"}],
+                "notebook_task": {
+                    "notebook_path": f"/Repos/{spec.name}/notebooks/validate_model"
+                }
+            }
+        ]
+
+        return {
+            "name": f"{spec.name}_workflow",
+            "tasks": tasks,
+            "schedule": {
+                "quartz_cron_expression": "0 0 0 * * ?",
+                "timezone_id": "UTC",
+                "pause_status": "UNPAUSED"
+            },
+            "email_notifications": {
+                "on_failure": []
+            },
+            "max_concurrent_runs": 1,
+            "trigger": {
+                "periodic": {
+                    "interval": 1,
+                    "unit": "DAYS"
+                }
+            }
+        }
+
+    def generate_feature_tables(self, spec: PipelineSpec) -> str:
+        primary_key = None
+        if spec.schema and spec.schema.primary_key:
+            primary_key = spec.schema.primary_key
+
+        lines = [
+            "from databricks.feature_store import FeatureStoreClient",
+            "from databricks.feature_store import FeatureLookup",
+            "from pyspark.sql import functions as F",
+            "",
+            f"# Feature Store Tables for {spec.name}",
+            "",
+            "fs = FeatureStoreClient()",
+            "",
+            '# Create feature table',
+            f'feature_table_name = "{self.catalog}.{self.schema}.{spec.name}_features"',
+            ""
+        ]
+
+        if primary_key:
+            lines.extend([
+                f'# Define the feature table with primary key: {primary_key}',
+                'fs.create_table(',
+                '    name=feature_table_name,',
+                f'    primary_keys=["{primary_key}"],',
+                f'    df=spark.table("{self.catalog}.{self.schema}.{spec.name}_gold"),',
+                f'    description="Features for {spec.name}"',
+                ')',
+                ""
+            ])
+        else:
+            lines.extend([
+                '# Write features to table',
+                f'df = spark.table("{self.catalog}.{self.schema}.{spec.name}_gold")',
+                'fs.write_table(',
+                '    name=feature_table_name,',
+                '    df=df,',
+                '    mode="overwrite"',
+                ')',
+                ""
+            ])
+
+        if spec.feature_definitions:
+            lines.extend([
+                "# Feature lookups for training",
+                "feature_lookups = ["
+            ])
+            for feature in spec.feature_definitions:
+                lines.append('    FeatureLookup(')
+                lines.append('        table_name=feature_table_name,')
+                lines.append(f'        feature_names=["{feature.name}"],')
+                if primary_key:
+                    lines.append(f'        lookup_key="{primary_key}"')
+                lines.append('    ),')
+            lines.append("]")
+
+        return "\n".join(lines)
+
+    def generate_mlflow_experiment(self, spec: PipelineSpec) -> str:
+        target = spec.model_config.target_column if spec.model_config else "target"
+        model_type = spec.model_config.model_type if spec.model_config else "gradient_boosting"
+        model_name = spec.model_config.name if spec.model_config else "model"
+
+        lines = [
+            "import mlflow",
+            "import mlflow.spark",
+            "from pyspark.ml.classification import GBTClassifier, RandomForestClassifier",
+            "from pyspark.ml.evaluation import BinaryClassificationEvaluator",
+            "",
+            f"# MLflow Experiment: {spec.name}",
+            "",
+            f'mlflow.set_experiment("/Users/{{username}}/{spec.name}_experiment")',
+            "",
+            f'with mlflow.start_run(run_name="{model_name}"):',
+            '    # Load training data',
+            f'    df = spark.table("{self.catalog}.{self.schema}.{spec.name}_gold")',
+            "",
+            '    # Log parameters',
+            f'    mlflow.log_param("target_column", "{target}")',
+            f'    mlflow.log_param("model_type", "{model_type}")',
+            ""
+        ]
+
+        if spec.model_config and spec.model_config.hyperparameters:
+            lines.append("    # Log hyperparameters")
+            for key, value in spec.model_config.hyperparameters.items():
+                lines.append(f'    mlflow.log_param("{key}", {repr(value)})')
+            lines.append("")
+
+        lines.extend([
+            '    # Train model',
+            '    model = GBTClassifier(',
+            '        featuresCol="features",',
+            f'        labelCol="{target}",',
+            '        maxIter=100',
+            '    )',
+            "",
+            '    trained_model = model.fit(df)',
+            "",
+            '    # Log model',
+            f'    mlflow.spark.log_model(trained_model, "{model_name}")',
+            "",
+            '    # Register model in Unity Catalog',
+            '    mlflow.register_model(',
+            f'        f"runs/{{mlflow.active_run().info.run_id}}/{model_name}",',
+            f'        "{self.catalog}.{self.schema}.{spec.name}_model"',
+            '    )',
+        ])
+
+        return "\n".join(lines)
+
+    def generate_unity_catalog_schema(self, spec: PipelineSpec) -> str:
+        lines = [
+            f"-- Unity Catalog Schema for {spec.name}",
+            f"-- Generated from PipelineSpec version {spec.version}",
+            "",
+            f"CREATE SCHEMA IF NOT EXISTS {self.catalog}.{self.schema};",
+            "",
+        ]
+
+        for layer in ["bronze", "silver", "gold"]:
+            table_name = f"{spec.name}_{layer}"
+            lines.extend([
+                f"CREATE OR REPLACE TABLE {self.catalog}.{self.schema}.{table_name} (",
+            ])
+
+            if spec.schema:
+                col_defs = []
+                for col in spec.schema.columns:
+                    spark_type = self._to_spark_type(col.data_type)
+                    nullable = "" if col.nullable else " NOT NULL"
+                    col_defs.append(f"    {col.name} {spark_type}{nullable}")
+
+                lines.append(",\n".join(col_defs))
+
+            lines.extend([
+                ")",
+                "USING DELTA",
+                f"COMMENT '{layer.title()} layer table for {spec.name}';",
+                ""
+            ])
+
+        return "\n".join(lines)
+
+    def _to_spark_type(self, data_type: str) -> str:
+        type_map = {
+            "string": "STRING",
+            "integer": "INT",
+            "float": "DOUBLE",
+            "timestamp": "TIMESTAMP",
+            "date": "DATE",
+            "boolean": "BOOLEAN"
+        }
+        return type_map.get(data_type.lower(), "STRING")
+
+    def generate_all(self, spec: PipelineSpec) -> Dict[str, Any]:
+        return {
+            "lakeflow_connect": self.generate_lakeflow_connect(spec),
+            "dlt_pipeline": self.generate_dlt_pipeline(spec),
+            "workflow_jobs": self.generate_workflow_jobs(spec),
+            "feature_tables": self.generate_feature_tables(spec),
+            "mlflow_experiment": self.generate_mlflow_experiment(spec),
+            "unity_catalog_schema": self.generate_unity_catalog_schema(spec)
+        }
+
+    def save_all(self, spec: PipelineSpec) -> List[str]:
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        saved_files = []
+
+        artifacts = self.generate_all(spec)
+
+        lakeflow_path = self.output_dir / f"{spec.name}_lakeflow_connect.json"
+        with open(lakeflow_path, "w") as f:
+            json.dump(artifacts["lakeflow_connect"], f, indent=2)
+        saved_files.append(str(lakeflow_path))
+
+        dlt_path = self.output_dir / f"{spec.name}_dlt_pipeline.py"
+        with open(dlt_path, "w") as f:
+            f.write(artifacts["dlt_pipeline"])
+        saved_files.append(str(dlt_path))
+
+        jobs_path = self.output_dir / f"{spec.name}_workflow_jobs.json"
+        with open(jobs_path, "w") as f:
+            json.dump(artifacts["workflow_jobs"], f, indent=2)
+        saved_files.append(str(jobs_path))
+
+        features_path = self.output_dir / f"{spec.name}_feature_tables.py"
+        with open(features_path, "w") as f:
+            f.write(artifacts["feature_tables"])
+        saved_files.append(str(features_path))
+
+        mlflow_path = self.output_dir / f"{spec.name}_mlflow_experiment.py"
+        with open(mlflow_path, "w") as f:
+            f.write(artifacts["mlflow_experiment"])
+        saved_files.append(str(mlflow_path))
+
+        schema_path = self.output_dir / f"{spec.name}_unity_catalog.sql"
+        with open(schema_path, "w") as f:
+            f.write(artifacts["unity_catalog_schema"])
+        saved_files.append(str(schema_path))
+
+        return saved_files