sibi-dst 2025.1.10__tar.gz → 2025.1.12__tar.gz

{sibi_dst-2025.1.10 → sibi_dst-2025.1.12}/PKG-INFO

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

{sibi_dst-2025.1.10 → sibi_dst-2025.1.12}/pyproject.toml

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

{sibi_dst-2025.1.10 → sibi_dst-2025.1.12}/sibi_dst/df_helper/_df_helper.py

@@ -68,6 +68,9 @@ class ParquetBackend(BaseBackend):
             df = self.helper.backend_parquet.load_files()
             if options and df is not None:
                 df = FilterHandler('dask', logger=self.logger, debug=False).apply_filters(df, filters=options)
+
+            df = df.persist()
+
             self.total_records = len(df) or -1  # If df is empty, set total_records to -1
             return self.total_records, df
         except Exception as e:

{sibi_dst-2025.1.10 → sibi_dst-2025.1.12}/sibi_dst/df_helper/backends/sqlalchemy/_db_connection.py

@@ -15,7 +15,7 @@ from sqlalchemy.engine import url as sqlalchemy_url
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import OperationalError, SQLAlchemyError
 from sqlalchemy.orm import sessionmaker, Session
-from sqlalchemy.pool import QueuePool, NullPool, StaticPool
+from sqlalchemy.pool import QueuePool, NullPool, StaticPool, Pool
 
 # Assuming these are your project's internal modules
 from sibi_dst.utils import Logger
@@ -54,7 +54,7 @@ class SqlAlchemyConnectionConfig(BaseModel):
     pool_timeout: int = int(os.environ.get("DB_POOL_TIMEOUT", 30))
     pool_recycle: int = int(os.environ.get("DB_POOL_RECYCLE", 1800))
     pool_pre_ping: bool = True
-    poolclass: Type[QueuePool] = QueuePool
+    poolclass: Type[Pool] = QueuePool
 
     # --- Internal & Runtime State ---
     model: Optional[Type[Any]] = None
@@ -195,7 +195,6 @@ class SqlAlchemyConnectionConfig(BaseModel):
             wrapper = self._engine_registry.get(self._engine_key_instance)
             if wrapper:
                 wrapper['active_connections'] += 1
-        # self.logger.debug(f"Connection checked out. Active: {self.active_connections}")
 
     def _on_checkin(self, *args) -> None:
         """Event listener for when a connection is returned to the pool."""
@@ -203,7 +202,6 @@ class SqlAlchemyConnectionConfig(BaseModel):
             wrapper = self._engine_registry.get(self._engine_key_instance)
             if wrapper:
                 wrapper['active_connections'] = max(0, wrapper['active_connections'] - 1)
-        # self.logger.debug(f"Connection checked in. Active: {self.active_connections}")
 
     @property
     def active_connections(self) -> int:

{sibi_dst-2025.1.10 → sibi_dst-2025.1.12}/sibi_dst/utils/date_utils.py

@@ -1,10 +1,12 @@
+from __future__ import annotations
+
 import datetime
 from typing import Union, Tuple, Callable, Dict, Optional
 
 import fsspec
 import numpy as np
 import pandas as pd
-
+import dask.dataframe as dd
 from .log_utils import Logger
 
 
@@ -305,154 +307,152 @@ class FileAgeChecker:
             raise ValueError(f"Unsupported modification time format for {file_path}") from e
 
 
-class BusinessDays:
+# --- Vectorized Helper Functions ---
+
+def _vectorized_busday_count(partition, begin_col, end_col, holidays):
     """
-    Provides functionality for handling business days calculations with a custom
-    holiday list. The class includes methods for calculating the number of
-    business days, modifying dates by adding business days, and applying these
-    operations to Dask DataFrames.
-
-    :ivar logger: Logger instance for logging error, warning, and debug messages.
-    :type logger: logging.Logger
-    :ivar HOLIDAY_LIST: Dictionary mapping years to lists of holiday dates.
-    :type HOLIDAY_LIST: dict
-    :ivar bd_cal: Numpy busdaycalendar object containing holidays and week mask.
-    :type bd_cal: numpy.busdaycalendar
-    :ivar holidays: Array of holiday dates used by the business day calendar.
-    :type holidays: numpy.ndarray
-    :ivar week_mask: Boolean array indicating working days within a week.
-    :type week_mask: numpy.ndarray
+    Calculates the number of business days between a start and end date.
     """
+    # Extract the raw columns
+    start_dates_raw = partition[begin_col]
+    end_dates_raw = partition[end_col]
 
-    def __init__(self, holiday_list, logger):
-        """
-        Initialize a BusinessDays object with a given holiday list.
-        """
-        self.logger = logger
-        self.HOLIDAY_LIST = holiday_list
-        bd_holidays = [day for year in self.HOLIDAY_LIST for day in self.HOLIDAY_LIST[year]]
-        self.bd_cal = np.busdaycalendar(holidays=bd_holidays, weekmask="1111100")
-        self.holidays = self.bd_cal.holidays
-        self.week_mask = self.bd_cal.weekmask
 
-    def get_business_days_count(self, begin_date, end_date):
-        """
-        Calculate the number of business days between two dates.
-        """
-        try:
-            begin_date = pd.to_datetime(begin_date)
-            end_date = pd.to_datetime(end_date)
-        except Exception as e:
-            raise ValueError(f"Invalid date format: {e}")
+    start_dates = pd.to_datetime(start_dates_raw, errors='coerce')
+    end_dates = pd.to_datetime(end_dates_raw, errors='coerce')
 
-        years = [str(year) for year in range(begin_date.year, end_date.year + 1)]
-        if not all(year in self.HOLIDAY_LIST for year in years):
-            raise ValueError("Not all years in date range are in the holiday list")
+    # Initialize the result Series with NaN, as the output is a number
+    result = pd.Series(np.nan, index=partition.index)
 
-        return np.busday_count(
-            begin_date.strftime("%Y-%m-%d"),
-            end_date.strftime("%Y-%m-%d"),
-            busdaycal=self.bd_cal,
-        )
+    # Create a mask for rows where both start and end dates are valid
+    valid_mask = pd.notna(start_dates) & pd.notna(end_dates)
 
-    def calc_business_days_from_df(self, df, begin_date_col, end_date_col, result_col="business_days"):
-        """
-        Add a column to a Dask DataFrame with the number of business days between two date columns.
-        """
-        if not all(col in df.columns for col in [begin_date_col, end_date_col]):
-            self.logger.error("Column names not found in DataFrame")
-            raise ValueError("Required columns are missing")
-
-        # Extract holidays and weekmask to recreate the busdaycalendar
-        holidays = self.bd_cal.holidays
-        weekmask = self.bd_cal.weekmask
-
-        # Define a function to calculate business days
-        def calculate_business_days(row, holidays, weekmask):
-            begin_date = pd.to_datetime(row[begin_date_col])
-            end_date = pd.to_datetime(row[end_date_col])
-            if pd.isna(begin_date) or pd.isna(end_date):
-                return np.nan
-            busdaycal = np.busdaycalendar(holidays=holidays, weekmask=weekmask)
-            return np.busday_count(
-                begin_date.strftime("%Y-%m-%d"),
-                end_date.strftime("%Y-%m-%d"),
-                busdaycal=busdaycal,
-            )
-
-        # Define a wrapper function for partition-wise operations
-        def apply_business_days(partition, holidays, weekmask):
-            return partition.apply(
-                calculate_business_days, axis=1, holidays=holidays, weekmask=weekmask
-            )
-
-        # Apply the function using map_partitions
-        df[result_col] = df.map_partitions(
-            apply_business_days,
-            holidays,
-            weekmask,
-            meta=(result_col, "int64"),
-        )
+    # Perform the vectorized calculation only on the valid subset
+    # Convert to NumPy arrays of date type for the calculation
+    result.loc[valid_mask] = np.busday_count(
+        start_dates[valid_mask].values.astype('datetime64[D]'),
+        end_dates[valid_mask].values.astype('datetime64[D]'),
+        holidays=holidays
+    )
 
-        return df
+    return result
+
+
+def _vectorized_sla_end_date(partition, start_col, n_days_col, holidays):
+    """
+    Calculates the end date of an SLA, skipping weekends and holidays.
+    """
+    # Extract the relevant columns as pandas Series
+    start_dates_raw = partition[start_col]
+    sla_days = partition[n_days_col]
 
-    def add_business_days(self, start_date, n_days):
-        """
-        Add n_days business days to start_date.
-        """
-        try:
-            start_date = datetime.datetime.strptime(start_date, "%Y-%m-%d")
-        except ValueError:
-            raise ValueError("Date should be a string in the format YYYY-MM-DD")
 
-        if str(start_date.year) not in self.HOLIDAY_LIST:
-            self.logger.warning(f"Year {start_date.year} is not in the holiday list")
+    start_dates = pd.to_datetime(start_dates_raw, errors='coerce')
 
+    # Initialize the result Series with NaT (Not a Time)
+    result = pd.Series(pd.NaT, index=partition.index, dtype='datetime64[ns]')
+
+    # Create a mask for rows that have valid start dates and SLA days
+    valid_mask = pd.notna(start_dates) & pd.notna(sla_days)
+
+    # Perform the vectorized calculation only on the valid subset
+    # Note: np.busday_offset requires a NumPy array, so we use .values
+    result.loc[valid_mask] = np.busday_offset(
+        start_dates[valid_mask].values.astype('datetime64[D]'),  # Convert to numpy array of dates
+        sla_days[valid_mask].astype(int),  # Ensure days are integers
+        roll='forward',
+        holidays=holidays
+    )
+
+    return result
+
+
+# --- Refactored BusinessDays Class ---
+
+class BusinessDays:
+    """
+    Business days calculations with a custom holiday list.
+    Supports scalar and efficient, vectorized Dask DataFrame operations.
+    """
+
+    def __init__(self, holiday_list: dict[str, list[str]], logger) -> None:
+        self.logger = logger
+        self.HOLIDAY_LIST = holiday_list
+
+        # Flatten and store as tuple for determinism
+        bd_holidays = [day for year in self.HOLIDAY_LIST for day in self.HOLIDAY_LIST[year]]
+        self.holidays = tuple(bd_holidays)
+
+    def get_business_days_count(
+            self,
+            begin_date: str | datetime.date | pd.Timestamp,
+            end_date: str | datetime.date | pd.Timestamp,
+    ) -> int:
+        """Scalar method to count business days between two dates."""
+        begin = pd.to_datetime(begin_date)
+        end = pd.to_datetime(end_date)
+        return int(np.busday_count(begin.date(), end.date(), holidays=list(self.holidays)))
+
+    def calc_business_days_from_df(
+            self,
+            df: dd.DataFrame,
+            begin_date_col: str,
+            end_date_col: str,
+            result_col: str = "business_days",
+    ) -> dd.DataFrame:
+        """Calculates business days between two columns in a Dask DataFrame."""
+        missing = {begin_date_col, end_date_col} - set(df.columns)
+        if missing:
+            self.logger.error(f"Missing columns: {missing}")
+            raise ValueError("Required columns are missing from DataFrame")
+
+        return df.assign(
+            **{result_col: df.map_partitions(
+                _vectorized_busday_count,
+                begin_col=begin_date_col,
+                end_col=end_date_col,
+                holidays=list(self.holidays),
+                meta=(result_col, 'f8')  # f8 is float64
+            )}
+        )
+
+    def add_business_days(
+            self,
+            start_date: str | datetime.date | pd.Timestamp,
+            n_days: int,
+    ) -> np.datetime64:
+        """Scalar method to add N business days to a start date."""
+        start = pd.to_datetime(start_date)
         return np.busday_offset(
-            start_date.strftime("%Y-%m-%d"),
+            start.date(),
             n_days,
-            roll="forward",
-            busdaycal=self.bd_cal,
+            roll='forward',
+            holidays=list(self.holidays),
         )
 
-    def calc_sla_end_date(self, df, start_date_col, n_days_col, result_col="sla_end_date"):
-        """
-        Add a column to a Dask DataFrame with SLA end dates based on start date and SLA days.
-        """
-        if not all(col in df.columns for col in [start_date_col, n_days_col]):
-            raise ValueError("Column names not found in DataFrame")
-
-        # Extract holidays and weekmask to recreate the busdaycalendar
-        holidays = self.bd_cal.holidays
-        weekmask = self.bd_cal.weekmask
-
-        # Define a function to calculate SLA end dates
-        def calculate_sla_end_date(row, holidays, weekmask):
-            start_date = pd.to_datetime(row[start_date_col])
-            n_days = row[n_days_col]
-            busdaycal = np.busdaycalendar(holidays=holidays, weekmask=weekmask)
-            return np.busday_offset(
-                start_date.strftime("%Y-%m-%d"),
-                n_days,
-                roll="forward",
-                busdaycal=busdaycal,
-            )
-
-        # Define a wrapper for partition-wise operation
-        def apply_sla_end_date(partition, holidays, weekmask):
-            return partition.apply(
-                calculate_sla_end_date, axis=1, holidays=holidays, weekmask=weekmask
-            )
-
-        # Apply the function using map_partitions
-        df[result_col] = df.map_partitions(
-            apply_sla_end_date,
-            holidays,
-            weekmask,
-            meta=(result_col, "object"),
+    def calc_sla_end_date(
+            self,
+            df: dd.DataFrame,
+            start_date_col: str,
+            n_days_col: str,
+            result_col: str = "sla_end_date",
+    ) -> dd.DataFrame:
+        """Calculates an SLA end date column for a Dask DataFrame."""
+        missing = {start_date_col, n_days_col} - set(df.columns)
+        if missing:
+            self.logger.error(f"Missing columns: {missing}")
+            raise ValueError("Required columns are missing from DataFrame")
+
+        return df.assign(
+            **{result_col: df.map_partitions(
+                _vectorized_sla_end_date,
+                start_col=start_date_col,
+                n_days_col=n_days_col,
+                holidays=list(self.holidays),
+                meta=(result_col, 'datetime64[ns]')
+            )}
         )
 
-        return df
 # Class enhancements
 # DateUtils.register_period('next_week', lambda: (datetime.date.today() + datetime.timedelta(days=7),
 #                                                 datetime.date.today() + datetime.timedelta(days=13)))

{sibi_dst-2025.1.10 → sibi_dst-2025.1.12}/sibi_dst/v2/df_helper/backends/sqlmodel/_io_dask.py

@@ -1,7 +1,8 @@
 import itertools
 import dask.dataframe as dd
 import pandas as pd
-from sqlmodel import create_engine, Session, select
+
+#from sqlmodel import create_engine, Session, select
 from sibi_dst.v2.df_helper.core import FilterHandler
 from sibi_dst.v2.utils import Logger
 
@@ -116,7 +117,7 @@ class SQLModelDask:
                 return dask_df
 
         except Exception as e:
-            self.logger.error(f"Error executing query: {str(e)}")
-            self.logger.error(self.query)
+            self.logger.error(f"_io_dask:Error executing query: {str(e)}")
+            self.logger.error(f"_io_dask:{self.query})
             # In case of error, return an empty Dask DataFrame with the expected columns.
             return dd.from_pandas(pd.DataFrame(columns=ordered_columns), npartitions=1)