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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batch-analytics
-Version: 0.3.35
+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.35 → batch_analytics-0.3.36}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "batch-analytics"
-version = "0.3.35"
+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.35 → 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,9 +106,8 @@ class TransformConfig:
             "/tmp/analytics_stage",
         )
     )
-    # Output format for load_staged when reading (parquet/delta/clickhouse/iceberg).
-    # Stage job always writes to ClickHouse; use clickhouse for analytics to read from staged table.
-    # Use iceberg to read Nessie Iceberg gold/silver tables (BATCH_STAGING_TABLE = catalog.ns.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")
     )

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.35 → 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,
@@ -356,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:

{batch_analytics-0.3.35 → batch_analytics-0.3.36}/src/batch_analytics/transform.py

@@ -200,10 +200,23 @@ def apply_spark_expr_columns(
     return out
 
 
+def _pivot_column_prefix(spec: dict) -> str:
+    """Prefix for renamed pivot columns.
+
+    Omitted ``column_name_prefix`` defaults to ``v`` (``v_ph``). An explicit empty
+    string stays empty so the column is the sanitized pivot value (``ph``).
+    """
+    if "column_name_prefix" in spec and spec.get("column_name_prefix") is not None:
+        return str(spec.get("column_name_prefix")).strip()
+    return "v"
+
+
 def _pivot_output_column_name(prefix: str, pivot_value: str) -> str:
     """Stable identifier for pivoted columns (e.g. imp_rm_001 from RM-001)."""
-    p = (prefix or "v").rstrip("_")
+    p = (prefix or "").rstrip("_")
     body = str(pivot_value).replace("-", "_").replace(" ", "_").lower()
+    if not p:
+        return body
     return f"{p}_{body}"
 
 
@@ -246,7 +259,7 @@ def apply_pivot(
     if not group_by or not pivot_col or not value_col:
         raise ValueError("pivot spec requires group_by, pivot_column, value_column")
     agg_name = (spec.get("agg") or "max").strip().lower()
-    prefix = (spec.get("column_name_prefix") or "v").strip() or "v"
+    prefix = _pivot_column_prefix(spec)
     rename_pivot = spec.get("rename_pivot_columns", True)
 
     from pyspark.sql import functions as F
@@ -494,6 +507,31 @@ def stage_to_clickhouse(
     )
 
 
+def stage_to_iceberg(
+    spark: SparkSession,
+    df: DataFrame,
+    config: BatchAnalyticsConfig,
+) -> None:
+    """Write transformed data to a Nessie Iceberg table (``BATCH_STAGING_FORMAT=iceberg``)."""
+    del spark  # catalog is already registered on the session that owns ``df``
+    n = df.count()
+    mode = _normalize_staging_write_mode(config.transform.staging_write_mode)
+    table = quote_iceberg_identifier(_resolve_iceberg_table_name(config))
+    logger.info(
+        "Staging to Iceberg %s (mode=%s); dataframe columns (%d): %s",
+        table,
+        mode,
+        len(df.columns),
+        sorted(df.columns),
+    )
+    writer = df.writeTo(table)
+    if mode == "overwrite":
+        writer.createOrReplace()
+    else:
+        writer.append()
+    logger.info("Staged data to Iceberg %s (%d rows)", table, n)
+
+
 def stage_to_path(
     spark: SparkSession,
     df: DataFrame,
@@ -570,6 +608,58 @@ def _staging_where_clause(config: BatchAnalyticsConfig) -> str:
     return f" WHERE {filt}" if filt else ""
 
 
+_SIMPLE_IDENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+def quote_iceberg_identifier(raw: str) -> str:
+    """Quote Iceberg path segments that are not plain Spark identifiers.
+
+    ``nessie.gold.batch_yield_features`` stays unquoted.
+    A hyphenated tenant namespace is wrapped in backticks.
+    """
+    text = (raw or "").strip()
+    if not text:
+        raise ValueError("Iceberg table name is empty")
+    parts = _split_iceberg_identifier(text)
+    return ".".join(_quote_iceberg_part(part) for part in parts)
+
+
+def _split_iceberg_identifier(text: str) -> List[str]:
+    parts: List[str] = []
+    buf: List[str] = []
+    in_quote = False
+    i = 0
+    while i < len(text):
+        ch = text[i]
+        if ch == "`":
+            if in_quote and i + 1 < len(text) and text[i + 1] == "`":
+                buf.append("`")
+                i += 2
+                continue
+            in_quote = not in_quote
+            i += 1
+            continue
+        if ch == "." and not in_quote:
+            parts.append("".join(buf))
+            buf = []
+            i += 1
+            continue
+        buf.append(ch)
+        i += 1
+    if in_quote:
+        raise ValueError(f"Unbalanced backticks in Iceberg identifier: {text}")
+    parts.append("".join(buf))
+    if any(part == "" for part in parts):
+        raise ValueError(f"Empty segment in Iceberg identifier: {text}")
+    return parts
+
+
+def _quote_iceberg_part(part: str) -> str:
+    if _SIMPLE_IDENT.match(part):
+        return part
+    return "`" + part.replace("`", "``") + "`"
+
+
 def _resolve_iceberg_table_name(config: BatchAnalyticsConfig) -> str:
     """Return fully-qualified Iceberg table id (catalog.namespace.table)."""
     raw = (config.transform.staging_table or "").strip()

{batch_analytics-0.3.35 → batch_analytics-0.3.36}/src/batch_analytics.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batch-analytics
-Version: 0.3.35
+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.35 → batch_analytics-0.3.36}/src/batch_analytics.egg-info/SOURCES.txt

@@ -33,4 +33,5 @@ src/batch_analytics/output/s3.py
 src/batch_analytics/utils/__init__.py
 src/batch_analytics/utils/gluon_autogluon_common.py
 tests/test_iceberg_staging.py
-tests/test_pca_mvda.py
+tests/test_pca_mvda.py
+tests/test_yield_quality_pivot.py

batch_analytics-0.3.36/tests/test_yield_quality_pivot.py

@@ -0,0 +1,218 @@
+"""Wide yield/quality pivot on local Spark, using gold-shaped fixture rows."""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+from unittest.mock import MagicMock
+
+_TESTS = Path(__file__).resolve().parent
+if str(_TESTS) not in sys.path:
+    sys.path.insert(0, str(_TESTS))
+
+from fixtures.yield_quality_wide import (  # noqa: E402
+    AGGREGATE_BEFORE_JOIN,
+    EXPECTED_WIDE_ROWS,
+    EXPR_COLUMNS,
+    EXTRACT_COLUMNS,
+    PIVOT_SPEC,
+    QUALITY_ROWS,
+    QUALITY_TABLE,
+    QUOTED_WIDE_TABLE,
+    WIDE_TABLE,
+    YIELD_ROWS,
+    YIELD_TABLE,
+)
+
+from batch_analytics.config import BatchAnalyticsConfig, ExtractConfig, TransformConfig
+from batch_analytics.extract import unify_dataframes
+from batch_analytics.transform import (
+    _pivot_output_column_name,
+    quote_iceberg_identifier,
+    stage_to_iceberg,
+    transform,
+)
+
+
+def _json_frame(spark, rows):
+    """Load rows through the JVM JSON reader.
+
+    ``createDataFrame`` pickles each row. PySpark 3.5 on Python 3.14 recurses
+    inside cloudpickle, so the fixture is written as JSON lines instead.
+    """
+    handle, path = tempfile.mkstemp(suffix=".jsonl")
+    os.close(handle)
+    with open(path, "w", encoding="utf-8") as fh:
+        for row in rows:
+            fh.write(json.dumps(row) + "\n")
+    return spark.read.json(path)
+
+
+def _frames(spark):
+    return {
+        QUALITY_TABLE: _json_frame(spark, QUALITY_ROWS),
+        YIELD_TABLE: _json_frame(spark, YIELD_ROWS),
+    }
+
+
+def _config() -> BatchAnalyticsConfig:
+    return BatchAnalyticsConfig(
+        extract=ExtractConfig(
+            source_tables=f"{QUALITY_TABLE},{YIELD_TABLE}",
+            source_format="iceberg",
+            join_keys="batch_no",
+            columns_json=json.dumps(EXTRACT_COLUMNS),
+            aggregate_before_join_json=json.dumps(AGGREGATE_BEFORE_JOIN),
+            filter_column="",
+            filter_values="",
+            use_native_connector=False,
+        ),
+        transform=TransformConfig(
+            dedup_columns="batch_no",
+            add_dimension_column="__absent__",
+            expr_columns_json=json.dumps(EXPR_COLUMNS),
+            pivot_json=json.dumps(PIVOT_SPEC),
+            staging_format="iceberg",
+            staging_table=WIDE_TABLE,
+            staging_write_mode="overwrite",
+        ),
+    )
+
+
+class PivotNameTests(unittest.TestCase):
+    def test_source_format_defaults_to_clickhouse(self) -> None:
+        with mock.patch.dict(os.environ, {}, clear=False):
+            os.environ.pop("BATCH_SOURCE_FORMAT", None)
+            cfg = ExtractConfig()
+        self.assertEqual(cfg.source_format, "clickhouse")
+
+    def test_empty_prefix_sanitizes_test_name(self) -> None:
+        self.assertEqual(_pivot_output_column_name("", "NaOCl Assay"), "naocl_assay")
+        self.assertEqual(_pivot_output_column_name("", "pH"), "ph")
+
+    def test_default_prefix_keeps_v(self) -> None:
+        self.assertEqual(_pivot_output_column_name("v", "pH"), "v_ph")
+
+    def test_hyphenated_tenant_namespace_is_quoted(self) -> None:
+        self.assertEqual(quote_iceberg_identifier(WIDE_TABLE), QUOTED_WIDE_TABLE)
+        self.assertEqual(
+            quote_iceberg_identifier("nessie.gold.batch_yield_features"),
+            "nessie.gold.batch_yield_features",
+        )
+
+
+class StageIcebergTests(unittest.TestCase):
+    def test_overwrite_uses_quoted_write_to(self) -> None:
+        df = MagicMock()
+        writer = MagicMock()
+        df.writeTo.return_value = writer
+        df.columns = ["batch_no", "actual_yield_qty", "ph"]
+        df.count.return_value = 3
+        stage_to_iceberg(MagicMock(), df, _config())
+        df.writeTo.assert_called_once_with(QUOTED_WIDE_TABLE)
+        writer.createOrReplace.assert_called_once()
+        writer.append.assert_not_called()
+
+    def test_append_mode(self) -> None:
+        cfg = _config()
+        cfg.transform.staging_write_mode = "append"
+        df = MagicMock()
+        writer = MagicMock()
+        df.writeTo.return_value = writer
+        df.columns = ["batch_no"]
+        df.count.return_value = 1
+        stage_to_iceberg(MagicMock(), df, cfg)
+        writer.append.assert_called_once()
+        writer.createOrReplace.assert_not_called()
+
+
+class YieldQualityWideTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        from pyspark.sql import SparkSession
+
+        cls.spark = (
+            SparkSession.builder.master("local[1]")
+            .appName("yield-quality-pivot")
+            .config("spark.ui.enabled", "false")
+            .config("spark.sql.shuffle.partitions", "2")
+            .getOrCreate()
+        )
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls.spark.stop()
+
+    def setUp(self) -> None:
+        self._env = {
+            "BATCH_TRANSFORM_EXPR_COLUMNS": os.environ.pop("BATCH_TRANSFORM_EXPR_COLUMNS", None),
+            "BATCH_TRANSFORM_PIVOT_JSON": os.environ.pop("BATCH_TRANSFORM_PIVOT_JSON", None),
+        }
+
+    def tearDown(self) -> None:
+        for key, value in self._env.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+
+    def test_pivot_matches_expected_wide_rows(self) -> None:
+        wide = transform(unify_dataframes(_frames(self.spark), _config()), _config())
+        rows = {row["batch_no"]: row.asDict() for row in wide.collect()}
+        self.assertEqual(set(rows), {"SL25264", "SL25265", "SL25266"})
+        self.assertEqual(
+            set(wide.columns),
+            {"batch_no", "actual_yield_qty", "ph", "naocl_assay", "temperature"},
+        )
+        for expected in EXPECTED_WIDE_ROWS:
+            got = rows[expected["batch_no"]]
+            self.assertEqual(got["batch_no"], expected["batch_no"])
+            for col_name in ("actual_yield_qty", "ph", "naocl_assay", "temperature"):
+                self._assert_close(got[col_name], expected[col_name], col_name, expected["batch_no"])
+
+    def test_duplicate_non_key_columns_are_rejected(self) -> None:
+        quality = _json_frame(self.spark, QUALITY_ROWS).select("batch_no", "document_id", "test_name")
+        yield_df = _json_frame(self.spark, YIELD_ROWS).select("batch_no", "document_id", "actual_yield_qty")
+        cfg = BatchAnalyticsConfig(
+            extract=ExtractConfig(
+                source_tables="q,y",
+                join_keys="batch_no",
+                columns_json="",
+                aggregate_before_join_json="",
+                filter_column="",
+                filter_values="",
+                use_native_connector=False,
+            )
+        )
+        with self.assertRaises(ValueError) as ctx:
+            unify_dataframes({"q": quality, "y": yield_df}, cfg)
+        self.assertIn("document_id", str(ctx.exception))
+
+    def test_omitted_prefix_uses_v(self) -> None:
+        spec = dict(PIVOT_SPEC)
+        del spec["column_name_prefix"]
+        cfg = _config()
+        cfg.transform.pivot_json = json.dumps(spec)
+        wide = transform(unify_dataframes(_frames(self.spark), cfg), cfg)
+        self.assertIn("v_ph", wide.columns)
+        self.assertNotIn("ph", wide.columns)
+
+    @staticmethod
+    def _assert_close(got, expected, col_name: str, batch_no: str) -> None:
+        if expected is None:
+            if got is not None:
+                raise AssertionError(f"{batch_no}.{col_name}: expected null, got {got}")
+            return
+        if got is None:
+            raise AssertionError(f"{batch_no}.{col_name}: expected {expected}, got null")
+        if abs(float(got) - float(expected)) > 1e-6:
+            raise AssertionError(f"{batch_no}.{col_name}: expected {expected}, got {got}")
+
+
+if __name__ == "__main__":
+    unittest.main()

batch_analytics-0.3.35/src/batch_analytics/extract.py

@@ -1,214 +0,0 @@
-"""
-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.functions import col
-
-from .config import BatchAnalyticsConfig
-
-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 table from ClickHouse.
-    Uses native connector if configured, otherwise JDBC.
-    """
-    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)
-
-    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 from ClickHouse.
-    Returns a dict mapping table name to DataFrame.
-    """
-    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
-
-
-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.
-    - 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)
-
-    if len(all_dfs) == 1:
-        return list(all_dfs.values())[0]
-
-    if join_keys:
-        dfs = list(all_dfs.values())
-        base = dfs[0]
-        for other in dfs[1:]:
-            base = base.join(other, on=join_keys, how="left")
-        return base
-
-    # Multiple tables, no join: use primary or first
-    if primary_table and primary_table in all_dfs:
-        return all_dfs[primary_table]
-    return list(all_dfs.values())[0]