churnkit 0.75.0a1__py3-none-any.whl

customer_retention/integrations/adapters/mlflow/__init__.py

@@ -0,0 +1,6 @@
+from .base import MLflowAdapter
+from .databricks import DatabricksMLflow
+from .experiment_tracker import ExperimentTracker
+from .local import LocalMLflow
+
+__all__ = ["MLflowAdapter", "LocalMLflow", "DatabricksMLflow", "ExperimentTracker"]

customer_retention/integrations/adapters/mlflow/base.py

@@ -0,0 +1,32 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+
+
+class MLflowAdapter(ABC):
+    @abstractmethod
+    def start_run(self, experiment_name: str, run_name: Optional[str] = None) -> str:
+        pass
+
+    @abstractmethod
+    def end_run(self) -> None:
+        pass
+
+    @abstractmethod
+    def log_params(self, params: Dict[str, Any]) -> None:
+        pass
+
+    @abstractmethod
+    def log_metrics(self, metrics: Dict[str, float]) -> None:
+        pass
+
+    @abstractmethod
+    def log_model(self, model: Any, artifact_path: str, registered_name: Optional[str] = None) -> str:
+        pass
+
+    @abstractmethod
+    def load_model(self, model_uri: str) -> Any:
+        pass
+
+    @abstractmethod
+    def transition_stage(self, model_name: str, version: str, stage: str) -> None:
+        pass

customer_retention/integrations/adapters/mlflow/databricks.py

@@ -0,0 +1,54 @@
+from typing import Any, Dict, Optional
+
+from customer_retention.core.compat.detection import is_spark_available
+
+from .base import MLflowAdapter
+
+try:
+    import mlflow
+    from mlflow.tracking import MlflowClient
+    MLFLOW_AVAILABLE = True
+except ImportError:
+    MLFLOW_AVAILABLE = False
+
+
+class DatabricksMLflow(MLflowAdapter):
+    def __init__(self, registry_uri: str = "databricks-uc"):
+        if not is_spark_available():
+            raise ImportError("PySpark required for DatabricksMLflow")
+        if not MLFLOW_AVAILABLE:
+            raise ImportError("mlflow package required")
+        mlflow.set_registry_uri(registry_uri)
+        self.registry_uri = registry_uri
+        self._client = MlflowClient()
+        self._run_id = None
+
+    def start_run(self, experiment_name: str, run_name: Optional[str] = None) -> str:
+        experiment = mlflow.get_experiment_by_name(experiment_name)
+        if experiment is None:
+            experiment_id = mlflow.create_experiment(experiment_name)
+        else:
+            experiment_id = experiment.experiment_id
+        run = mlflow.start_run(experiment_id=experiment_id, run_name=run_name)
+        self._run_id = run.info.run_id
+        return self._run_id
+
+    def end_run(self) -> None:
+        mlflow.end_run()
+        self._run_id = None
+
+    def log_params(self, params: Dict[str, Any]) -> None:
+        mlflow.log_params(params)
+
+    def log_metrics(self, metrics: Dict[str, float]) -> None:
+        mlflow.log_metrics(metrics)
+
+    def log_model(self, model: Any, artifact_path: str, registered_name: Optional[str] = None) -> str:
+        info = mlflow.sklearn.log_model(model, artifact_path, registered_model_name=registered_name)
+        return info.model_uri
+
+    def load_model(self, model_uri: str) -> Any:
+        return mlflow.sklearn.load_model(model_uri)
+
+    def transition_stage(self, model_name: str, version: str, stage: str) -> None:
+        self._client.set_model_version_tag(name=model_name, version=version, key="stage", value=stage)

customer_retention/integrations/adapters/mlflow/experiment_tracker.py

@@ -0,0 +1,161 @@
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+try:
+    import mlflow
+    from mlflow.tracking import MlflowClient
+    MLFLOW_AVAILABLE = True
+except ImportError:
+    MLFLOW_AVAILABLE = False
+
+if TYPE_CHECKING:
+    from customer_retention.analysis.auto_explorer.findings import ExplorationFindings
+    from customer_retention.analysis.recommendations.pipeline import RecommendationPipeline
+
+
+class ExperimentTracker:
+    def __init__(self, tracking_uri: str = "./mlruns", experiment_name: str = "customer_retention"):
+        if not MLFLOW_AVAILABLE:
+            raise ImportError("mlflow package required. Install with: uv sync --extra ml")
+        mlflow.set_tracking_uri(tracking_uri)
+        self.tracking_uri = tracking_uri
+        self.experiment_name = experiment_name
+        self._client = MlflowClient(tracking_uri=tracking_uri)
+        self._ensure_experiment()
+
+    def log_exploration(self, findings: "ExplorationFindings", run_name: Optional[str] = None) -> str:
+        with mlflow.start_run(
+            run_name=run_name or f"exploration_{Path(findings.source_path).stem}",
+            experiment_id=self._ensure_experiment()
+        ) as run:
+            self._log_exploration_params(findings)
+            self._log_exploration_metrics(findings)
+            self._log_column_metrics(findings)
+            mlflow.log_dict(findings.to_dict(), "exploration_findings.json")
+            self._set_exploration_tags(findings)
+            return run.info.run_id
+
+    def log_pipeline_execution(
+        self, pipeline: "RecommendationPipeline", run_name: Optional[str] = None,
+        parent_run_id: Optional[str] = None
+    ) -> str:
+        with mlflow.start_run(
+            run_name=run_name or "recommendation_pipeline",
+            experiment_id=self._ensure_experiment(),
+            nested=parent_run_id is not None
+        ) as run:
+            self._log_pipeline_params(pipeline)
+            self._log_pipeline_metrics(pipeline)
+            self._log_pipeline_artifacts(pipeline)
+            mlflow.set_tags({"stage": "transformation", "pipeline_fitted": str(pipeline._is_fitted)})
+            return run.info.run_id
+
+    def log_model_training(
+        self, model: Any, metrics: Dict[str, float], params: Dict[str, Any],
+        model_name: str = "churn_model", run_name: Optional[str] = None
+    ) -> str:
+        with mlflow.start_run(
+            run_name=run_name or f"training_{model_name}",
+            experiment_id=self._ensure_experiment()
+        ) as run:
+            mlflow.log_params(params)
+            mlflow.log_metrics(metrics)
+            mlflow.sklearn.log_model(model, "model", registered_model_name=model_name)
+            mlflow.set_tags({"stage": "training", "model_name": model_name})
+            return run.info.run_id
+
+    def get_best_run(self, metric: str = "overall_quality_score", ascending: bool = False) -> Optional[Dict]:
+        experiment = mlflow.get_experiment_by_name(self.experiment_name)
+        if experiment is None:
+            return None
+        runs = self._client.search_runs(
+            experiment_ids=[experiment.experiment_id],
+            order_by=[f"metrics.{metric} {'ASC' if ascending else 'DESC'}"],
+            max_results=1
+        )
+        return runs[0].to_dictionary() if runs else None
+
+    def list_exploration_runs(self) -> List[Dict]:
+        experiment = mlflow.get_experiment_by_name(self.experiment_name)
+        if experiment is None:
+            return []
+        runs = self._client.search_runs(
+            experiment_ids=[experiment.experiment_id],
+            filter_string="tags.stage = 'exploration'"
+        )
+        return [r.to_dictionary() for r in runs]
+
+    @staticmethod
+    def serve_ui(host: str = "127.0.0.1", port: int = 5000, tracking_uri: str = "./mlruns"):
+        import subprocess
+        import sys
+        subprocess.run([
+            sys.executable, "-m", "mlflow", "ui",
+            "--backend-store-uri", tracking_uri, "--host", host, "--port", str(port)
+        ])
+
+    def _ensure_experiment(self) -> str:
+        experiment = mlflow.get_experiment_by_name(self.experiment_name)
+        if experiment is None:
+            return mlflow.create_experiment(self.experiment_name)
+        return experiment.experiment_id
+
+    def _log_exploration_params(self, findings: "ExplorationFindings") -> None:
+        mlflow.log_params({
+            "source_path": findings.source_path,
+            "source_format": findings.source_format,
+            "target_column": findings.target_column or "none",
+        })
+
+    def _log_exploration_metrics(self, findings: "ExplorationFindings") -> None:
+        mlflow.log_metrics({
+            "row_count": findings.row_count,
+            "column_count": findings.column_count,
+            "memory_usage_mb": findings.memory_usage_mb,
+            "overall_quality_score": findings.overall_quality_score,
+            "modeling_ready": 1.0 if findings.modeling_ready else 0.0,
+            "critical_issues_count": len(findings.critical_issues),
+            "warnings_count": len(findings.warnings),
+        })
+
+    def _log_column_metrics(self, findings: "ExplorationFindings") -> None:
+        type_counts: Dict[str, int] = {}
+        cleaning_needed_count = 0
+        for col in findings.columns.values():
+            type_name = col.inferred_type.value
+            type_counts[type_name] = type_counts.get(type_name, 0) + 1
+            if col.cleaning_needed:
+                cleaning_needed_count += 1
+        for type_name, count in type_counts.items():
+            mlflow.log_metric(f"columns_{type_name}", count)
+        mlflow.log_metric("columns_needing_cleaning", cleaning_needed_count)
+
+    def _set_exploration_tags(self, findings: "ExplorationFindings") -> None:
+        mlflow.set_tags({
+            "stage": "exploration",
+            "modeling_ready": str(findings.modeling_ready),
+            "is_time_series": str(findings.is_time_series),
+        })
+
+    def _log_pipeline_params(self, pipeline: "RecommendationPipeline") -> None:
+        mlflow.log_params({
+            "recommendation_count": len(pipeline.recommendations),
+            "is_fitted": pipeline._is_fitted,
+        })
+
+    def _log_pipeline_metrics(self, pipeline: "RecommendationPipeline") -> None:
+        rec_types: Dict[str, int] = {}
+        rec_categories: Dict[str, int] = {}
+        for rec in pipeline.recommendations:
+            rec_types[rec.recommendation_type] = rec_types.get(rec.recommendation_type, 0) + 1
+            rec_categories[rec.category] = rec_categories.get(rec.category, 0) + 1
+        for rec_type, count in rec_types.items():
+            mlflow.log_metric(f"rec_type_{rec_type}", count)
+        for category, count in rec_categories.items():
+            mlflow.log_metric(f"rec_category_{category}", count)
+
+    def _log_pipeline_artifacts(self, pipeline: "RecommendationPipeline") -> None:
+        from customer_retention.analysis.recommendations.base import Platform
+        mlflow.log_dict(pipeline.to_dict(), "pipeline_config.json")
+        mlflow.log_text(pipeline.generate_code(), "generated_code_local.py")
+        mlflow.log_text(pipeline.generate_code(Platform.DATABRICKS), "generated_code_databricks.py")

customer_retention/integrations/adapters/mlflow/local.py

@@ -0,0 +1,50 @@
+from typing import Any, Dict, Optional
+
+from .base import MLflowAdapter
+
+try:
+    import mlflow
+    from mlflow.tracking import MlflowClient
+    MLFLOW_AVAILABLE = True
+except ImportError:
+    MLFLOW_AVAILABLE = False
+
+
+class LocalMLflow(MLflowAdapter):
+    def __init__(self, tracking_uri: str = "./mlruns"):
+        if not MLFLOW_AVAILABLE:
+            raise ImportError("mlflow package required: pip install mlflow")
+        mlflow.set_tracking_uri(tracking_uri)
+        self.tracking_uri = tracking_uri
+        self._client = MlflowClient(tracking_uri=tracking_uri)
+        self._run_id = None
+
+    def start_run(self, experiment_name: str, run_name: Optional[str] = None) -> str:
+        experiment = mlflow.get_experiment_by_name(experiment_name)
+        if experiment is None:
+            experiment_id = mlflow.create_experiment(experiment_name)
+        else:
+            experiment_id = experiment.experiment_id
+        run = mlflow.start_run(experiment_id=experiment_id, run_name=run_name)
+        self._run_id = run.info.run_id
+        return self._run_id
+
+    def end_run(self) -> None:
+        mlflow.end_run()
+        self._run_id = None
+
+    def log_params(self, params: Dict[str, Any]) -> None:
+        mlflow.log_params(params)
+
+    def log_metrics(self, metrics: Dict[str, float]) -> None:
+        mlflow.log_metrics(metrics)
+
+    def log_model(self, model: Any, artifact_path: str, registered_name: Optional[str] = None) -> str:
+        info = mlflow.sklearn.log_model(model, artifact_path, registered_model_name=registered_name)
+        return info.model_uri
+
+    def load_model(self, model_uri: str) -> Any:
+        return mlflow.sklearn.load_model(model_uri)
+
+    def transition_stage(self, model_name: str, version: str, stage: str) -> None:
+        self._client.transition_model_version_stage(name=model_name, version=version, stage=stage)

customer_retention/integrations/adapters/storage/__init__.py

@@ -0,0 +1,5 @@
+from .base import DeltaStorage
+from .databricks import DatabricksDelta
+from .local import LocalDelta
+
+__all__ = ["DeltaStorage", "LocalDelta", "DatabricksDelta"]

customer_retention/integrations/adapters/storage/base.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+
+
+class DeltaStorage(ABC):
+    @abstractmethod
+    def read(self, path: str, version: Optional[int] = None) -> pd.DataFrame:
+        pass
+
+    @abstractmethod
+    def write(self, df: pd.DataFrame, path: str, mode: str = "overwrite",
+              partition_by: Optional[List[str]] = None,
+              metadata: Optional[Dict[str, str]] = None) -> None:
+        pass
+
+    @abstractmethod
+    def merge(self, df: pd.DataFrame, path: str, condition: str,
+              update_cols: Optional[List[str]] = None) -> None:
+        pass
+
+    @abstractmethod
+    def history(self, path: str) -> List[Dict[str, Any]]:
+        pass
+
+    @abstractmethod
+    def vacuum(self, path: str, retention_hours: int = 168) -> None:
+        pass
+
+    @abstractmethod
+    def exists(self, path: str) -> bool:
+        pass

customer_retention/integrations/adapters/storage/databricks.py

@@ -0,0 +1,76 @@
+import json
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+
+from customer_retention.core.compat.detection import get_spark_session, is_spark_available
+
+from .base import DeltaStorage
+
+
+class DatabricksDelta(DeltaStorage):
+    def __init__(self):
+        if not is_spark_available():
+            raise ImportError("PySpark required for DatabricksDelta")
+        self._spark = None
+
+    @property
+    def spark(self) -> Any:
+        if self._spark is None:
+            self._spark = get_spark_session()
+            if self._spark is None:
+                from pyspark.sql import SparkSession
+                self._spark = SparkSession.builder.getOrCreate()
+        return self._spark
+
+    def read(self, path: str, version: Optional[int] = None) -> pd.DataFrame:
+        reader = self.spark.read.format("delta")
+        if version is not None:
+            reader = reader.option("versionAsOf", version)
+        return reader.load(path).toPandas()
+
+    def write(self, df: pd.DataFrame, path: str, mode: str = "overwrite",
+              partition_by: Optional[List[str]] = None,
+              metadata: Optional[Dict[str, str]] = None) -> None:
+        if metadata:
+            self.spark.conf.set(
+                "spark.databricks.delta.commitInfo.userMetadata",
+                json.dumps(metadata),
+            )
+        spark_df = self.spark.createDataFrame(df)
+        writer = spark_df.write.format("delta").mode(mode)
+        if partition_by:
+            writer = writer.partitionBy(*partition_by)
+        writer.save(path)
+
+    def merge(self, df: pd.DataFrame, path: str, condition: str,
+              update_cols: Optional[List[str]] = None) -> None:
+        from delta.tables import DeltaTable
+        spark_df = self.spark.createDataFrame(df)
+        target = DeltaTable.forPath(self.spark, path)
+        merge_builder = target.alias("target").merge(spark_df.alias("source"), condition)
+        if update_cols:
+            update_dict = {col: f"source.{col}" for col in update_cols}
+            merge_builder = merge_builder.whenMatchedUpdate(set=update_dict)
+        else:
+            merge_builder = merge_builder.whenMatchedUpdateAll()
+        merge_builder.whenNotMatchedInsertAll().execute()
+
+    def history(self, path: str) -> List[Dict[str, Any]]:
+        from delta.tables import DeltaTable
+        dt = DeltaTable.forPath(self.spark, path)
+        history_df = dt.history()
+        return history_df.toPandas().to_dict("records")
+
+    def vacuum(self, path: str, retention_hours: int = 168) -> None:
+        from delta.tables import DeltaTable
+        dt = DeltaTable.forPath(self.spark, path)
+        dt.vacuum(retention_hours)
+
+    def exists(self, path: str) -> bool:
+        from delta.tables import DeltaTable
+        try:
+            DeltaTable.forPath(self.spark, path)
+            return True
+        except Exception:
+            return False

customer_retention/integrations/adapters/storage/local.py

@@ -0,0 +1,59 @@
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+
+from .base import DeltaStorage
+
+try:
+    import deltalake
+    from deltalake import DeltaTable, write_deltalake
+    DELTA_RS_AVAILABLE = True
+except ImportError:
+    DELTA_RS_AVAILABLE = False
+
+
+class LocalDelta(DeltaStorage):
+    def __init__(self):
+        if not DELTA_RS_AVAILABLE:
+            raise ImportError("deltalake package required: pip install deltalake")
+
+    def read(self, path: str, version: Optional[int] = None) -> pd.DataFrame:
+        if version is not None:
+            dt = DeltaTable(path, version=version)
+        else:
+            dt = DeltaTable(path)
+        return dt.to_pandas()
+
+    def write(self, df: pd.DataFrame, path: str, mode: str = "overwrite",
+              partition_by: Optional[List[str]] = None,
+              metadata: Optional[Dict[str, str]] = None) -> None:
+        kwargs = {"mode": mode}
+        if partition_by:
+            kwargs["partition_by"] = partition_by
+        if metadata:
+            from deltalake import CommitProperties
+            kwargs["commit_properties"] = CommitProperties(custom_metadata=metadata)
+        write_deltalake(path, df, **kwargs)
+
+    def merge(self, df: pd.DataFrame, path: str, condition: str,
+              update_cols: Optional[List[str]] = None) -> None:
+        dt = DeltaTable(path)
+        merge_builder = dt.merge(df, predicate=condition, source_alias="source", target_alias="target")
+        if update_cols:
+            update_dict = {col: f"source.{col}" for col in update_cols}
+            merge_builder = merge_builder.when_matched_update(updates=update_dict)
+        else:
+            merge_builder = merge_builder.when_matched_update_all()
+        merge_builder.when_not_matched_insert_all().execute()
+
+    def history(self, path: str) -> List[Dict[str, Any]]:
+        dt = DeltaTable(path)
+        return dt.history()
+
+    def vacuum(self, path: str, retention_hours: int = 168) -> None:
+        dt = DeltaTable(path)
+        dt.vacuum(retention_hours=retention_hours, enforce_retention_duration=False, dry_run=False)
+
+    def exists(self, path: str) -> bool:
+        from pathlib import Path
+        return Path(path).joinpath("_delta_log").is_dir()

customer_retention/integrations/feature_store/__init__.py

@@ -0,0 +1,47 @@
+"""Feature store module for leakage-safe feature management.
+
+This module provides a unified interface for managing ML features with
+point-in-time correctness. It supports both local development (Feast)
+and production (Databricks Feature Engineering) backends.
+
+Key Components:
+    - TemporalFeatureDefinition: Feature definition with temporal metadata
+    - FeatureRegistry: Central registry for all feature definitions
+    - FeatureStoreManager: Unified interface for feature store operations
+
+Example:
+    >>> from customer_retention.integrations.feature_store import (
+    ...     FeatureStoreManager, TemporalFeatureDefinition, FeatureRegistry
+    ... )
+    >>>
+    >>> # Create feature definitions
+    >>> registry = FeatureRegistry()
+    >>> registry.register(TemporalFeatureDefinition(
+    ...     name="tenure_months",
+    ...     description="Customer tenure in months",
+    ...     entity_key="customer_id",
+    ...     timestamp_column="feature_timestamp",
+    ...     source_columns=["tenure"],
+    ... ))
+    >>>
+    >>> # Create feature store manager
+    >>> manager = FeatureStoreManager.create(backend="feast")
+    >>> manager.publish_features(df, registry)
+"""
+
+from .definitions import (
+    FeatureComputationType,
+    TemporalAggregation,
+    TemporalFeatureDefinition,
+)
+from .manager import FeatureStoreManager, get_feature_store_manager
+from .registry import FeatureRegistry
+
+__all__ = [
+    "TemporalFeatureDefinition",
+    "FeatureComputationType",
+    "TemporalAggregation",
+    "FeatureRegistry",
+    "FeatureStoreManager",
+    "get_feature_store_manager",
+]