sibi-dst 2025.9.6__tar.gz → 2025.9.8__tar.gz

{sibi_dst-2025.9.6 → sibi_dst-2025.9.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sibi-dst
-Version: 2025.9.6
+Version: 2025.9.8
 Summary: Data Science Toolkit
 Author: Luis Valverde
 Author-email: lvalverdeb@gmail.com

{sibi_dst-2025.9.6 → sibi_dst-2025.9.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sibi-dst"
-version = "2025.9.6"
+version = "2025.9.8"
 description = "Data Science Toolkit"
 authors = ["Luis Valverde <lvalverdeb@gmail.com>"]
 readme = "README.md"

{sibi_dst-2025.9.6 → sibi_dst-2025.9.8}/sibi_dst/utils/boilerplate/__init__.py

@@ -3,6 +3,7 @@ from .base_data_cube import BaseDataCube
 from .base_attacher import make_attacher
 from .base_parquet_reader import BaseParquetReader
 from .hybrid_data_loader import HybridDataLoader
+from .base_pipeline import BasePipeline
 
 __all__ = [
     "BaseDataCube",
@@ -10,5 +11,6 @@ __all__ = [
     "make_attacher",
     "BaseParquetReader",
     "HybridDataLoader",
+    "BasePipeline",
 ]
 

sibi_dst-2025.9.8/sibi_dst/utils/boilerplate/base_pipeline.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from typing import Type, Any, Callable, List
+
+import pandas as pd
+import dask.dataframe as dd
+
+from sibi_dst.utils import ManagedResource, ParquetSaver
+from sibi_dst.df_helper import ParquetReader
+from sibi_dst.utils.dask_utils import dask_is_empty
+
+
+class DateRangeHelper:
+    @staticmethod
+    def generate_daily_ranges(start_date: str, end_date: str, date_format: str = "%Y-%m-%d") -> List[str]:
+        start = pd.to_datetime(start_date)
+        end = pd.to_datetime(end_date)
+        return [d.strftime(date_format) for d in pd.date_range(start, end, freq="D")]
+
+    @staticmethod
+    def generate_monthly_ranges(start_date: str, end_date: str, date_format: str = "%Y-%m-%d") -> List[tuple[str, str]]:
+        """
+        Generate (start_date, end_date) tuples for each calendar month in range.
+        Always includes the first and last month, even if partial.
+        """
+        start = pd.to_datetime(start_date)
+        end = pd.to_datetime(end_date)
+        ranges = []
+        current = start.replace(day=1)
+        while current <= end:
+            month_end = (current + pd.offsets.MonthEnd(0)).normalize()
+            ranges.append((
+                current.strftime(date_format),
+                min(month_end, end).strftime(date_format)
+            ))
+            current += pd.DateOffset(months=1)
+        return ranges
+
+class BasePipeline(ManagedResource):
+    def __init__(
+        self,
+        start_date: str,
+        end_date: str,
+        dataset_cls: Type,
+        parquet_storage_path: str,
+        *,
+        fs: Any,
+        filename: str = "dataset",
+        date_field: str = "date",
+        max_workers: int = 4,
+        dataset_kwargs: dict = None,
+        **kwargs,
+    ):
+        kwargs["fs"] = fs
+        super().__init__(**kwargs)
+
+        self.start_date = start_date
+        self.end_date = end_date
+        self.fs = fs
+        self.filename = filename
+        self.date_field = date_field
+        self.max_workers = max_workers
+        self.storage_path = parquet_storage_path.rstrip("/")
+        self.df: dd.DataFrame | None = None
+
+        self.ds = dataset_cls(
+            start_date=self.start_date,
+            end_date=self.end_date,
+            debug=self.debug,
+            logger=self.logger,
+            **(dataset_kwargs or {}),
+        )
+
+    def _get_storage_path_for_date(self, date: pd.Timestamp) -> str:
+        return f"{self.storage_path}/{date.year}/{date.month:02d}/{date.day:02d}"
+
+    def _get_output_filename(self, fmt: str = "parquet") -> str:
+        return f"{self.filename}.{fmt}"
+
+    async def aload(self, **kwargs) -> dd.DataFrame:
+        await self.emit("status", message="Loading dataset...", progress=5)
+        self.df = await self.ds.aload(**kwargs)
+        return self.df
+
+    async def to_parquet(self, **kwargs) -> None:
+        df = await self.aload(**kwargs)
+        if dask_is_empty(df):
+            self.logger.warning("No data to save.")
+            return
+
+        df[self.date_field] = dd.to_datetime(df[self.date_field], errors="coerce")
+        dates = DateRangeHelper.generate_daily_ranges(self.start_date, self.end_date)
+
+        tasks = []
+        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            for date_str in dates:
+                date_obj = pd.to_datetime(date_str).date()
+                df_day = df[df[self.date_field].dt.date == date_obj]
+                if dask_is_empty(df_day):
+                    self.logger.info(f"No data for {date_obj}, skipping.")
+                    continue
+
+                path = self._get_storage_path_for_date(pd.Timestamp(date_obj))
+                await self.emit("status", message=f"Saving data for {date_obj}")
+
+                saver = ParquetSaver(
+                    df_result=df_day,
+                    parquet_storage_path=path,
+                    fs=self.fs,
+                    debug=self.debug,
+                    logger=self.logger,
+                )
+
+                tasks.append(
+                    asyncio.get_running_loop().run_in_executor(
+                        executor, saver.save_to_parquet, self._get_output_filename()
+                    )
+                )
+
+            await asyncio.gather(*tasks)
+        await self.emit("complete", message="All partitions written.")
+
+    async def from_parquet(self, **kwargs) -> dd.DataFrame:
+        reader = ParquetReader(
+            parquet_start_date=self.start_date,
+            parquet_end_date=self.end_date,
+            parquet_storage_path=self.storage_path,
+            parquet_filename=self._get_output_filename(),
+            fs=self.fs,
+            debug=self.debug,
+            logger=self.logger,
+        )
+        return await reader.aload(**kwargs)
+
+    async def to_clickhouse(self, clk_conf: dict, **kwargs):
+        """
+        Writes daily-partitioned data to ClickHouse using concurrent threads.
+        """
+        from sibi_dst.utils import ClickHouseWriter
+
+        df = await self.from_parquet(**kwargs)
+        if dask_is_empty(df):
+            self.logger.warning("No data to write to ClickHouse.")
+            return
+
+        df[self.date_field] = dd.to_datetime(df[self.date_field], errors="coerce")
+        df = df.persist()
+
+        unique_dates = df[self.date_field].dt.date.dropna().unique().compute()
+        if len(unique_dates)==0:
+            self.logger.warning("No valid dates found for partitioning.")
+            return
+
+        clk = ClickHouseWriter(**clk_conf)
+        loop = asyncio.get_running_loop()
+        tasks = []
+
+        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            for date in unique_dates:
+                df_day = df[df[self.date_field].dt.date == date]
+                if dask_is_empty(df_day):
+                    self.logger.info(f"[ClickHouse] No data for {date}, skipping.")
+                    continue
+
+                self.logger.info(f"[ClickHouse] Writing {len(df_day)} rows for {date}")
+
+                tasks.append(
+                    loop.run_in_executor(executor, clk.save_to_clickhouse, df_day)
+                )
+
+            await asyncio.gather(*tasks)
+
+        self.logger.info(f"ClickHouse write complete for {len(unique_dates)} daily partitions.")
+
+
+__all__ = ["BasePipeline"]

{sibi_dst-2025.9.6 → sibi_dst-2025.9.8}/sibi_dst/utils/clickhouse_writer.py

@@ -90,10 +90,14 @@ class ClickHouseWriter(ManagedResource):
 
         # one client per thread to avoid session contention
         self._tlocal = threading.local()
+        ow = self.overwrite
+        if ow:
+            self._command(f"DROP TABLE IF EXISTS {self._ident(self.table)}")
+            self.logger.info(f"Dropped table {self.table} (overwrite=True)")
 
     # ------------- public -------------
 
-    def save_to_clickhouse(self, df: dd.DataFrame, *, overwrite: Optional[bool] = None) -> None:
+    def save_to_clickhouse(self, df: dd.DataFrame) -> None:
         """
         Persist a Dask DataFrame into ClickHouse.
 
@@ -118,15 +122,10 @@ class ClickHouseWriter(ManagedResource):
         )
 
         # (re)create table
-        ow = self.overwrite if overwrite is None else bool(overwrite)
         dtypes = df._meta_nonempty.dtypes  # metadata-only types (no compute)
         schema_sql = self._generate_clickhouse_schema(dtypes)
         engine_sql = self._default_engine_sql() if not self.engine else self.engine
 
-        if ow:
-            self._command(f"DROP TABLE IF EXISTS {self._ident(self.table)}")
-            self.logger.info(f"Dropped table {self.table} (overwrite=True)")
-
         create_sql = f"CREATE TABLE IF NOT EXISTS {self._ident(self.table)} ({schema_sql}) {engine_sql}"
         self._command(create_sql)
         self.logger.info(f"Ensured table {self.table} exists")