batch-analytics 0.3.34__tar.gz → 0.3.36__tar.gz

{batch_analytics-0.3.34 → batch_analytics-0.3.36}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batch-analytics
-Version: 0.3.34
+Version: 0.3.36
 Summary: PySpark batch analytics: Extract, Transform, Stage, and analytical modules (linear regression, correlation, PCA, t-test, LLM classification).
 Author: Litewave Analytics Team
 License: MIT

{batch_analytics-0.3.34 → batch_analytics-0.3.36}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "batch-analytics"
-version = "0.3.34"
+version = "0.3.36"
 description = "PySpark batch analytics: Extract, Transform, Stage, and analytical modules (linear regression, correlation, PCA, t-test, LLM classification)."
 readme = "README.md"
 requires-python = ">=3.8"

{batch_analytics-0.3.34 → batch_analytics-0.3.36}/src/batch_analytics/config.py

@@ -66,6 +66,27 @@ class ExtractConfig:
     # filter_values: comma-separated list, or JSON array e.g. ["a","b"] for values containing commas.
     filter_column: str = os.environ.get("BATCH_EXTRACT_FILTER_COLUMN", "").strip()
     filter_values: str = os.environ.get("BATCH_EXTRACT_FILTER_VALUES", "").strip()
+    # clickhouse (default) or iceberg. Iceberg reads BATCH_SOURCE_TABLES via spark.table.
+    source_format: str = field(
+        default_factory=lambda: (
+            os.environ.get("BATCH_SOURCE_FORMAT") or "clickhouse"
+        ).strip().lower()
+        or "clickhouse"
+    )
+    # Comma-separated left-join keys when more than one source table is set.
+    join_keys: str = field(
+        default_factory=lambda: os.environ.get("BATCH_EXTRACT_JOIN_KEYS", "").strip()
+    )
+    # JSON object: table name or short name -> comma-separated columns (or a list).
+    columns_json: str = field(
+        default_factory=lambda: os.environ.get("BATCH_EXTRACT_COLUMNS_JSON", "").strip()
+    )
+    # JSON object: table -> {"group_by": "batch_no", "aggs": {"actual_yield_qty": "max"}}.
+    aggregate_before_join_json: str = field(
+        default_factory=lambda: os.environ.get(
+            "BATCH_EXTRACT_AGG_BEFORE_JOIN_JSON", ""
+        ).strip()
+    )
 
 
 @dataclass
@@ -85,15 +106,23 @@ class TransformConfig:
             "/tmp/analytics_stage",
         )
     )
-    # Output format for load_staged when reading (parquet/delta/clickhouse).
-    # Stage job always writes to ClickHouse; use clickhouse for analytics to read from staged table.
+    # Staging format for write (ETL) and load_staged (analytics): parquet, delta, clickhouse, iceberg.
+    # iceberg writes BATCH_STAGING_TABLE with DataFrame.writeTo and reads it back via spark.sql.
     staging_format: str = field(
         default_factory=lambda: os.environ.get("BATCH_STAGING_FORMAT", "clickhouse")
     )
-    # Staging table name in ClickHouse (when format=clickhouse)
+    # Staging table: ClickHouse table name, or Iceberg identifier e.g. nessie.gold.batch_yield_features
     staging_table: str = field(
         default_factory=lambda: os.environ.get("BATCH_STAGING_TABLE", "analytics_staging")
     )
+    # Optional SQL WHERE fragment when reading staging (no leading WHERE).
+    staging_filter: str = field(
+        default_factory=lambda: os.environ.get("BATCH_STAGING_FILTER", "").strip()
+    )
+    # Optional comma-separated columns when reading staging (empty = SELECT *).
+    staging_columns: str = field(
+        default_factory=lambda: os.environ.get("BATCH_STAGING_COLUMNS", "").strip()
+    )
     # Spark save mode for ClickHouse staging (and path staging): overwrite | append
     staging_write_mode: str = field(
         default_factory=lambda: os.environ.get("BATCH_STAGING_WRITE_MODE", "overwrite")
@@ -211,11 +240,18 @@ class SparkK8sConfig:
     executor_cores: int = int(os.environ.get("SPARK_EXECUTOR_CORES", "1"))
     executor_memory: str = os.environ.get("SPARK_EXECUTOR_MEMORY", "512m")
     executor_memory_overhead: str = os.environ.get("SPARK_EXECUTOR_MEMORY_OVERHEAD", "128m")
-    # S3 (optional; set for s3a:// paths)
+    # S3 (optional; set for s3a:// paths / Iceberg warehouse)
     s3_access_key: str = os.environ.get("AWS_ACCESS_KEY_ID", "")
     s3_secret_key: str = os.environ.get("AWS_SECRET_ACCESS_KEY", "")
-    s3_endpoint: str = os.environ.get("AWS_ENDPOINT", "s3.amazonaws.com")
-    s3_region: str = os.environ.get("AWS_REGION", "us-east-2")
+    # Prefer S3_ENDPOINT (lakehouse MinIO); fall back to AWS_ENDPOINT for legacy.
+    s3_endpoint: str = os.environ.get(
+        "S3_ENDPOINT",
+        os.environ.get("AWS_ENDPOINT", ""),
+    )
+    s3_region: str = os.environ.get(
+        "AWS_DEFAULT_REGION",
+        os.environ.get("AWS_REGION", "us-east-2"),
+    )
 
 
 @dataclass

batch_analytics-0.3.36/src/batch_analytics/extract.py

@@ -0,0 +1,375 @@
+"""
+Extract stage: Load data from ClickHouse using Spark ClickHouse connector or JDBC.
+"""
+
+import json
+import logging
+import os
+from typing import Dict, List, Optional, Tuple
+
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql import functions as F
+from pyspark.sql.functions import col
+
+from .config import BatchAnalyticsConfig
+from .transform import quote_iceberg_identifier
+
+logger = logging.getLogger(__name__)
+
+
+def parse_extract_filter_values(raw: str) -> List[str]:
+    """
+    Parse BATCH_EXTRACT_FILTER_VALUES: comma-separated tokens, or JSON array string.
+
+    Examples:
+      a,b,c -> ["a","b","c"]
+      ["GP/A","GP/B"] -> JSON list (values may contain commas)
+    """
+    text = (raw or "").strip()
+    if not text:
+        return []
+    if text.startswith("["):
+        try:
+            data = json.loads(text)
+            if isinstance(data, list):
+                out = [str(x).strip() for x in data if str(x).strip()]
+                return out
+        except json.JSONDecodeError:
+            logger.warning("BATCH_EXTRACT_FILTER_VALUES looks like JSON but failed to parse; using comma split")
+    return [p.strip() for p in text.split(",") if p.strip()]
+
+
+def _apply_extract_filter(df: DataFrame, config: BatchAnalyticsConfig) -> DataFrame:
+    """Apply col IN (values) when filter_column is set; empty column = no filter."""
+    col_name = (config.extract.filter_column or "").strip()
+    if not col_name:
+        return df
+    if col_name not in df.columns:
+        logger.warning(
+            "BATCH_EXTRACT_FILTER_COLUMN=%r not in extracted columns %s; skipping filter",
+            col_name,
+            df.columns,
+        )
+        return df
+    values = parse_extract_filter_values(config.extract.filter_values)
+    if not values:
+        logger.warning(
+            "BATCH_EXTRACT_FILTER_COLUMN=%r set but BATCH_EXTRACT_FILTER_VALUES is empty; skipping IN filter",
+            col_name,
+        )
+        return df
+    filtered = df.filter(col(col_name).isin(values))
+    return filtered
+
+
+def _parse_ch_database_table(table: str, default_database: str) -> Tuple[str, str]:
+    """
+    Resolve ``table`` reference to (database, table_name).
+
+    ``batch_metric_facts`` → (default_database, batch_metric_facts)
+    ``analytics.batch_metric_facts`` → (analytics, batch_metric_facts)
+    """
+    t = (table or "").strip()
+    if not t:
+        return default_database, t
+    if "." in t and not t.startswith("("):
+        db, tbl = t.split(".", 1)
+        db, tbl = db.strip(), tbl.strip()
+        if db and tbl:
+            return db, tbl
+    return default_database, t
+
+
+def _read_via_catalog(spark: SparkSession, cfg: BatchAnalyticsConfig, table: str) -> Optional[DataFrame]:
+    """
+    Read via Spark SQL catalog (ClickHouseCatalog), registered in job_runner.create_spark_session.
+
+    clickhouse-spark-runtime does **not** register legacy short name ``format(\"clickhouse\")`` /
+    ``clickhouse.DefaultSource``; catalog + ``spark.table(catalog.db.table)`` is the supported path.
+    """
+    cat = os.environ.get("BATCH_CLICKHOUSE_CATALOG", "batch_ch").strip()
+    if not cat:
+        return None
+    db, tbl = _parse_ch_database_table(table, cfg.clickhouse.database)
+    ident = f"{cat}.{db}.{tbl}"
+    try:
+        return spark.table(ident)
+    except Exception as e:
+        logger.warning(
+            "Catalog read failed for %s (%s): %s. Trying other readers.",
+            ident,
+            table,
+            e,
+        )
+        return None
+
+
+def _read_via_format(spark: SparkSession, cfg: BatchAnalyticsConfig, table: str) -> Optional[DataFrame]:
+    """
+    Read from ClickHouse using the native format API (clickhouse-spark-runtime).
+    Requires: com.clickhouse.spark:clickhouse-spark-runtime in spark.jars.packages
+    """
+    try:
+        rd = (
+            spark.read.format("clickhouse")
+            .option("host", cfg.clickhouse.host)
+            .option("protocol", cfg.clickhouse.protocol)
+            .option("http_port", str(cfg.clickhouse.port))
+            .option("database", cfg.clickhouse.database)
+            .option("table", table)
+            .option("user", cfg.clickhouse.user)
+        )
+        if cfg.clickhouse.password:
+            rd = rd.option("password", cfg.clickhouse.password)
+        df = rd.load()
+        return df
+    except Exception as e:
+        logger.warning(
+            "Native ClickHouse connector failed for table %s: %s. Falling back to JDBC.",
+            table,
+            e,
+        )
+        return None
+
+
+def _read_via_jdbc(spark: SparkSession, cfg: BatchAnalyticsConfig, table: str) -> DataFrame:
+    """Read from ClickHouse via JDBC."""
+    return spark.read.jdbc(
+        cfg.clickhouse.jdbc_url,
+        table,
+        properties=cfg.clickhouse.jdbc_properties,
+    )
+
+
+def extract_table(
+    spark: SparkSession,
+    table: str,
+    config: BatchAnalyticsConfig,
+) -> DataFrame:
+    """
+    Extract a single source table.
+
+    ``BATCH_SOURCE_FORMAT=iceberg`` reads a Nessie Iceberg identifier via ``spark.table``.
+    The default ``clickhouse`` path uses the native connector when configured, otherwise JDBC.
+    """
+    fmt = (config.extract.source_format or "clickhouse").strip().lower()
+    if fmt == "iceberg":
+        ident = quote_iceberg_identifier(table)
+        logger.info("Extracting Iceberg table %s", ident)
+        df = spark.table(ident)
+    elif fmt == "clickhouse":
+        if config.extract.use_native_connector:
+            # Prefer catalog (matches clickhouse-spark-runtime); avoid legacy DefaultSource path.
+            df = _read_via_catalog(spark, config, table)
+            if df is None:
+                df = _read_via_format(spark, config, table)
+            if df is None:
+                df = _read_via_jdbc(spark, config, table)
+        else:
+            df = _read_via_jdbc(spark, config, table)
+    else:
+        raise ValueError(
+            f"BATCH_SOURCE_FORMAT must be clickhouse or iceberg; got {fmt!r}"
+        )
+
+    df = _apply_extract_filter(df, config)
+    logger.info("Extracted table %s: %d rows", table, df.count())
+    return df
+
+
+def extract_all(
+    spark: SparkSession,
+    config: BatchAnalyticsConfig,
+) -> Dict[str, DataFrame]:
+    """
+    Extract all configured source tables (ClickHouse or Iceberg).
+    Returns a dict mapping table name to DataFrame, in source_tables order.
+    """
+    tables = [t.strip() for t in config.extract.source_tables.split(",") if t.strip()]
+    if not tables:
+        raise ValueError("No source tables configured in BATCH_SOURCE_TABLES")
+
+    result: Dict[str, DataFrame] = {}
+    for table in tables:
+        df = extract_table(spark, table, config)
+        result[table] = df
+
+    return result
+
+
+_AGG_FUNCS = {
+    "max": F.max,
+    "min": F.min,
+    "sum": F.sum,
+    "avg": F.avg,
+    "mean": F.avg,
+    "first": F.first,
+}
+
+
+def _parse_json_object(raw: str, env_name: str) -> Dict:
+    text = (raw or "").strip()
+    if not text:
+        return {}
+    try:
+        data = json.loads(text)
+    except json.JSONDecodeError as e:
+        raise ValueError(f"{env_name} must be valid JSON: {e}") from e
+    if not isinstance(data, dict):
+        raise ValueError(f"{env_name} must be a JSON object")
+    return data
+
+
+def _table_short_name(table_name: str) -> str:
+    text = (table_name or "").strip().rstrip("`")
+    if not text:
+        return ""
+    return text.split(".")[-1].strip("`")
+
+
+def _spec_for_table(table_name: str, spec: Dict) -> Optional[object]:
+    if not spec:
+        return None
+    if table_name in spec:
+        return spec[table_name]
+    short = _table_short_name(table_name)
+    if short and short in spec:
+        return spec[short]
+    return None
+
+
+def _column_list(raw: object, table_name: str) -> List[str]:
+    if isinstance(raw, str):
+        names = [part.strip() for part in raw.split(",") if part.strip()]
+    elif isinstance(raw, list):
+        names = [str(part).strip() for part in raw if str(part).strip()]
+    else:
+        raise ValueError(
+            f"extract.columns for {table_name} must be a comma-separated string or a list"
+        )
+    if not names:
+        raise ValueError(f"extract.columns for {table_name} is empty")
+    return names
+
+
+def prepare_source_frame(
+    df: DataFrame,
+    table_name: str,
+    config: BatchAnalyticsConfig,
+) -> DataFrame:
+    """Project columns and optionally aggregate one source before a join."""
+    columns_spec = _parse_json_object(
+        config.extract.columns_json, "BATCH_EXTRACT_COLUMNS_JSON"
+    )
+    selected = _spec_for_table(table_name, columns_spec)
+    if selected is not None:
+        names = _column_list(selected, table_name)
+        missing = [name for name in names if name not in df.columns]
+        if missing:
+            raise ValueError(
+                f"extract columns {missing} not in {table_name}; have {df.columns}"
+            )
+        df = df.select(*names)
+
+    agg_spec_all = _parse_json_object(
+        config.extract.aggregate_before_join_json,
+        "BATCH_EXTRACT_AGG_BEFORE_JOIN_JSON",
+    )
+    agg_spec = _spec_for_table(table_name, agg_spec_all)
+    if not agg_spec:
+        return df
+    if not isinstance(agg_spec, dict):
+        raise ValueError(f"aggregate_before_join for {table_name} must be an object")
+    group_by = agg_spec.get("group_by") or agg_spec.get("groupBy")
+    gcols = [part.strip() for part in str(group_by or "").split(",") if part.strip()]
+    aggs = agg_spec.get("aggs") or {}
+    if not gcols or not isinstance(aggs, dict) or not aggs:
+        raise ValueError(
+            f"aggregate_before_join for {table_name} requires group_by and aggs"
+        )
+    for name in gcols:
+        if name not in df.columns:
+            raise ValueError(
+                f"aggregate group_by column {name!r} not in {table_name}; have {df.columns}"
+            )
+    exprs = []
+    for col_name, agg_name in aggs.items():
+        key = str(col_name).strip()
+        fn_name = str(agg_name).strip().lower()
+        if key not in df.columns:
+            raise ValueError(
+                f"aggregate column {key!r} not in {table_name}; have {df.columns}"
+            )
+        fn = _AGG_FUNCS.get(fn_name)
+        if fn is None:
+            raise ValueError(
+                f"aggregate {fn_name!r} must be one of {sorted(_AGG_FUNCS)}; table {table_name}"
+            )
+        exprs.append(fn(F.col(key)).alias(key))
+    return df.groupBy(*gcols).agg(*exprs)
+
+
+def _join_keys(config: BatchAnalyticsConfig, join_keys: Optional[List[str]]) -> Optional[List[str]]:
+    if join_keys:
+        return [key.strip() for key in join_keys if key and str(key).strip()]
+    raw = (config.extract.join_keys or "").strip()
+    if not raw:
+        return None
+    return [part.strip() for part in raw.split(",") if part.strip()]
+
+
+def unify_dataframes(
+    all_dfs: Dict[str, DataFrame],
+    config: BatchAnalyticsConfig,
+    join_keys: Optional[List[str]] = None,
+    primary_table: Optional[str] = None,
+) -> DataFrame:
+    """
+    Project, aggregate, and left-join already loaded source frames.
+
+    First table is the join base. Non-key columns that exist on both sides are
+    rejected so a wide pivot is not built on ambiguous names.
+    """
+    if not all_dfs:
+        raise ValueError("No source tables to unify")
+    prepared: Dict[str, DataFrame] = {
+        name: prepare_source_frame(df, name, config) for name, df in all_dfs.items()
+    }
+    keys = _join_keys(config, join_keys)
+    frames = list(prepared.values())
+    if len(frames) == 1:
+        return frames[0]
+    if keys:
+        base = frames[0]
+        for other in frames[1:]:
+            missing = [key for key in keys if key not in base.columns or key not in other.columns]
+            if missing:
+                raise ValueError(f"join keys {missing} missing after extract projection")
+            overlap = (set(base.columns) & set(other.columns)) - set(keys)
+            if overlap:
+                raise ValueError(
+                    "join would duplicate columns "
+                    f"{sorted(overlap)}; set extract.columns to drop them"
+                )
+            base = base.join(other, on=keys, how="left")
+        return base
+    if primary_table and primary_table in prepared:
+        return prepared[primary_table]
+    return frames[0]
+
+
+def extract_unified(
+    spark: SparkSession,
+    config: BatchAnalyticsConfig,
+    join_keys: Optional[List[str]] = None,
+    primary_table: Optional[str] = None,
+) -> DataFrame:
+    """
+    Extract and unify source tables into one DataFrame.
+    - Single table: returns it directly (after optional projection and aggregate).
+    - Multiple tables + join_keys: joins on those keys (left join, first table base).
+    - Multiple tables, no join_keys: returns the primary (or first) table.
+      Use primary_table to choose which table to use for analytics.
+    """
+    all_dfs = extract_all(spark, config)
+    return unify_dataframes(all_dfs, config, join_keys=join_keys, primary_table=primary_table)

{batch_analytics-0.3.34 → batch_analytics-0.3.36}/src/batch_analytics/job_runner.py

@@ -19,7 +19,7 @@ from .extract import extract_unified
 from .log import log_dataframe_summary, log_run
 from .modules import DEFAULT_MODULES, MODULE_REGISTRY, VALID_MODULES
 from .output import write_analytics_output
-from .transform import load_staged, stage_to_clickhouse, transform
+from .transform import load_staged, stage_to_clickhouse, stage_to_iceberg, transform
 
 logging.basicConfig(
     level=logging.INFO,
@@ -161,6 +161,9 @@ def create_spark_session(
                 read_codec,
             )
 
+    builder = _apply_iceberg_nessie_catalog(builder, config)
+    builder = _apply_s3a_lakehouse(builder, config)
+
     if cfg.master.startswith("k8s://"):
         driver_host = socket.gethostbyname(socket.gethostname())
         builder = (
@@ -188,19 +191,91 @@ def create_spark_session(
             .config("spark.kubernetes.executor.serviceAccountName", cfg.service_account)
             .config("spark.kubernetes.container.image.pullPolicy", "IfNotPresent")
         )
-        if cfg.s3_access_key and cfg.s3_secret_key:
-            builder = (
-                builder.config("spark.hadoop.fs.s3a.impl", "org.apache.hadoop.fs.s3a.S3AFileSystem")
-                .config("spark.hadoop.fs.s3a.access.key", cfg.s3_access_key)
-                .config("spark.hadoop.fs.s3a.secret.key", cfg.s3_secret_key)
-                .config("spark.hadoop.fs.s3a.endpoint", cfg.s3_endpoint)
-                .config("spark.hadoop.fs.s3a.endpoint.region", cfg.s3_region)
-            )
         logger.info("Spark on Kubernetes: master=%s", cfg.master)
 
     return builder.getOrCreate()
 
 
+def _apply_iceberg_nessie_catalog(builder, config: BatchAnalyticsConfig):
+    """Register Nessie Iceberg catalog when lakehouse env is present."""
+    nessie_uri = (os.environ.get("NESSIE_URI") or "").strip()
+    warehouse = (os.environ.get("ICEBERG_WAREHOUSE") or "").strip()
+    if not nessie_uri and not warehouse:
+        return builder
+    catalog = (os.environ.get("ICEBERG_CATALOG") or "nessie").strip() or "nessie"
+    ref = (os.environ.get("NESSIE_REF") or "main").strip() or "main"
+    if not nessie_uri:
+        nessie_uri = "http://nessie:19120/api/v2"
+    if not warehouse:
+        warehouse = "s3a://lake/"
+    logger.info(
+        "Configuring Iceberg catalog %s uri=%s ref=%s warehouse=%s",
+        catalog,
+        nessie_uri,
+        ref,
+        warehouse,
+    )
+    return (
+        builder.config(
+            "spark.sql.extensions",
+            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions",
+        )
+        .config(f"spark.sql.catalog.{catalog}", "org.apache.iceberg.spark.SparkCatalog")
+        .config(
+            f"spark.sql.catalog.{catalog}.catalog-impl",
+            "org.apache.iceberg.nessie.NessieCatalog",
+        )
+        .config(f"spark.sql.catalog.{catalog}.uri", nessie_uri)
+        .config(f"spark.sql.catalog.{catalog}.ref", ref)
+        .config(f"spark.sql.catalog.{catalog}.warehouse", warehouse)
+        .config(
+            f"spark.sql.catalog.{catalog}.io-impl",
+            "org.apache.iceberg.hadoop.HadoopFileIO",
+        )
+    )
+
+
+def _apply_s3a_lakehouse(builder, config: BatchAnalyticsConfig):
+    """S3A settings for Iceberg warehouse / s3a staging (AWS or MinIO via S3_ENDPOINT)."""
+    cfg = config.spark_k8s
+    if not (cfg.s3_access_key or cfg.s3_secret_key or os.environ.get("ICEBERG_WAREHOUSE")):
+        return builder
+
+    builder = builder.config(
+        "spark.hadoop.fs.s3a.impl",
+        "org.apache.hadoop.fs.s3a.S3AFileSystem",
+    )
+    if cfg.s3_access_key:
+        builder = builder.config("spark.hadoop.fs.s3a.access.key", cfg.s3_access_key)
+    if cfg.s3_secret_key:
+        builder = builder.config("spark.hadoop.fs.s3a.secret.key", cfg.s3_secret_key)
+
+    endpoint = (cfg.s3_endpoint or "").strip()
+    region = (cfg.s3_region or "us-east-2").strip()
+    if endpoint:
+        # Custom/MinIO endpoint
+        ssl = not endpoint.startswith("http://")
+        builder = (
+            builder.config("spark.hadoop.fs.s3a.endpoint", endpoint)
+            .config("spark.hadoop.fs.s3a.path.style.access", "true")
+            .config(
+                "spark.hadoop.fs.s3a.connection.ssl.enabled",
+                "true" if ssl else "false",
+            )
+        )
+    else:
+        builder = (
+            builder.config(
+                "spark.hadoop.fs.s3a.endpoint",
+                f"https://s3.{region}.amazonaws.com",
+            )
+            .config("spark.hadoop.fs.s3a.endpoint.region", region)
+            .config("spark.hadoop.fs.s3a.path.style.access", "false")
+            .config("spark.hadoop.fs.s3a.connection.ssl.enabled", "true")
+        )
+    return builder
+
+
 def run_pipeline(
     config: Optional[BatchAnalyticsConfig] = None,
     spark: Optional[SparkSession] = None,
@@ -281,12 +356,19 @@ def run_pipeline(
             df_transformed = remove_duplicates(df_transformed, key_columns=dedup_cols)
         result["stages"]["transform"] = {"skipped": True}
 
-    # ----- Stage (to ClickHouse) - separate job before analytics -----
+    # ----- Stage (ClickHouse or Iceberg) - separate job before analytics -----
     if run_stage:
-        logger.info("Stage: Stage to ClickHouse")
-        stage_to_clickhouse(spark, df_transformed, config)
+        staging_format = (config.transform.staging_format or "").strip().lower()
+        if staging_format == "iceberg":
+            logger.info("Stage: Stage to Iceberg")
+            stage_to_iceberg(spark, df_transformed, config)
+            destination = config.transform.staging_table
+        else:
+            logger.info("Stage: Stage to ClickHouse")
+            stage_to_clickhouse(spark, df_transformed, config)
+            destination = f"{config.clickhouse.database}.{config.transform.staging_table}"
         result["stages"]["stage"] = {
-            "destination": f"{config.clickhouse.database}.{config.transform.staging_table}",
+            "destination": destination,
             "row_count": df_transformed.count(),
         }
     else: