sibi-dst 2025.8.6__py3-none-any.whl → 2025.8.8__py3-none-any.whl

sibi_dst/df_helper/_df_helper.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from typing import Any, Dict, Optional, TypeVar, Union
 
 import dask.dataframe as dd
@@ -15,6 +16,15 @@ from .backends.sqlalchemy import SqlAlchemyConnectionConfig, SqlAlchemyLoadFromD
 
 T = TypeVar("T", bound=BaseModel)
 
+def _is_dask_df(x) -> bool:
+    return isinstance(x, dd.DataFrame)
+
+def _maybe_persist(df, persist: bool):
+    return df.persist() if persist and _is_dask_df(df) else df
+
+def _maybe_compute(df, as_pandas: bool):
+    return df.compute() if as_pandas and _is_dask_df(df) else df
+
 
 # ---- Backend Strategy Pattern ----
 class BaseBackend:
@@ -22,13 +32,13 @@ class BaseBackend:
         self.helper = helper
         self.logger = helper.logger
         self.debug = helper.debug
-        self.total_records = helper.total_records
+        self.total_records = -1
 
     def load(self, **options) -> Union[tuple[Any, Any], dd.DataFrame, pd.DataFrame]:
         raise NotImplementedError
 
     async def aload(self, **options) -> Union[tuple[Any, Any], dd.DataFrame, pd.DataFrame]:
-        return self.load(**options)
+        return await asyncio.to_thread(self.load,**options)
 
 
 class SqlAlchemyBackend(BaseBackend):
@@ -47,7 +57,7 @@ class SqlAlchemyBackend(BaseBackend):
                 self.total_records, result = db_loader.build_and_load()
                 return self.total_records, result
         except Exception as e:
-            self.logger.error(f"Failed to load data from sqlalchemy: {e}", exc_info=self.debug)
+            self.logger.error(f"Failed to load data from sqlalchemy: {e}", exc_info=self.debug, extra=self.helper.logger_extra)
             return -1, dd.from_pandas(pd.DataFrame(), npartitions=1)
 
 
@@ -55,56 +65,59 @@ class ParquetBackend(BaseBackend):
     def load(self, **options):
         try:
             df = self.helper.backend_parquet.load_files(**options)
-            if self._is_empty(df):
-                return -1, self._empty_like(df)
-            nrows = self._row_count(df)
-            if nrows == 0:
-                self.logger.debug("No records after filters; returning empty DataFrame.")
+            if not self.helper._has_any_rows(df):
+                self.total_records = 0
                 return 0, self._empty_like(df)
 
-            df = df.persist()
-            self.total_records = self._row_count(df) or -1
+                # Let DfHelper decide about persist
+            self.total_records = -1  # unknown without full count
             return self.total_records, df
 
         except Exception as e:
             self.total_records = -1  # Reset total_records on failure
-            self.logger.error(f"Failed to load data from parquet: {e}", exc_info=self.debug)
+            self.logger.error(f"Failed to load data from parquet: {e}", exc_info=self.debug, extra=self.helper.logger_extra)
             return -1, dd.from_pandas(pd.DataFrame(), npartitions=1)
 
-    def _is_empty(self, ddf) -> bool:
-        """True if no rows across all partitions."""
-        try:
-            # head with npartitions=-1 walks partitions until it gets n rows
-            return ddf.head(1, npartitions=-1).shape[0] == 0
-        except Exception:
-            return True
-
-    def _row_count(self, ddf) -> int:
-        """Reliable row count for Dask DataFrame."""
-        return int(ddf.map_partitions(len).sum().compute())
-
-    def _empty_like(self, ddf):
-        """Return an empty Dask DF with the SAME columns/dtypes."""
+    @staticmethod
+    def _empty_like(ddf):
         empty_pdf = ddf._meta.iloc[0:0]
         return dd.from_pandas(empty_pdf, npartitions=1)
 
 
 class HttpBackend(BaseBackend):
     def load(self, **options):
-        # Will raise NotImplementedError from helper.backend_http if sync not supported
-        return self.helper.backend_http.fetch_data(**options)
+        # Avoid event-loop problems in sync code paths.
+        # If someone calls .load() on an async backend, make it explicit.
+        raise RuntimeError(
+            "HttpBackend.load() is sync but this backend is async-only. "
+            "Call `await helper.aload(...)` or `await helper.load_async(prefer_native=True, ...)`."
+        )
 
     async def aload(self, **options):
         if not self.helper.backend_http:
-            self.logger.warning("HTTP plugin not configured properly.")
+            self.logger.warning("HTTP plugin not configured properly.", extra=self.helper.logger_extra)
             self.total_records = -1
             return self.total_records, dd.from_pandas(pd.DataFrame(), npartitions=1)
+
         result = await self.helper.backend_http.fetch_data(**options)
-        self.total_records = len(result)
-        return self.total_records, result
+
+        # Normalize to DataFrame if the plugin returns list/dict
+        if isinstance(result, (list, dict)):
+            pdf = pd.DataFrame(result)
+            ddf = dd.from_pandas(pdf, npartitions=max(1, min(32, len(pdf) // 50_000 or 1)))
+            self.total_records = len(pdf)
+            return self.total_records, ddf
+
+        if isinstance(result, pd.DataFrame):
+            self.total_records = len(result)
+            ddf = dd.from_pandas(result, npartitions=max(1, min(32, len(result) // 50_000 or 1)))
+            return self.total_records, ddf
+
+        # Fallback
+        self.total_records = -1
+        return self.total_records, dd.from_pandas(pd.DataFrame(), npartitions=1)
 
 
-# ---- Main DfHelper ----
 class DfHelper(ManagedResource):
     _BACKEND_STRATEGIES = {
         "sqlalchemy": SqlAlchemyBackend,
@@ -119,6 +132,7 @@ class DfHelper(ManagedResource):
     }
 
     default_config: Dict[str, Any] = None
+    logger_extra: Dict[str, Any] = {"sibi_dst_component": __name__}
 
     def __init__(self, backend="sqlalchemy", **kwargs):
         self.default_config = self.default_config or {}
@@ -155,24 +169,25 @@ class DfHelper(ManagedResource):
     def _cleanup(self):
         attr_name = self._BACKEND_ATTR_MAP.get(self.backend)
         if not attr_name:
-            self.logger.warning(f"No attribute mapping found for backend '{self.backend}'. Cleanup skipped.")
+            self.logger.warning(f"No attribute mapping found for backend '{self.backend}'. Cleanup skipped.", extra=self.logger_extra)
             return
         active_config = getattr(self, attr_name, None)
         if active_config and hasattr(active_config, "close"):
-            self.logger.debug(f"Closing resources for '{self.backend}' backend using attribute '{attr_name}'.")
+            self.logger.debug(f"{self.__class__.__name__} is closing resources for backend '{self.backend}' backend using attribute '{attr_name}'.", extra=self.logger_extra)
             active_config.close()
 
     async def _acleanup(self):
         self.logger.warning(
-            "DfHelper instance was not used in an async context manager; cleanup is being called manually."
+            "DfHelper instance was not used in an async context manager; cleanup is being called manually.",
+            extra=self.logger_extra,
         )
         attr_name = self._BACKEND_ATTR_MAP.get(self.backend)
         if not attr_name:
-            self.logger.warning(f"No attribute mapping found for backend '{self.backend}'. Cleanup skipped.")
+            self.logger.warning(f"No attribute mapping found for backend '{self.backend}'. Cleanup skipped.", extra=self.logger_extra)
             return
         active_config = getattr(self, attr_name, None)
         if active_config and hasattr(active_config, "aclose"):
-            self.logger.debug(f"Closing resources for '{self.backend}' backend using attribute '{attr_name}'.")
+            self.logger.debug(f"Closing resources for '{self.backend}' backend using attribute '{attr_name}'.", extra=self.logger_extra)
             await active_config.aclose()
 
     # ---------- config helpers ----------
@@ -183,24 +198,49 @@ class DfHelper(ManagedResource):
 
     # ---------- load/aload ----------
     def load(self, *, persist: bool = False, as_pandas: bool = False, **options) -> Union[pd.DataFrame, dd.DataFrame]:
-        self.logger.debug(f"Loading data from {self.backend} backend with options: {options}")
+        self.logger.debug(f"Loading data from {self.backend} backend with options: {options}", extra=self.logger_extra)
         self.total_records, df = self.backend_strategy.load(**options)
         df = self._process_loaded_data(df)
         df = self._post_process_df(df)
-        #self.logger.debug(f"Finished loading data from {self.backend} backend with options: {options}")
-        df = df.persist() if persist else df
-        return df.compute() if as_pandas else df
+        df = _maybe_persist(df, persist)
+        return _maybe_compute(df, as_pandas)
+
+    async def aload(
+        self,
+        *,
+        persist: bool = False,
+        as_pandas: bool = False,
+        timeout: Optional[float] = None,
+        **options
+    ) -> Union[pd.DataFrame, dd.DataFrame]:
+        # 1) Async load if available, else run sync load in a thread.
+        if hasattr(self.backend_strategy, "aload"):
+            load_awaitable = self.backend_strategy.aload(**options)
+        else:
+            # Run ONLY the backend load step in a thread to avoid event-loop blocking.
+            load_awaitable = asyncio.to_thread(self.backend_strategy.load, **options)
 
-    async def aload(self, *, persist: bool = False, as_pandas: bool = False, **options) -> Union[pd.DataFrame, dd.DataFrame]:
-        self.total_records, df = await self.backend_strategy.aload(**options)
-        df = self._process_loaded_data(df)
-        df = self._post_process_df(df)
-        df = df.persist() if persist else df
-        return df.compute() if as_pandas else df
+        total, df = await (asyncio.wait_for(load_awaitable, timeout) if timeout else load_awaitable)
+        self.total_records = total
+
+        # 2) Post-processing steps are sync; offload to threads.
+        df = await asyncio.to_thread(self._process_loaded_data, df)
+        df = await asyncio.to_thread(self._post_process_df, df)
+
+        # 3) Persist and compute can block; offload when needed.
+        if persist and _is_dask_df(df):
+            df = await asyncio.to_thread(df.persist)
+
+        if as_pandas and _is_dask_df(df):
+            # Allow separate timeout for compute if desired; reuse same timeout here.
+            compute_awaitable = asyncio.to_thread(df.compute)
+            return await (asyncio.wait_for(compute_awaitable, timeout) if timeout else compute_awaitable)
+
+        return df
 
     # ---------- dataframe post-processing ----------
     def _post_process_df(self, df: dd.DataFrame) -> dd.DataFrame:
-        self.logger.debug("Post-processing DataFrame.")
+        self.logger.debug(f"{self.__class__.__name__} is post-processing resulting dataframe with {len(df)} records.", extra=self.logger_extra)
         df_params = self._backend_params.df_params
         if not df_params:
             return df
@@ -211,7 +251,7 @@ class DfHelper(ManagedResource):
         if fieldnames:
             valid = [f for f in fieldnames if f in df.columns]
             if len(valid) < len(fieldnames):
-                self.logger.warning(f"Missing columns for filtering: {set(fieldnames) - set(valid)}")
+                self.logger.warning(f"Missing columns for filtering: {set(fieldnames) - set(valid)}", extra=self.logger_extra)
             df = df[valid]
         if column_names:
             if len(df.columns) != len(column_names):
@@ -224,7 +264,7 @@ class DfHelper(ManagedResource):
                 raise ValueError(f"Index column '{index_col}' not found in DataFrame.")
             df = df.set_index(index_col)
 
-        self.logger.debug("Post-processing complete.")
+        self.logger.debug("Post-processing complete.", extra=self.logger_extra)
         return df
 
     def _process_loaded_data(self, df: dd.DataFrame) -> dd.DataFrame:
@@ -233,22 +273,25 @@ class DfHelper(ManagedResource):
             return df
         if hasattr(df, "npartitions") and df.npartitions == 1 and not len(df.head(1)):
             return df
-        self.logger.debug("Applying rename mapping if necessary.")
+        self.logger.debug(f"{self.__class__.__name__} is applying rename mapping if/when necessary.", extra=self.logger_extra)
         rename_map = {k: v for k, v in field_map.items() if k in df.columns}
         if rename_map:
             df = df.rename(columns=rename_map)
         return df
 
     # ---------- sinks ----------
-    def save_to_parquet(self, df: dd.DataFrame, parquet_filename: str, **kwargs):
+    def save_to_parquet(self, df: dd.DataFrame, **kwargs):
         fs: AbstractFileSystem = kwargs.pop("fs", self.fs)
-        path: str = kwargs.pop("parquet_storage_path")
+        path: str = kwargs.pop("parquet_storage_path", self.backend_parquet.parquet_storage_path if self.backend_parquet else None)
+        parquet_filename = kwargs.pop("parquet_filename", self.backend_parquet.parquet_filename if self.backend_parquet else None)
+        if not parquet_filename:
+            raise ValueError("A 'parquet_filename' keyword argument must be provided.")
         if not fs:
             raise ValueError("A filesystem (fs) must be provided to save the parquet file.")
         if not path:
             raise ValueError("A 'parquet_storage_path' keyword argument must be provided.")
-        if len(df.head(1)) == 0:
-            self.logger.warning("Skipping save: The provided DataFrame is empty.")
+        if not self._has_any_rows(df):
+            self.logger.warning("Skipping save: The provided DataFrame is empty.", extra=self.logger_extra)
             return
 
         with ParquetSaver(
@@ -262,17 +305,17 @@ class DfHelper(ManagedResource):
         ) as saver:
             saver.save_to_parquet(parquet_filename)
 
-        self.logger.debug(f"Successfully saved '{parquet_filename}' to '{path}'.")
+        self.logger.debug(f"Successfully saved '{parquet_filename}' to '{path}'.", extra=self.logger_extra)
 
     def save_to_clickhouse(self, df: dd.DataFrame, **credentials):
-        if hasattr(df, "npartitions") and df.npartitions == 1 and not len(df.head(1)):
-            self.logger.warning("Cannot write to ClickHouse; DataFrame is empty.")
+        if not self._has_any_rows(df):
+            self.logger.warning("Skipping save to ClickHouse: The provided DataFrame is empty.", extra=self.logger_extra)
             return
         with ClickHouseWriter(debug=self.debug, logger=self.logger, verbose=self.verbose, **credentials) as writer:
             writer.save_to_clickhouse(df)
-            self.logger.debug("Save to ClickHouse completed.")
+            self.logger.debug("Save to ClickHouse completed.", extra=self.logger_extra)
 
-    # ---------- convenience period loaders ----------
+    # ---------- period loaders ----------
     def load_period(self, dt_field: str, start: str, end: str, **kwargs):
         final_kwargs = self._prepare_period_filters(dt_field, start, end, **kwargs)
         return self.load(**final_kwargs)
@@ -288,13 +331,20 @@ class DfHelper(ManagedResource):
         field_map = self._backend_params.field_map or {}
         reverse_map = {v: k for k, v in field_map.items()} if field_map else {}
         if len(reverse_map) != len(field_map):
-            self.logger.warning("field_map values are not unique; reverse mapping may be unreliable.")
+            self.logger.warning("field_map values are not unique; reverse mapping may be unreliable.", extra=self.logger_extra)
         mapped_field = reverse_map.get(dt_field, dt_field)
         if start_date == end_date:
             kwargs[f"{mapped_field}__date"] = start_date
         else:
             kwargs[f"{mapped_field}__date__range"] = [start_date, end_date]
-        self.logger.debug(f"Period load generated filters: {kwargs}")
+        self.logger.debug(f"Period load generated filters: {kwargs}", extra=self.logger_extra)
         return kwargs
 
+    @staticmethod
+    def _has_any_rows(ddf: dd.DataFrame) -> bool:
+        try:
+            return bool(ddf.head(1, npartitions=-1).shape[0])
+        except Exception:
+            return False
+
 

sibi_dst/df_helper/_parquet_artifact.py

@@ -23,6 +23,7 @@ class ParquetArtifact(ManagedResource):
 
     _global_lock = threading.RLock()
     _active_runs: set[tuple[str, str]] = set()
+    logger_extra = {"sibi_dst_component": __name__}
 
     def __init__(self, **kwargs: Any):
         # Merge defaults from ManagedResource and caller kwargs
@@ -49,7 +50,7 @@ class ParquetArtifact(ManagedResource):
     # ---------- lazy members ----------
     @cached_property
     def mmanifest(self) -> MissingManifestManager:
-        self.logger.info("Initializing MissingManifestManager...")
+        self.logger.info("Initializing MissingManifestManager...", extra=self.logger_extra)
         manifest_path = self._build_manifest_path()
 
         # ensure manifest directory exists
@@ -66,16 +67,16 @@ class ParquetArtifact(ManagedResource):
         )
 
         if not mgr._safe_exists(mgr.manifest_path):
-            self.logger.info(f"Creating new manifest at {mgr.manifest_path}")
+            self.logger.info(f"Creating new manifest at {mgr.manifest_path}", extra=self.logger_extra)
             mgr.save()
         else:
-            self.logger.info(f"Manifest already exists at {mgr.manifest_path}")
+            self.logger.info(f"Manifest already exists at {mgr.manifest_path}", extra=self.logger_extra)
 
         return mgr
 
     @cached_property
     def update_planner(self) -> UpdatePlanner:
-        self.logger.info("Initializing UpdatePlanner...")
+        self.logger.info("Initializing UpdatePlanner...", extra=self.logger_extra)
         skipped_files = self.mmanifest.load_existing() or []
 
         cfg = {
@@ -91,7 +92,7 @@ class ParquetArtifact(ManagedResource):
 
     @cached_property
     def data_wrapper(self) -> DataWrapper:
-        self.logger.info("Initializing DataWrapper...")
+        self.logger.info("Initializing DataWrapper...", extra=self.logger_extra)
 
         # Ensure the planner has a plan
         if getattr(self.update_planner, "plan", None) is None:
@@ -170,7 +171,7 @@ class ParquetArtifact(ManagedResource):
         with ParquetArtifact._global_lock:
             if key in ParquetArtifact._active_runs:
                 self.logger.info(
-                    f"Run already in progress for {key}; skipping this invocation."
+                    f"Run already in progress for {key}; skipping this invocation.", extra=self.logger_extra
                 )
                 return
             ParquetArtifact._active_runs.add(key)
@@ -182,7 +183,7 @@ class ParquetArtifact(ManagedResource):
             plan = getattr(self.update_planner, "plan", None)
             if plan is None or (hasattr(plan, "empty") and plan.empty):
                 # Planning uses Pandas; this is safe to check.
-                self.logger.info("No updates needed. Skipping Parquet generation.")
+                self.logger.info("No updates needed. Skipping Parquet generation.", extra=self.logger_extra)
                 return
 
             # Print plan once per run
@@ -286,7 +287,7 @@ class ParquetArtifact(ManagedResource):
 
         final_kwargs.update(period_params)
         self.logger.debug(
-            f"kwargs passed to update_parquet/generate_parquet: {final_kwargs}"
+            f"kwargs passed to update_parquet/generate_parquet: {final_kwargs}", extra=self.logger_extra
         )
 
         # Delegate to generator (handles cache invalidation + forwarding knobs)
@@ -297,7 +298,7 @@ class ParquetArtifact(ManagedResource):
         """Ensure the directory exists across fsspec backends."""
         with self._lock:
             if not self.fs.exists(path):
-                self.logger.info(f"Creating directory: {path}")
+                self.logger.info(f"Creating directory: {path}", extra=self.logger_extra)
                 try:
                     self.fs.makedirs(path, exist_ok=True)
                 except TypeError:
@@ -317,4 +318,4 @@ class ParquetArtifact(ManagedResource):
             if "data_wrapper" in self.__dict__ and hasattr(self.data_wrapper, "close"):
                 self.data_wrapper.close()
         except Exception as e:
-            self.logger.warning(f"Error during resource cleanup: {e}")
+            self.logger.warning(f"Error during resource cleanup: {e}", extra=self.logger_extra)

sibi_dst/df_helper/_parquet_reader.py

@@ -70,6 +70,10 @@ class ParquetReader(DfHelper):
         self.df = super().load(**kwargs)
         return self.df
 
+    async def aload(self, **kwargs) -> Union[pd.DataFrame, dd.DataFrame]:
+        self.df = await super().aload(**kwargs)
+        return self.df
+
     def directory_exists(self):
         try:
             info = self.fs.info(self.parquet_storage_path)