churnkit 0.75.1a3__py3-none-any.whl → 0.76.0a2__py3-none-any.whl

customer_retention/integrations/adapters/feature_store/databricks.py

@@ -8,6 +8,24 @@ from ..base import AdapterResult
 from .base import FeatureStoreAdapter, FeatureViewConfig
 
 
+def _import_feature_engineering_client() -> Any:
+    try:
+        from databricks.feature_engineering import FeatureEngineeringClient
+
+        return FeatureEngineeringClient
+    except ImportError:
+        from databricks.feature_store import FeatureStoreClient
+
+        return FeatureStoreClient
+
+
+def _validate_write_mode(mode: str) -> None:
+    if mode == "overwrite":
+        raise ValueError(
+            "FeatureEngineeringClient.write_table only supports mode='merge'. Use mode='merge' instead of 'overwrite'."
+        )
+
+
 class DatabricksFeatureStore(FeatureStoreAdapter):
     def __init__(self, catalog: str = "main", schema: str = "default"):
         if not is_spark_available():
@@ -19,27 +37,46 @@ class DatabricksFeatureStore(FeatureStoreAdapter):
     @property
     def fe_client(self) -> Any:
         if self._fe_client is None:
-            from databricks.feature_engineering import FeatureEngineeringClient
-            self._fe_client = FeatureEngineeringClient()
+            client_cls = _import_feature_engineering_client()
+            self._fe_client = client_cls()
         return self._fe_client
 
     def _full_name(self, name: str) -> str:
         return f"{self.catalog}.{self.schema}.{name}"
 
-    def create_table(self, name: str, schema: Dict[str, str], primary_keys: List[str]) -> AdapterResult:
+    def create_table(
+        self, name: str, schema: Dict[str, str], primary_keys: List[str], timeseries_column: Optional[str] = None
+    ) -> AdapterResult:
         full_name = self._full_name(name)
         spark = get_spark_session()
         df = spark.createDataFrame([], self._schema_to_spark(schema))
-        self.fe_client.create_table(name=full_name, primary_keys=primary_keys, df=df)
+        kwargs: Dict[str, Any] = {"name": full_name, "primary_keys": primary_keys, "df": df}
+        if timeseries_column:
+            kwargs["timeseries_columns"] = [timeseries_column]
+        self.fe_client.create_table(**kwargs)
         return AdapterResult(success=True, metadata={"name": full_name})
 
     def _schema_to_spark(self, schema: Dict[str, str]) -> Any:
-        from pyspark.sql.types import FloatType, IntegerType, StringType, StructField, StructType
-        type_map = {"int": IntegerType(), "float": FloatType(), "string": StringType()}
+        from pyspark.sql.types import (
+            FloatType,
+            IntegerType,
+            StringType,
+            StructField,
+            StructType,
+            TimestampType,
+        )
+
+        type_map = {
+            "int": IntegerType(),
+            "float": FloatType(),
+            "string": StringType(),
+            "timestamp": TimestampType(),
+        }
         fields = [StructField(name, type_map.get(dtype, StringType()), True) for name, dtype in schema.items()]
         return StructType(fields)
 
     def write_table(self, name: str, df: pd.DataFrame, mode: str = "merge") -> AdapterResult:
+        _validate_write_mode(mode)
         full_name = self._full_name(name)
         spark = get_spark_session()
         spark_df = spark.createDataFrame(df)
@@ -72,14 +109,22 @@ class DatabricksFeatureStore(FeatureStoreAdapter):
         table_name = self._full_name(config.name)
         spark = get_spark_session()
         spark_df = spark.createDataFrame(df)
-        self.fe_client.create_table(name=table_name, primary_keys=[config.entity_key], df=spark_df)
+        kwargs: Dict[str, Any] = {"name": table_name, "primary_keys": [config.entity_key], "df": spark_df}
+        if hasattr(config, "timeseries_column") and config.timeseries_column:
+            kwargs["timeseries_columns"] = [config.timeseries_column]
+        self.fe_client.create_table(**kwargs)
         return table_name
 
     def get_historical_features(self, entity_df: pd.DataFrame, feature_refs: List[str]) -> pd.DataFrame:
         from databricks.feature_engineering import FeatureLookup
+
         spark = get_spark_session()
-        lookups = [FeatureLookup(table_name=ref.split(":")[0], lookup_key=[entity_df.columns[0]]) for ref in feature_refs]
-        training_set = self.fe_client.create_training_set(df=spark.createDataFrame(entity_df), feature_lookups=lookups, label=None)
+        lookups = [
+            FeatureLookup(table_name=ref.split(":")[0], lookup_key=[entity_df.columns[0]]) for ref in feature_refs
+        ]
+        training_set = self.fe_client.create_training_set(
+            df=spark.createDataFrame(entity_df), feature_lookups=lookups, label=None
+        )
         return training_set.load_df().toPandas()
 
     def materialize(self, feature_views: List[str], start_date: str, end_date: str) -> None:
@@ -89,6 +134,9 @@ class DatabricksFeatureStore(FeatureStoreAdapter):
         entity_df = pd.DataFrame(entity_keys)
         spark = get_spark_session()
         from databricks.feature_engineering import FeatureLookup
-        lookups = [FeatureLookup(table_name=ref.split(":")[0], lookup_key=list(entity_keys.keys())) for ref in feature_refs]
+
+        lookups = [
+            FeatureLookup(table_name=ref.split(":")[0], lookup_key=list(entity_keys.keys())) for ref in feature_refs
+        ]
         result = self.fe_client.score_batch(df=spark.createDataFrame(entity_df), feature_lookups=lookups)
         return result.toPandas().to_dict()

customer_retention/integrations/adapters/mlflow/base.py

@@ -30,3 +30,11 @@ class MLflowAdapter(ABC):
     @abstractmethod
     def transition_stage(self, model_name: str, version: str, stage: str) -> None:
         pass
+
+    @abstractmethod
+    def set_alias(self, model_name: str, alias: str, version: str) -> None:
+        pass
+
+    @abstractmethod
+    def get_model_by_alias(self, model_name: str, alias: str) -> Any:
+        pass

customer_retention/integrations/adapters/mlflow/databricks.py

@@ -7,9 +7,12 @@ from .base import MLflowAdapter
 try:
     import mlflow
     from mlflow.tracking import MlflowClient
+
     MLFLOW_AVAILABLE = True
+    MLFLOW_MAJOR_VERSION = int(mlflow.__version__.split(".")[0])
 except ImportError:
     MLFLOW_AVAILABLE = False
+    MLFLOW_MAJOR_VERSION = 0
 
 
 class DatabricksMLflow(MLflowAdapter):
@@ -18,7 +21,8 @@ class DatabricksMLflow(MLflowAdapter):
             raise ImportError("PySpark required for DatabricksMLflow")
         if not MLFLOW_AVAILABLE:
             raise ImportError("mlflow package required")
-        mlflow.set_registry_uri(registry_uri)
+        if MLFLOW_MAJOR_VERSION < 3:
+            mlflow.set_registry_uri(registry_uri)
         self.registry_uri = registry_uri
         self._client = MlflowClient()
         self._run_id = None
@@ -44,7 +48,10 @@ class DatabricksMLflow(MLflowAdapter):
         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)
+        if MLFLOW_MAJOR_VERSION >= 3:
+            info = mlflow.sklearn.log_model(model, name=artifact_path, registered_model_name=registered_name)
+        else:
+            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:
@@ -52,3 +59,9 @@ class DatabricksMLflow(MLflowAdapter):
 
     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)
+
+    def set_alias(self, model_name: str, alias: str, version: str) -> None:
+        self._client.set_registered_model_alias(name=model_name, alias=alias, version=version)
+
+    def get_model_by_alias(self, model_name: str, alias: str) -> Any:
+        return self._client.get_model_version_by_alias(name=model_name, alias=alias)

customer_retention/integrations/adapters/mlflow/local.py

@@ -5,6 +5,7 @@ from .base import MLflowAdapter
 try:
     import mlflow
     from mlflow.tracking import MlflowClient
+
     MLFLOW_AVAILABLE = True
 except ImportError:
     MLFLOW_AVAILABLE = False
@@ -48,3 +49,9 @@ class LocalMLflow(MLflowAdapter):
 
     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)
+
+    def set_alias(self, model_name: str, alias: str, version: str) -> None:
+        self._client.set_registered_model_alias(name=model_name, alias=alias, version=version)
+
+    def get_model_by_alias(self, model_name: str, alias: str) -> Any:
+        return self._client.get_model_version_by_alias(name=model_name, alias=alias)

customer_retention/integrations/databricks_init.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import os
+import shutil
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class DatabricksInitResult:
+    catalog: str
+    schema: str
+    experiment_name: str
+    workspace_path: str | None
+    model_name: str
+    notebooks_copied: list[str] = field(default_factory=list)
+
+    @property
+    def environment_variables(self) -> dict[str, str]:
+        env_vars = {
+            "CR_CATALOG": self.catalog,
+            "CR_SCHEMA": self.schema,
+            "CR_EXPERIMENT_NAME": self.experiment_name,
+            "CR_EXPERIMENTS_DIR": f"/Workspace/{self.workspace_path}/experiments" if self.workspace_path else "",
+        }
+        if self.workspace_path:
+            env_vars["CR_WORKSPACE_PATH"] = self.workspace_path
+        return env_vars
+
+
+def databricks_init(
+    catalog: str = "main",
+    schema: str = "default",
+    experiment_name: str | None = None,
+    workspace_path: str | None = None,
+    copy_notebooks: bool = True,
+    model_name: str = "customer_retention",
+) -> DatabricksInitResult:
+    _validate_databricks_environment()
+    _set_environment_variables(catalog, schema, workspace_path)
+    resolved_experiment_name = experiment_name or _resolve_experiment_name_from_notebook_path()
+    resolved_experiment_name = _make_absolute_experiment_path(resolved_experiment_name, workspace_path)
+    _set_experiment_name_env_var(resolved_experiment_name)
+    _configure_mlflow_experiment(resolved_experiment_name)
+    notebooks_copied: list[str] = []
+    if copy_notebooks and workspace_path:
+        notebooks_copied = _copy_exploration_notebooks(workspace_path)
+    result = DatabricksInitResult(
+        catalog=catalog,
+        schema=schema,
+        experiment_name=resolved_experiment_name,
+        workspace_path=workspace_path,
+        model_name=model_name,
+        notebooks_copied=notebooks_copied,
+    )
+    _display_init_summary(result)
+    return result
+
+
+def _validate_databricks_environment() -> None:
+    if not os.environ.get("DATABRICKS_RUNTIME_VERSION"):
+        raise RuntimeError(
+            "databricks_init() must be called from a Databricks notebook. "
+            "DATABRICKS_RUNTIME_VERSION not found in environment."
+        )
+
+
+def _set_environment_variables(catalog: str, schema: str, workspace_path: str | None) -> None:
+    os.environ["CR_CATALOG"] = catalog
+    os.environ["CR_SCHEMA"] = schema
+    if workspace_path:
+        os.environ["CR_WORKSPACE_PATH"] = workspace_path
+        os.environ["CR_EXPERIMENTS_DIR"] = f"/Workspace/{workspace_path}/experiments"
+
+
+def _set_experiment_name_env_var(experiment_name: str) -> None:
+    os.environ["CR_EXPERIMENT_NAME"] = experiment_name
+
+
+def _resolve_experiment_name_from_notebook_path() -> str:
+    try:
+        dbutils = _get_dbutils()
+        if dbutils:
+            notebook_path = dbutils.notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get()
+            return notebook_path.rsplit("/", 1)[-1]
+    except Exception:
+        pass
+    return "customer_retention"
+
+
+def _get_dbutils() -> Any | None:
+    try:
+        from customer_retention.core.compat.detection import get_dbutils
+
+        return get_dbutils()
+    except Exception:
+        return None
+
+
+def _make_absolute_experiment_path(experiment_name: str, workspace_path: str | None) -> str:
+    if experiment_name.startswith("/"):
+        return experiment_name
+    if not workspace_path:
+        return experiment_name
+    base = workspace_path.removeprefix("/Workspace")
+    if not base.startswith("/"):
+        base = f"/{base}"
+    return f"{base}/{experiment_name}"
+
+
+def _configure_mlflow_experiment(experiment_name: str) -> None:
+    try:
+        import mlflow
+
+        mlflow.set_experiment(experiment_name)
+    except ImportError:
+        pass
+
+
+def _copy_exploration_notebooks(workspace_path: str) -> list[str]:
+    from customer_retention.generators.notebook_generator.project_init import ProjectInitializer
+
+    source_dir = ProjectInitializer(project_name="")._get_exploration_source_dir()
+    if not source_dir or not source_dir.exists():
+        return []
+
+    dest_dir = Path(f"/Workspace/{workspace_path}/exploration_notebooks")
+    dest_dir.mkdir(parents=True, exist_ok=True)
+
+    copied = []
+    for notebook in source_dir.glob("*.ipynb"):
+        dest_path = dest_dir / notebook.name
+        if not dest_path.exists():
+            shutil.copy2(notebook, dest_path)
+            copied.append(str(dest_path))
+
+    return copied
+
+
+def _display_init_summary(result: DatabricksInitResult) -> None:
+    print("ChurnKit Databricks Initialization Complete")
+    print("=" * 45)
+    print(f"  Catalog:          {result.catalog}")
+    print(f"  Schema:           {result.schema}")
+    print(f"  Experiment:       {result.experiment_name}")
+    print(f"  Workspace Path:   {result.workspace_path or '(not set)'}")
+    print(f"  Model Name:       {result.model_name}")
+    if result.notebooks_copied:
+        print(f"  Notebooks Copied: {len(result.notebooks_copied)}")
+        for nb in result.notebooks_copied:
+            print(f"    - {nb}")
+    print("=" * 45)

customer_retention/stages/profiling/temporal_feature_analyzer.py

@@ -5,7 +5,7 @@ from typing import Any, Dict, List, Optional, Tuple
 import numpy as np
 from scipy import stats
 
-from customer_retention.core.compat import DataFrame, ensure_datetime_column, pd, qcut, to_pandas
+from customer_retention.core.compat import DataFrame, ensure_datetime_column, native_pd, pd, qcut, to_pandas
 from customer_retention.core.utils import compute_effect_size
 
 
@@ -642,7 +642,7 @@ class TemporalFeatureAnalyzer:
             )
 
     def _calculate_iv(self, feature: pd.Series, target: pd.Series, bins: int = 10) -> float:
-        df_iv = pd.DataFrame({"feature": feature, "target": target}).dropna()
+        df_iv = native_pd.DataFrame({"feature": feature, "target": target}).dropna()
         if len(df_iv) < bins * 2:
             return 0.0
         try:
@@ -666,7 +666,7 @@ class TemporalFeatureAnalyzer:
         return float(grouped["iv"].sum())
 
     def _calculate_ks(self, feature: pd.Series, target: pd.Series) -> Tuple[float, float]:
-        df_ks = pd.DataFrame({"feature": feature, "target": target}).dropna()
+        df_ks = native_pd.DataFrame({"feature": feature, "target": target}).dropna()
         group0, group1 = df_ks[df_ks["target"] == 0]["feature"], df_ks[df_ks["target"] == 1]["feature"]
         if len(group0) == 0 or len(group1) == 0:
             return 0.0, 1.0

customer_retention/stages/profiling/temporal_feature_engineer.py

@@ -25,7 +25,7 @@ from typing import Any, Dict, List, Optional
 
 import numpy as np
 
-from customer_retention.core.compat import Timedelta, pd, to_datetime, to_pandas
+from customer_retention.core.compat import Timedelta, native_pd, pd, to_datetime, to_pandas
 
 
 class ReferenceMode(Enum):
@@ -307,7 +307,7 @@ class TemporalFeatureEngineer:
 
         if self.config.reference_mode == ReferenceMode.GLOBAL_DATE:
             ref_date = self.config.global_reference_date or datetime.now()
-            return pd.DataFrame({
+            return native_pd.DataFrame({
                 entity_col: entities,
                 "reference_date": ref_date,
             })

customer_retention/stages/profiling/temporal_pattern_analyzer.py

@@ -10,7 +10,7 @@ from customer_retention.core.compat import (
     Timestamp,
     cut,
     ensure_datetime_column,
-    pd,
+    native_pd,
     safe_to_datetime,
     to_pandas,
 )
@@ -316,7 +316,7 @@ def _diagnose_anomaly_pattern(
     entity_first = df.groupby(entity_column)[time_column].min()
     entity_last = df.groupby(entity_column)[time_column].max()
     tenure = (entity_last - entity_first).dt.days
-    tenure_by_target = pd.DataFrame({"target": entity_target, "tenure": tenure})
+    tenure_by_target = native_pd.DataFrame({"target": entity_target, "tenure": tenure})
     retained_tenure = tenure_by_target[tenure_by_target["target"] == 1]["tenure"]
     churned_tenure = tenure_by_target[tenure_by_target["target"] == 0]["tenure"]
     retained_median_tenure = float(retained_tenure.median()) if len(retained_tenure) > 0 else None
@@ -597,7 +597,7 @@ class TemporalPatternAnalyzer:
 
     def analyze_cohorts(self, df: DataFrame, entity_column: str, cohort_column: str, target_column: Optional[str] = None, period: str = "M") -> DataFrame:
         if len(df) == 0:
-            return pd.DataFrame()
+            return native_pd.DataFrame()
 
         df_copy = to_pandas(df).copy()
         ensure_datetime_column(df_copy, cohort_column)
@@ -638,7 +638,7 @@ class TemporalPatternAnalyzer:
         target_correlation = None
         if target_column and target_column in df.columns:
             entity_target = df.groupby(entity_column)[target_column].first()
-            combined = pd.DataFrame({"recency": recency_days, "target": entity_target}).dropna()
+            combined = native_pd.DataFrame({"recency": recency_days, "target": entity_target}).dropna()
 
             if len(combined) > 2:
                 corr, _ = stats.pearsonr(combined["recency"], combined["target"])

customer_retention/stages/profiling/time_series_profiler.py

@@ -7,7 +7,7 @@ from customer_retention.core.compat import (
     DataFrame,
     Timestamp,
     ensure_datetime_column,
-    pd,
+    native_pd,
     to_pandas,
 )
 
@@ -95,7 +95,7 @@ def classify_lifecycle_quadrants(entity_lifecycles: DataFrame) -> LifecycleQuadr
         lifecycles=lc,
         tenure_threshold=tenure_threshold,
         intensity_threshold=intensity_threshold,
-        recommendations=pd.DataFrame(rows),
+        recommendations=native_pd.DataFrame(rows),
     )
 
 
@@ -160,7 +160,7 @@ def classify_activity_segments(entity_lifecycles: DataFrame) -> ActivitySegmentR
         lifecycles=lc,
         q25_threshold=q25,
         q75_threshold=q75,
-        recommendations=pd.DataFrame(rows),
+        recommendations=native_pd.DataFrame(rows),
     )
 
 
@@ -237,7 +237,7 @@ class TimeSeriesProfiler:
     def _compute_entity_lifecycles(self, df: DataFrame) -> DataFrame:
         grouped = df.groupby(self.entity_column)[self.time_column]
 
-        lifecycles = pd.DataFrame({
+        lifecycles = native_pd.DataFrame({
             "entity": grouped.first().index.tolist(),
             "first_event": grouped.min().values,
             "last_event": grouped.max().values,
@@ -302,7 +302,7 @@ class TimeSeriesProfiler:
             events_per_entity=DistributionStats(
                 min=0, max=0, mean=0, median=0, std=0, q25=0, q75=0
             ),
-            entity_lifecycles=pd.DataFrame(columns=[
+            entity_lifecycles=native_pd.DataFrame(columns=[
                 "entity", "first_event", "last_event", "duration_days", "event_count"
             ]),
             avg_inter_event_days=None,

customer_retention/stages/profiling/time_window_aggregator.py

@@ -12,7 +12,7 @@ from customer_retention.core.compat import (
     Timestamp,
     ensure_datetime_column,
     is_numeric_dtype,
-    pd,
+    native_pd,
     to_pandas,
 )
 
@@ -85,7 +85,7 @@ class TimeWindowAggregator:
     ) -> DataFrame:
         df = to_pandas(df)
         if len(df) == 0:
-            return pd.DataFrame()
+            return native_pd.DataFrame()
 
         df = df.copy()
         ensure_datetime_column(df, self.time_column)
@@ -113,7 +113,7 @@ class TimeWindowAggregator:
         if include_tenure:
             result_data["days_since_first_event"] = self._compute_tenure(df, entities, reference_date)
 
-        result = pd.DataFrame(result_data)
+        result = native_pd.DataFrame(result_data)
         result.attrs["aggregation_reference_date"] = (
             reference_date.isoformat() if hasattr(reference_date, "isoformat") else str(reference_date))
         result.attrs["aggregation_timestamp"] = Timestamp.now().isoformat()