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

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

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "batch-analytics"
-version = "0.3.34"
+version = "0.3.35"
 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.35}/src/batch_analytics/config.py

@@ -85,15 +85,24 @@ class TransformConfig:
             "/tmp/analytics_stage",
         )
     )
-    # Output format for load_staged when reading (parquet/delta/clickhouse).
+    # 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: 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 +220,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.34 → batch_analytics-0.3.35}/src/batch_analytics/job_runner.py

@@ -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,

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

@@ -534,33 +534,108 @@ def load_staged(
 ) -> DataFrame:
     """
     Load previously staged data (e.g. when running only analytics modules).
+
+    Formats:
+      parquet / delta — path under BATCH_STAGING_PATH
+      clickhouse — BATCH_STAGING_TABLE in ClickHouse
+      iceberg — Nessie Iceberg table (BATCH_STAGING_TABLE = catalog.ns.table)
     """
     staging_path = config.transform.staging_path
-    fmt = config.transform.staging_format
+    fmt = (config.transform.staging_format or "").strip().lower()
 
     if fmt == "parquet":
         return spark.read.parquet(staging_path)
     if fmt == "delta":
         return spark.read.format("delta").load(staging_path)
+    if fmt == "iceberg":
+        return _load_staged_iceberg(spark, config)
     if fmt == "clickhouse":
-        ch = config.clickhouse
-        tbl = config.transform.staging_table
-        cat = os.environ.get("BATCH_CLICKHOUSE_CATALOG", "batch_ch").strip()
-        if cat:
-            try:
-                return spark.table(f"{cat}.{ch.database}.{tbl}")
-            except Exception as e:
-                logger.warning(
-                    "load_staged: catalog table %s.%s.%s failed (%s), using JDBC",
-                    cat,
-                    ch.database,
-                    tbl,
-                    e,
-                )
-        dbtable = f"(SELECT * FROM `{ch.database}`.`{tbl}`) AS _stg"
-        return spark.read.jdbc(
-            ch.jdbc_url,
-            dbtable,
-            properties=ch.jdbc_properties,
-        )
+        return _load_staged_clickhouse(spark, config)
     return spark.read.format(fmt).load(staging_path)
+
+
+def _staging_select_list(config: BatchAnalyticsConfig) -> str:
+    cols = (config.transform.staging_columns or "").strip()
+    if not cols:
+        return "*"
+    return ", ".join(c.strip() for c in cols.split(",") if c.strip())
+
+
+def _staging_where_clause(config: BatchAnalyticsConfig) -> str:
+    filt = (config.transform.staging_filter or "").strip()
+    if not filt:
+        return ""
+    if filt.lower().startswith("where "):
+        filt = filt[6:].strip()
+    return f" WHERE {filt}" if filt else ""
+
+
+def _resolve_iceberg_table_name(config: BatchAnalyticsConfig) -> str:
+    """Return fully-qualified Iceberg table id (catalog.namespace.table)."""
+    raw = (config.transform.staging_table or "").strip()
+    if not raw:
+        raise ValueError(
+            "BATCH_STAGING_TABLE is required when BATCH_STAGING_FORMAT=iceberg "
+            "(e.g. nessie.gold.batch_yield_features)"
+        )
+    catalog = (os.environ.get("ICEBERG_CATALOG") or "nessie").strip()
+    parts = [p for p in raw.replace("/", ".").split(".") if p]
+    if len(parts) == 1:
+        return f"{catalog}.gold.{parts[0]}"
+    if len(parts) == 2:
+        return f"{catalog}.{parts[0]}.{parts[1]}"
+    return ".".join(parts)
+
+
+def _load_staged_iceberg(
+    spark: SparkSession,
+    config: BatchAnalyticsConfig,
+) -> DataFrame:
+    table = _resolve_iceberg_table_name(config)
+    sql = (
+        f"SELECT {_staging_select_list(config)} FROM {table}"
+        f"{_staging_where_clause(config)}"
+    )
+    logger.info("load_staged iceberg: %s", sql)
+    return spark.sql(sql)
+
+
+def _load_staged_clickhouse(
+    spark: SparkSession,
+    config: BatchAnalyticsConfig,
+) -> DataFrame:
+    ch = config.clickhouse
+    tbl = config.transform.staging_table
+    where = _staging_where_clause(config)
+    select_list = _staging_select_list(config)
+    cat = os.environ.get("BATCH_CLICKHOUSE_CATALOG", "batch_ch").strip()
+    if cat and not where and select_list == "*":
+        try:
+            return spark.table(f"{cat}.{ch.database}.{tbl}")
+        except Exception as e:
+            logger.warning(
+                "load_staged: catalog table %s.%s.%s failed (%s), using JDBC",
+                cat,
+                ch.database,
+                tbl,
+                e,
+            )
+    elif cat and (where or select_list != "*"):
+        sql = f"SELECT {select_list} FROM {cat}.{ch.database}.{tbl}{where}"
+        try:
+            logger.info("load_staged clickhouse catalog: %s", sql)
+            return spark.sql(sql)
+        except Exception as e:
+            logger.warning(
+                "load_staged: catalog SQL failed (%s), using JDBC",
+                e,
+            )
+    dbtable = (
+        f"(SELECT {select_list} FROM `{ch.database}`.`{tbl}`{where}) AS _stg"
+    )
+    logger.info("load_staged clickhouse jdbc: %s", dbtable)
+    return spark.read.jdbc(
+        ch.jdbc_url,
+        dbtable,
+        properties=ch.jdbc_properties,
+    )

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

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

@@ -32,4 +32,5 @@ src/batch_analytics/output/local.py
 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

batch_analytics-0.3.35/tests/test_iceberg_staging.py

@@ -0,0 +1,85 @@
+"""Unit tests for Iceberg staging load helpers."""
+
+from __future__ import annotations
+
+import os
+import unittest
+from unittest.mock import MagicMock, patch
+
+from batch_analytics.config import BatchAnalyticsConfig, TransformConfig
+from batch_analytics.transform import _load_staged_iceberg, _resolve_iceberg_table_name
+
+
+class ResolveIcebergTableNameTests(unittest.TestCase):
+    def test_full_name_passthrough(self) -> None:
+        cfg = BatchAnalyticsConfig(
+            transform=TransformConfig(staging_table="nessie.gold.batch_yield_features")
+        )
+        self.assertEqual(
+            _resolve_iceberg_table_name(cfg),
+            "nessie.gold.batch_yield_features",
+        )
+
+    def test_two_part_gets_catalog(self) -> None:
+        cfg = BatchAnalyticsConfig(
+            transform=TransformConfig(staging_table="gold.batch_yield_features")
+        )
+        with patch.dict(os.environ, {"ICEBERG_CATALOG": "nessie"}, clear=False):
+            self.assertEqual(
+                _resolve_iceberg_table_name(cfg),
+                "nessie.gold.batch_yield_features",
+            )
+
+    def test_one_part_defaults_to_gold(self) -> None:
+        cfg = BatchAnalyticsConfig(
+            transform=TransformConfig(staging_table="batch_yield_features")
+        )
+        with patch.dict(os.environ, {"ICEBERG_CATALOG": "nessie"}, clear=False):
+            self.assertEqual(
+                _resolve_iceberg_table_name(cfg),
+                "nessie.gold.batch_yield_features",
+            )
+
+
+class LoadStagedIcebergTests(unittest.TestCase):
+    def test_filter_and_columns(self) -> None:
+        spark = MagicMock()
+        spark.sql.return_value = MagicMock(name="df")
+        cfg = BatchAnalyticsConfig(
+            transform=TransformConfig(
+                staging_table="nessie.gold.batch_yield_features",
+                staging_filter="product = 'Gabapentin'",
+                staging_columns="batch_no,actual_yield_qty,batch_seq",
+            )
+        )
+        _load_staged_iceberg(spark, cfg)
+        spark.sql.assert_called_once()
+        sql = spark.sql.call_args[0][0]
+        self.assertIn("SELECT batch_no, actual_yield_qty, batch_seq FROM", sql)
+        self.assertIn("WHERE product = 'Gabapentin'", sql)
+
+
+class LoadStagedClickhouseTests(unittest.TestCase):
+    def test_filter_pushed_to_jdbc(self) -> None:
+        from batch_analytics.transform import _load_staged_clickhouse
+
+        spark = MagicMock()
+        spark.read.jdbc.return_value = MagicMock(name="df")
+        cfg = BatchAnalyticsConfig(
+            transform=TransformConfig(
+                staging_table="analytics_staging",
+                staging_filter="batch_id = 'B1'",
+                staging_columns="batch_id,qty",
+            )
+        )
+        cfg.clickhouse.database = "example_db"
+        with patch.dict(os.environ, {"BATCH_CLICKHOUSE_CATALOG": ""}, clear=False):
+            _load_staged_clickhouse(spark, cfg)
+        spark.read.jdbc.assert_called_once()
+        dbtable = spark.read.jdbc.call_args[0][1]
+        self.assertIn("SELECT batch_id, qty FROM", dbtable)
+        self.assertIn("WHERE batch_id = 'B1'", dbtable)
+
+
+if __name__ == "__main__":
+    unittest.main()