sibi-dst 0.3.10__py3-none-any.whl → 0.3.11__py3-none-any.whl

sibi_dst/df_helper/_df_helper.py

@@ -101,6 +101,8 @@ class DfHelper:
 
     def _load_from_sqlalchemy(self, **options):
         try:
+            options.setdefault("debug", self.debug)
+            options.setdefault("verbose_debug", self.verbose_debug)
             db_loader = SqlAlchemyLoadFromDb(
                 self.plugin_sqlalchemy,
                 self.plugin_query,
@@ -120,6 +122,8 @@ class DfHelper:
 
     def _load_from_db(self, **options) -> Union[pd.DataFrame, dd.DataFrame]:
         try:
+            options.setdefault("debug", self.debug)
+            options.setdefault("verbose_debug", self.verbose_debug)
             db_loader = DjangoLoadFromDb(
                 self.plugin_django_connection,
                 self.plugin_query,

sibi_dst/df_helper/core/__init__.py

@@ -2,10 +2,17 @@ from __future__ import annotations
 
 from ._params_config import ParamsConfig
 from ._query_config import QueryConfig
-from ._defaults import django_field_conversion_map
+from ._defaults import (
+    django_field_conversion_map_pandas,
+    django_field_conversion_map_dask,
+    sqlalchemy_field_conversion_map_dask,
+    normalize_sqlalchemy_type)
 
 __all__ = [
     "ParamsConfig",
     "QueryConfig",
-    "django_field_conversion_map"
+    "django_field_conversion_map_pandas",
+    "django_field_conversion_map_dask",
+    "sqlalchemy_field_conversion_map_dask",
+    "normalize_sqlalchemy_type"
 ]

sibi_dst/df_helper/core/_defaults.py

@@ -1,10 +1,12 @@
 #  Copyright (c) 2023. ISTMO Center S.A.  All Rights Reserved
 #
 import json
-from dataclasses import dataclass
-from typing import Dict, Union, Optional
+from typing import Dict
 
 import pandas as pd
+from sqlalchemy import String, Text, Integer, BigInteger, SmallInteger, Float, Boolean, DateTime, Date, Time, JSON, \
+    Numeric, UUID
+from sqlalchemy.dialects.mysql import TINYINT, MEDIUMTEXT
 
 # This is the defaults configuration file for the df_helper module.
 
@@ -13,11 +15,12 @@ import pandas as pd
 # conversion. This mapping is used to convert the values in a pandas DataFrame to the appropriate data types based on
 # the Django field type.
 
-django_field_conversion_map: Dict[str, callable] = {
+django_field_conversion_map_pandas: Dict[str, callable] = {
     "CharField": lambda x: x.astype(str),
     "TextField": lambda x: x.astype(str),
     "IntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
     "AutoField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "BigAutoField": lambda x: pd.to_numeric(x, errors="coerce"),
     "BigIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
     "SmallIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
     "PositiveIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
@@ -35,3 +38,96 @@ django_field_conversion_map: Dict[str, callable] = {
     "ArrayField": lambda x: x.apply(eval),
     "UUIDField": lambda x: x.astype(str),
 }
+
+django_field_conversion_map_dask: Dict[str, callable] = {
+    "CharField": lambda x: x.astype(str),
+    "TextField": lambda x: x.astype(str),
+    "IntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "AutoField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "BigAutoField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "BigIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "SmallIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "PositiveIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "PositiveSmallIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "FloatField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "DecimalField": lambda x: pd.to_numeric(x, errors="coerce"),
+    "BooleanField": lambda x: x.astype(bool),
+    "NullBooleanField": lambda x: x.astype(bool),
+    "DateTimeField": lambda x: pd.to_datetime(x, errors="coerce"),
+    "DateField": lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda x: x.dt.date, meta=("date", "object")),
+    "TimeField": lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda x: x.dt.time, meta=("time", "object")),
+    "DurationField": lambda x: pd.to_timedelta(x, errors="coerce"),
+    "JSONField": lambda x: x.map_partitions(lambda s: s.apply(json.loads), meta=("json", "object")),
+    "ArrayField": lambda x: x.map_partitions(lambda s: s.apply(eval), meta=("array", "object")),
+    "UUIDField": lambda x: x.astype(str),
+}
+
+sqlalchemy_field_conversion_map_dask: Dict[str, callable] = {
+    String.__name__: lambda x: x.astype(str).fillna(""),
+    Text.__name__: lambda x: x.fillna('').astype(str),
+    Integer.__name__: lambda x: x.fillna(0).astype(int),
+    BigInteger.__name__: lambda x: pd.to_numeric(x, errors="coerce"),
+    SmallInteger.__name__: lambda x: pd.to_numeric(x, errors="coerce"),
+    Float.__name__: lambda x: pd.to_numeric(x, errors="coerce"),
+    Numeric.__name__: lambda x: pd.to_numeric(x, errors="coerce"),
+    Boolean.__name__: lambda x: x.astype(bool),
+    DateTime.__name__: lambda x: pd.to_datetime(x, errors="coerce"),
+    Date.__name__: lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda x: x.dt.date, meta=("date", "object")),
+    Time.__name__: lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda x: x.dt.time, meta=("time", "object")),
+    JSON.__name__: lambda x: x.map_partitions(lambda s: s.apply(json.loads), meta=("json", "object")),
+    UUID.__name__: lambda x: x.astype(str),
+}
+
+# Conversion map with normalized SQLAlchemy field types
+# sqlalchemy_field_conversion_map_dask: Dict[str, callable] = {
+#     "String": lambda x: x.map_partitions(lambda s: s.astype(str), meta=("string", "string")),
+#     "Text": lambda x: x.map_partitions(lambda s: s.astype(str), meta=("text", "string")),
+#     "Integer": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "SmallInteger": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "BigInteger": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "Float": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "Numeric": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "Boolean": lambda x: x.map_partitions(lambda s: s.fillna(False).astype(bool), meta=("boolean", "bool")),
+#     "DateTime": lambda x: pd.to_datetime(x, errors="coerce"),
+#     "Date": lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda s: s.dt.date, meta=("date", "object")),
+#     "Time": lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda s: s.dt.time, meta=("time", "object")),
+#     "JSON": lambda x: x.map_partitions(lambda s: s.apply(json.loads), meta=("json", "object")),
+# }
+
+
+def normalize_sqlalchemy_type(field_type):
+    """
+    Normalize SQLAlchemy field types to generic type names.
+    Handles dialect-specific types (e.g., MySQL).
+    """
+    # Map of generic SQLAlchemy types
+    type_mapping = {
+        String: "String",
+        Text: "Text",
+        Integer: "Integer",
+        SmallInteger: "SmallInteger",
+        BigInteger: "BigInteger",
+        Float: "Float",
+        Numeric: "Numeric",
+        Boolean: "Boolean",
+        DateTime: "DateTime",
+        Date: "Date",
+        Time: "Time",
+        JSON: "JSON",
+    }
+
+    # Dialect-specific types
+    dialect_mapping = {
+        TINYINT: "SmallInteger",
+        MEDIUMTEXT: "Text",
+    }
+
+    # Check if the field matches a generic or dialect-specific type
+    for sql_type, name in {**type_mapping, **dialect_mapping}.items():
+        if isinstance(field_type, sql_type):
+            return name
+
+    # Fallback to raw class name
+    return field_type.__class__.__name__
+
+

sibi_dst/df_helper/plugins/django/_django_load_from_db.py

@@ -4,6 +4,7 @@ from django.db.models import Q
 
 from sibi_dst.df_helper.plugins.django import ReadFrameDask
 from sibi_dst.utils import Logger
+from sibi_dst.df_helper.core import django_field_conversion_map_dask
 
 class DjangoLoadFromDb:
     df: dd.DataFrame
@@ -26,10 +27,10 @@ class DjangoLoadFromDb:
 
     def build_and_load(self):
         self.df = self._build_and_load()
-        if self.df is not None:
-            self._process_loaded_data()
+        #self.df = self._convert_columns(self.df)
         return self.df
 
+
     def _build_and_load(self) -> dd.DataFrame:
         query = self.connection_config.model.objects.using(self.connection_config.connection_name)
         if not self.params_config.filters:
@@ -61,10 +62,60 @@ class DjangoLoadFromDb:
                 q_objects.add(~Q(**{key: value}), Q.AND)
         return q_objects
 
-    def _process_loaded_data(self):
-        field_map = self.params_config.field_map
-        if field_map is not None:
-            rename_mapping = {k: v for k, v in field_map.items() if k in self.df.columns}
-            if rename_mapping:
-                # Apply renaming
-                self.df = self.df.rename(columns=rename_mapping)
+    def _convert_columns(self, df: dd.DataFrame) -> dd.DataFrame:
+        """
+        Convert the data types of columns in a Dask DataFrame based on the field type in the Django model.
+
+        :param df: Dask DataFrame whose columns' data types are to be converted.
+        :return: Dask DataFrame with converted column data types.
+        """
+
+        def log_debug(message: str, is_verbose: bool = False):
+            """Helper to handle debug and verbose debug logging."""
+            if self.debug:
+                self.logger.debug(message)
+                if is_verbose and self.verbose_debug:
+                    print(message)
+
+        if self.debug:
+            self.logger.info(f'Converting columns: {list(df.columns)}')
+
+        # Get field information from the Django model
+        model_fields = self.connection_config.model._meta.get_fields()
+        field_type_map = {field.name: type(field).__name__ for field in model_fields}
+        # Simplified loop to apply conversions partition-wise
+        for field_name, field_type in field_type_map.items():
+            if field_name not in df.columns:
+
+                log_debug(f"Column '{field_name}' not found in DataFrame columns.")
+                continue
+
+            conversion_func = django_field_conversion_map_dask.get(field_type)
+            if not conversion_func:
+                message=f"Field type '{field_type}' not found in conversion_map."
+                log_debug(message, is_verbose=True)
+                continue
+
+            def apply_conversion(partition):
+                """
+                Apply the conversion function to a single partition for the given column.
+                """
+                try:
+                    if field_name in partition.columns:
+                        partition[field_name] = conversion_func(partition[field_name])
+                except Exception as e:
+                    self.logger.error(f"Error converting column '{field_name}' in partition: {str(e)}")
+                return partition
+
+            try:
+                # Apply conversion lazily to each partition
+                df = df.map_partitions(
+                    apply_conversion,
+                    meta=df,
+                )
+                log_debug(f"Successfully queued conversion for column '{field_name}' to type '{field_type}'.",
+                          is_verbose=True)
+            except Exception as e:
+                log_debug(f"Failed to queue conversion for column '{field_name}': {str(e)}", is_verbose=True)
+
+        return df

sibi_dst/df_helper/plugins/sql_alchemy/_io_sqlalchemy_dask.py

@@ -1,143 +1,133 @@
 import itertools
+
 import dask.dataframe as dd
 import pandas as pd
-from sqlalchemy.orm import Query
-from sqlalchemy.inspection import inspect
-
-
-class ReadFrameSqlAlchemy:
-    def __init__(
-        self,
-        query,
-        session,
-        fieldnames=None,
-        index_col=None,
-        coerce_float=False,
-        verbose=True,
-        datetime_index=False,
-        column_names=None,
-        chunk_size=1000,
-    ):
+from sqlalchemy import create_engine, inspect, select
+from sqlalchemy.orm import sessionmaker
+
+from sibi_dst.df_helper.plugins.sql_alchemy._sqlachemy_filter_handler import SqlAlchemyFilterHandler
+from sibi_dst.utils import Logger
+
+
+class SQLAlchemyDask:
+    def __init__(self, model, filters, engine_url, chunk_size=1000, logger=None, verbose=True):
         """
-        Initialize the loader for SQLAlchemy queries.
-
-        Args:
-            query: SQLAlchemy query (ORM or Select).
-            session: SQLAlchemy session for executing the query.
-            fieldnames: Optional list of field names to include in the result.
-            index_col: Column to use as the index of the DataFrame.
-            coerce_float: Attempt to coerce values to float where applicable.
-            verbose: Whether to include verbose processing like handling choices.
-            datetime_index: Whether to convert the index to a datetime index.
-            column_names: Optional mapping of fieldnames to custom column names.
-            chunk_size: Number of records to fetch in each chunk.
+        Initialize with an SQLAlchemy query and database engine URL.
+
+        :param model: SQLAlchemy ORM model.
+        :param filters: Filters to apply on the query.
+        :param engine_url: Database connection string for SQLAlchemy engine.
+        :param chunk_size: Number of records per chunk for Dask partitions.
+        :param logger: Logger instance for logging.
+        :param verbose: Whether to print detailed logs.
         """
-        self.query = query
-        self.session = session
-        self.fieldnames = fieldnames
-        self.index_col = index_col
-        self.coerce_float = coerce_float
-        self.verbose = verbose
-        self.datetime_index = datetime_index
-        self.column_names = column_names
+        self.query = None
+        self.model = model
+        self.filters = filters
         self.chunk_size = chunk_size
+        self.verbose = verbose
+        self.engine = create_engine(engine_url)
+        self.Session = sessionmaker(bind=self.engine)
+        self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__)
 
     @staticmethod
-    def row_to_dict(row, fields=None):
+    def infer_dtypes_from_model(model):
         """
-        Convert a SQLAlchemy result row to a dictionary.
-
-        Args:
-            row: SQLAlchemy ORM object, Row object, or tuple.
-            fields: List of fields to extract.
-
-        Returns:
-            A dictionary representation of the row.
+        Infer data types for Dask DataFrame based on SQLAlchemy ORM model columns.
         """
-        # Handle ORM instances
-        if hasattr(row, "__dict__"):  # For ORM instances
-            data = row.__dict__.copy()
-            data.pop("_sa_instance_state", None)  # Remove SQLAlchemy internal state
-        # Handle SQLAlchemy Row objects
-        elif hasattr(row, "_mapping"):  # For SQLAlchemy result rows
-            data = dict(row._mapping)
-        # Handle tuples (e.g., raw query results)
-        elif isinstance(row, tuple):
-            if fields:
-                data = dict(zip(fields, row))
-            else:
-                raise ValueError("Cannot map tuple row without field names.")
-        else:
-            raise ValueError(f"Unsupported row type: {type(row)}. Expected ORM instance, dict-like object, or tuple.")
-
-        # Filter by specified fields
-        if fields:
-            return {field: data.get(field, None) for field in fields}
-        else:
-            return data
+        mapper = inspect(model)
+        sqlalchemy_to_dask_dtype = {
+            #'INTEGER': pd.to_numeric(x, errors="coerce"),
+            'INTEGER': 'Int64',
+            'SMALLINT': 'Int64',
+            'BIGINT': 'Int64',
+            'FLOAT': 'float64',
+            'NUMERIC': 'float64',
+            'BOOLEAN': 'bool',
+            'VARCHAR': 'object',
+            'TEXT': 'object',
+            'DATE': 'datetime64[ns]',
+            'DATETIME': 'datetime64[ns]',
+            'TIME': 'object',
+            'UUID': 'object',
+        }
+
+        dtypes = {}
+        for column in mapper.columns:
+            dtype = sqlalchemy_to_dask_dtype.get(str(column.type).upper(), 'object')
+            dtypes[column.name] = dtype
+
+        return dtypes
 
     def read_frame(self, fillna_value=None):
         """
-        Convert the query results to a Dask DataFrame.
-
-        Args:
-            fillna_value: Value to use for filling missing values.
+        Load data from an SQLAlchemy query into a Dask DataFrame.
 
-        Returns:
-            A Dask DataFrame.
+        :param fillna_value: Value to replace NaN or NULL values with, if any.
+        :return: Dask DataFrame.
         """
-        # Infer fieldnames if not provided
-        if not self.fieldnames:
-            if hasattr(self.query, "selected_columns"):
-                self.fieldnames = [col.key for col in self.query.selected_columns]
-            else:
-                self.fieldnames = [col.name for col in inspect(self.query._entity_zero().class_).columns]
-
-        partitions = []
-        results = self.session.execute(self.query)  # Execute the query
-
-        # Debugging raw results
-        print("Results fetched:", results)
-
-        # Chunk processing
-        iterator = iter(results)
-        while True:
-            chunk = list(itertools.islice(iterator, self.chunk_size))
-            if not chunk:
-                break
-
-            # Convert chunk to DataFrame
-            df = pd.DataFrame.from_records(
-                [self.row_to_dict(row, self.fieldnames) for row in chunk],
-                columns=self.fieldnames,
-                coerce_float=self.coerce_float,
-            )
-
-            # Handle missing values
-            if fillna_value is not None:
-                df = df.fillna(fillna_value)
-
-            # Convert datetime columns to timezone-naive
-            for col in df.columns:
-                if isinstance(df[col].dtype, pd.DatetimeTZDtype):
-                    df[col] = df[col].dt.tz_localize(None)
-
-            partitions.append(dd.from_pandas(df, npartitions=1))
-
-        # Concatenate partitions
-        dask_df = dd.concat(partitions, axis=0, ignore_index=True)
-
-        # Handle index column
-        if self.index_col and self.index_col in dask_df.columns:
-            dask_df = dask_df.set_index(self.index_col)
-
-        # Convert index to datetime if required
-        if self.datetime_index and self.index_col in dask_df.columns:
-            dask_df = dask_df.map_partitions(lambda df: df.set_index(pd.to_datetime(df.index)))
-
-        # Handle column renaming
-        if self.column_names:
-            rename_mapping = dict(zip(self.fieldnames, self.column_names))
-            dask_df = dask_df.rename(columns=rename_mapping)
-
-        return dask_df
+        with self.Session() as session:
+            try:
+                # Build query
+                self.query = select(self.model)
+                if self.filters:
+                    self.query = SqlAlchemyFilterHandler.apply_filters_sqlalchemy(self.query, self.model, self.filters)
+                else:
+                    n_records = 100
+                    self.query = self.query.limit(n_records)
+
+                # Infer dtypes
+                dtypes = self.infer_dtypes_from_model(self.model)
+                # Get the column order from the SQLAlchemy model
+                ordered_columns = [column.name for column in self.model.__table__.columns]
+
+                # Execute query and fetch results in chunks
+                result_proxy = session.execute(self.query)
+                results = result_proxy.scalars().all()  # Fetch all rows
+                iterator = iter(results)
+
+                partitions = []
+
+                while True:
+                    chunk = list(itertools.islice(iterator, self.chunk_size))
+                    if not chunk:
+                        break
+
+                    # Convert chunk to Pandas DataFrame
+                    df = pd.DataFrame.from_records(
+                        [row._asdict() if hasattr(row, '_asdict') else row.__dict__ for row in chunk]
+                    )
+                    # Drop internal SQLAlchemy state if it exists
+                    df = df.loc[:, ~df.columns.str.contains('_sa_instance_state')]
+
+                    # Reorder columns to match the model's order
+                    df = df[ordered_columns]
+
+                    # Fill NaN values
+                    if fillna_value is not None:
+                        df = df.fillna(fillna_value)
+
+                    # Convert timezone-aware columns to naive
+                    for col in df.columns:
+                        if isinstance(df[col].dtype, pd.DatetimeTZDtype):
+                            df[col] = df[col].dt.tz_localize(None)
+
+                    # Apply inferred dtypes
+                    df = df.astype(dtypes)
+                    # Create a Dask partition
+                    partitions.append(dd.from_pandas(df, npartitions=1))
+
+                # Concatenate all partitions
+                # print(partitions)
+                if partitions:
+                    dask_df = dd.concat(partitions, axis=0, ignore_index=True)
+                else:
+                    dask_df = dd.from_pandas(pd.DataFrame(columns=ordered_columns), npartitions=1)
+
+                if self.verbose:
+                    self.logger.info(f"Loaded {len(dask_df)} rows into Dask DataFrame.")
+                return dask_df
+
+            except Exception as e:
+                self.logger.error(f"Error executing query: {str(e)}")
+                return dd.from_pandas(pd.DataFrame(columns=ordered_columns), npartitions=1)

sibi_dst/df_helper/plugins/sql_alchemy/_sqlachemy_filter_handler.py

@@ -50,6 +50,7 @@ class SqlAlchemyFilterHandler:
                     return [datetime.date.fromisoformat(v) for v in value]
             return value
 
+
         def handle_date_operator(column, date_op):
             """
             Handle filtering on specific datetime parts (e.g., year, month).
@@ -90,6 +91,7 @@ class SqlAlchemyFilterHandler:
 
             # Get the column from the model
             column = getattr(model, field_name, None)
+            #column = model.__table__.columns.get(field_name)
             if not column:
                 raise AttributeError(f"Field '{field_name}' not found in model '{model.__name__}'")
 

sibi_dst/df_helper/plugins/sql_alchemy/_sqlalchemy_load_from_db.py

@@ -1,23 +1,30 @@
+from typing import Dict
+
 import dask.dataframe as dd
-from sqlmodel import Session, select
-from typing import Any, Dict, Optional
-import logging
 import pandas as pd
+from sqlalchemy.inspection import inspect
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy import select
+#from sqlmodel import Session, select
+
+from sibi_dst.df_helper.core import ParamsConfig, QueryConfig, sqlalchemy_field_conversion_map_dask, \
+    normalize_sqlalchemy_type
 from sibi_dst.utils import Logger
-from ._sqlalchemy_db_connection import SqlAlchemyConnectionConfig
-from sibi_dst.df_helper.core import ParamsConfig, QueryConfig
+from ._io_sqlalchemy_dask import SQLAlchemyDask
 from ._sqlachemy_filter_handler import SqlAlchemyFilterHandler
+from ._sqlalchemy_db_connection import SqlAlchemyConnectionConfig
+
 
 class SqlAlchemyLoadFromDb:
     df: dd.DataFrame
 
     def __init__(
-        self,
-        plugin_sqlalchemy: SqlAlchemyConnectionConfig,  # Expected to be an instance of SqlAlchemyConnection
-        plugin_query: QueryConfig = None,
-        plugin_params: ParamsConfig = None,
-        logger: Logger = None,
-        **kwargs,
+            self,
+            plugin_sqlalchemy: SqlAlchemyConnectionConfig,  # Expected to be an instance of SqlAlchemyConnection
+            plugin_query: QueryConfig = None,
+            plugin_params: ParamsConfig = None,
+            logger: Logger = None,
+            **kwargs,
     ):
         """
         Initialize the loader with database connection, query, and parameters.
@@ -27,7 +34,7 @@ class SqlAlchemyLoadFromDb:
         self.model = self.db_connection.model
         self.engine = self.db_connection.engine
         self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__)
-        self.query_config =  plugin_query
+        self.query_config = plugin_query
         self.params_config = plugin_params
         self.debug = kwargs.pop("debug", False)
         self.verbose_debug = kwargs.pop("verbose_debug", False)
@@ -40,44 +47,13 @@ class SqlAlchemyLoadFromDb:
         return self.df
 
     def _build_and_load(self) -> dd.DataFrame:
-        """
-        Query the database and load results into a Dask DataFrame.
-        """
-        with Session(self.engine) as session:
-            try:
-                query = select(self.model)
-                filters = self.params_config.filters
-                if filters:
-                    n_records = 0
-                    query = SqlAlchemyFilterHandler.apply_filters_sqlalchemy(query, self.model,self.params_config.filters)
-                else:
-                    n_records = self.query_config.n_records or 100
-
-                if n_records:
-                    query = query.limit(n_records)
-
-                # Debug: Log the SQL query
-                if self.debug:
-                    self.logger.info(f"Executing query: {str(query)}")
-
-                # Execute the query
-                try:
-                    results = session.exec(query).fetchall()
-                    if results:
-                        records = [
-                            {key: getattr(result, key) for key in result.__table__.columns.keys()}
-                                for result in results
-                        ]
-                    df = dd.from_pandas(pd.DataFrame(records), npartitions=1)
-                except Exception as e:
-                    self.logger.info(results)
-                    self.logger.warning("Query returned no results.")
-                    df = dd.from_pandas(pd.DataFrame(), npartitions=1)
-
-            except Exception as e:
-                print(query)
-                self.logger.error(f"Error loading data: {e}")
-                df = dd.from_pandas(pd.DataFrame(), npartitions=1)
-
+        try:
+            reader = SQLAlchemyDask(model=self.model, filters=self.params_config.filters,engine_url=self.engine.url, logger=self.logger, chunk_size=1000, verbose=self.debug)
+            df = reader.read_frame()
+            if df is None or len(df.index) == 0:
+                self.logger.warning("Query returned no results.")
+                return dd.from_pandas(pd.DataFrame(), npartitions=1)
             return df
-
+        except Exception as e:
+            self.logger.error(f"Failed to load data into Dask DataFrame.{e}")
+            return dd.from_pandas(pd.DataFrame(), npartitions=1)

sibi_dst/utils/_data_utils.py

@@ -8,7 +8,7 @@ class DataUtils:
         self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__)
 
     @staticmethod
-    def transform_numeric_columns(df, fill_value=0, transform_func=None):
+    def transform_numeric_columns(df, columns=None, fill_value=0, transform_func=None):
         """
         Transform numeric columns in a DataFrame (Pandas or Dask), handling missing values and applying optional transformations.
 
@@ -21,10 +21,11 @@ class DataUtils:
         Returns:
         - pandas.DataFrame or dask.dataframe.DataFrame: Updated DataFrame with transformed numeric columns.
         """
-        # Detect numeric columns
-        numeric_columns = df.select_dtypes(include=['number']).columns.tolist()
+        if columns is None:
+            # Detect numeric columns
+            columns = df.select_dtypes(include=['number']).columns.tolist()
 
-        if not numeric_columns:
+        if not columns:
             return df
 
         # Default transformation function (identity) if none is provided
@@ -32,7 +33,7 @@ class DataUtils:
             transform_func = lambda x: x
 
         # Apply transformations
-        for col in numeric_columns:
+        for col in columns:
             dtype = df[col].dtype
             if pd.api.types.is_integer_dtype(dtype):
                 meta_type = 'int64'
@@ -51,38 +52,36 @@ class DataUtils:
         return df
 
     @staticmethod
-    def transform_boolean_columns(df, threshold=1):
+    def transform_boolean_columns(df, columns=None, sample_size=100):
         """
-        Transform boolean-like columns in a DataFrame (Pandas or Dask) to actual booleans.
+        Detect if the provided columns in a DataFrame (Pandas or Dask) contain only 0 and 1
+        and convert them to boolean. Detection is performed using a sample.
 
         Parameters:
         - df (pandas.DataFrame or dask.dataframe.DataFrame): The DataFrame.
-        - threshold (int or float): The value to evaluate as `True`.
+        - columns (list of str): List of columns to check and transform.
+        - sample_size (int): Number of rows to sample for detection. Ignored for Pandas DataFrames.
 
         Returns:
         - pandas.DataFrame or dask.dataframe.DataFrame: Updated DataFrame with transformed boolean columns.
         """
+        # Apply transformation to each specified column
+        for col in columns:
+            if col in df.columns:
+                if isinstance(df, dd.DataFrame):
+                    # Replace NaN with 0, then convert to boolean
+                    df[col] = df[col].map_partitions(
+                        lambda s: pd.to_numeric(s, errors='coerce')  # Convert to numeric, invalid to NaN
+                        .fillna(0)  # Replace NaN with 0
+                        .astype(int)  # Ensure integer type
+                        .astype(bool),  # Convert to boolean
+                        meta=(col, 'bool')
+                    )
+                else:
+                    # For Pandas DataFrame, handle mixed types and invalid values
+                    df[col] = pd.to_numeric(df[col], errors='coerce')  # Convert to numeric, invalid to NaN
+                    df[col] = df[col].fillna(0).astype(int).astype(bool)
 
-        def is_boolean_like(col):
-            """
-            Check if a column is boolean-like (contains only two unique values).
-            """
-            unique_values = col.dropna().unique()
-            if isinstance(col, dd.Series):
-                unique_values = unique_values.compute()
-            return len(unique_values) <= 2 and set(unique_values).issubset({0, 1, True, False})
-
-        # Detect boolean-like columns
-        boolean_columns = [col for col in df.columns if is_boolean_like(df[col])]
-
-        # Apply transformation to each detected column
-        for col in boolean_columns:
-            if isinstance(df, dd.DataFrame):
-                df[col] = df[col].map_partitions(
-                    lambda s: s == threshold, meta=(col, 'bool')
-                )
-            else:
-                df[col] = df[col] == threshold
         return df
 
     def merge_lookup_data(self, classname, df, **kwargs):

{sibi_dst-0.3.10.dist-info → sibi_dst-0.3.11.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: sibi-dst
-Version: 0.3.10
-Summary: 
+Version: 0.3.11
+Summary: Data Science Toolkit
 Author: Luis Valverde
 Author-email: lvalverdeb@gmail.com
 Requires-Python: >=3.11,<4.0

{sibi_dst-0.3.10.dist-info → sibi_dst-0.3.11.dist-info}/RECORD

@@ -1,15 +1,15 @@
 sibi_dst/__init__.py,sha256=1KaC0LYTHxjpENq-NXI325WcEYZ8GCBrHGkLoFxEcu0,251
 sibi_dst/df_helper/__init__.py,sha256=JXJBY47G6wOYhzNI646OBl3pSGWIy4282-3qPGYHU7w,167
-sibi_dst/df_helper/_df_helper.py,sha256=Pvu1kByZhUCAY9LGKFrcyasTq1MeeIBeMoeCgScStPM,12507
+sibi_dst/df_helper/_df_helper.py,sha256=e-ptCEDYt5dx8byNiA0ca8Eejl1DG1V5pioZUzabEnY,12747
 sibi_dst/df_helper/_parquet_artifact.py,sha256=f5oHwXtsNW6-ONSFsRB0AniVefA0THzP92J-nugp9vo,4973
-sibi_dst/df_helper/core/__init__.py,sha256=UXGUGGSjjrcJRrs25zPV-xgJoyYy1WjVQAExcJDWgV0,254
-sibi_dst/df_helper/core/_defaults.py,sha256=AVNT_Vk8K7dLKOnPX_-Cygi-Nuku65CIn0baE0Wn6dI,1877
+sibi_dst/df_helper/core/__init__.py,sha256=NSYY_evzq6XEkO06Nz6xLH5KznzRGI44cLbrnN3zHXQ,503
+sibi_dst/df_helper/core/_defaults.py,sha256=pJU-lX7w4nrt0Anx35j08mVr_0oMGn1bTA_iCl_p1qI,6700
 sibi_dst/df_helper/core/_params_config.py,sha256=hO-PddoaGjFebqJFgtn76WwVHcCjzPW3z5i3NyK6mDw,3475
 sibi_dst/df_helper/core/_query_config.py,sha256=HEiyR_fBJjIMum-PSQroY3KaefQ2SpW1w1SQS8oT-NU,489
 sibi_dst/df_helper/plugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sibi_dst/df_helper/plugins/django/__init__.py,sha256=VkvYql-PUyCKOfoKx5aGdUAki8X-kULfhnCvgSQzHDQ,311
 sibi_dst/df_helper/plugins/django/_django_db_connection.py,sha256=9fGvXXgqPk_TC7BvaSljxTgNDWoXi_LZxKabEgExznM,1654
-sibi_dst/df_helper/plugins/django/_django_load_from_db.py,sha256=ZWVwJCJY7xmlZwDCZG3vNlEMyTGKJ8CoEtwgKYX0ofQ,2918
+sibi_dst/df_helper/plugins/django/_django_load_from_db.py,sha256=NrkBb5LM1A_vo3wAotqj2sVVYIuTfFGrQqIXk3xOoDs,5162
 sibi_dst/df_helper/plugins/django/_django_sql_model_builder.py,sha256=GprCh2c6PFlRBUCir4hh8pmf4Eqb-4OIK6Vz_xXQjMw,14864
 sibi_dst/df_helper/plugins/django/_io_dask.py,sha256=jryDojeA62rB3seRaWWMjsAmekKacK5xctwCQGVklPQ,9063
 sibi_dst/df_helper/plugins/django/_io_dask_alt.py,sha256=zDjLyYxBeL0ffn3yfE_7vqMLMpeEEk2o-zMr66sKkDw,6827
@@ -19,10 +19,10 @@ sibi_dst/df_helper/plugins/parquet/__init__.py,sha256=ClkyIsIh_ovEwqm0dTrkXImbPj
 sibi_dst/df_helper/plugins/parquet/_parquet_filter_handler.py,sha256=45mHID1azAg5PmaYWbuRlghoRd3H2aTLj1XcycfLJo0,3497
 sibi_dst/df_helper/plugins/parquet/_parquet_options.py,sha256=cKyRj0UCby9-iYPPFnlel1H03x8MnAoEv8k1tp7kHXw,4277
 sibi_dst/df_helper/plugins/sql_alchemy/__init__.py,sha256=FHorj40SbHc0OBzQ_ieG6MG-HLbf0tw6I_5eoIjJkOI,369
-sibi_dst/df_helper/plugins/sql_alchemy/_io_sqlalchemy_dask.py,sha256=ET9cih0Frc5izMOsdvNlLhjJWtUQbwZhRtsdo5dRckQ,5059
-sibi_dst/df_helper/plugins/sql_alchemy/_sqlachemy_filter_handler.py,sha256=UXcZ1shS7shcjkSqIIduAnb1Lhzc6pZ6NEcbkcnwgWk,4606
+sibi_dst/df_helper/plugins/sql_alchemy/_io_sqlalchemy_dask.py,sha256=6IjQEREXqTAzSJE95FKfXjRkTlEjRMS4hJ_yMpyKDTg,5223
+sibi_dst/df_helper/plugins/sql_alchemy/_sqlachemy_filter_handler.py,sha256=H8ypUjLKzYYl9BerfJjX_Uv9qBVkBR-wZiQlh3uRQXg,4669
 sibi_dst/df_helper/plugins/sql_alchemy/_sqlalchemy_db_connection.py,sha256=HtMsfH5com4dLVJxh3wdMUpQI3mz0cKDJz0CmFS2S8U,1648
-sibi_dst/df_helper/plugins/sql_alchemy/_sqlalchemy_load_from_db.py,sha256=f1zqBISVn8OjZJs8hu6IvRZSwMX7_DIZMIbhxV6uV80,3179
+sibi_dst/df_helper/plugins/sql_alchemy/_sqlalchemy_load_from_db.py,sha256=myrtEzK12DvA73x7QFaqXFb_TxOPMrsVj-mxYHJD2dg,2371
 sibi_dst/df_helper/plugins/sql_alchemy/_sqlalchemy_model_builder.py,sha256=vrTTeFcrf8sFpzqLlQfspjduKuR7Wgn7mDPUASsQs6s,4355
 sibi_dst/df_helper/plugins/sql_model/__init__.py,sha256=MXd4OOdTqR4cENSV733SGodPO6eQMCexANs-3w0qL5U,226
 sibi_dst/df_helper/plugins/sql_model/_sqlmodel_db_connection.py,sha256=6jmMjKIv5Btysj3kZMaXQ98IqKQkhnOC-JWtb1B8rus,4265
@@ -31,7 +31,7 @@ sibi_dst/utils/__init__.py,sha256=jiXJSnmsaGZTRhUThtIo6cssWXBWXNij8ffYmv77QK4,79
 sibi_dst/utils/_airflow_manager.py,sha256=rlt3eolR5QvtxWhAtBTCpHXvxftnKM-ibPMv3fVwNZk,7524
 sibi_dst/utils/_clickhouse_writer.py,sha256=mdgszbyVluhGvDmvsHY4XDTZrp42L3xtdmiyn3z2bYM,8534
 sibi_dst/utils/_credentials.py,sha256=8i6z7y3y5S-6mSk4xrT2AwhzCA32mTn1n1iYX9IVyHk,1724
-sibi_dst/utils/_data_utils.py,sha256=XB0xjW2LbCmoZjgDbNQQpWaf4upmSoTXeJZ3QMVqbsQ,7056
+sibi_dst/utils/_data_utils.py,sha256=3hBMg852ANpS5bOtlU-F4H-Q91WIGga5LrKWWyDvnAA,7354
 sibi_dst/utils/_data_wrapper.py,sha256=pZnylBFTvsLGfYGv2tTyQHzyb6IbIahfaXR-PxHdivk,24099
 sibi_dst/utils/_date_utils.py,sha256=6HCrcTiuYLNsbgrNB3eAVAAgXbfx7Ce1qNc3OJla9nM,5621
 sibi_dst/utils/_df_utils.py,sha256=o2bK5-xMGKqIG4i9xfavYRxIkiHLA0nz5TQTN78998k,7350
@@ -40,6 +40,6 @@ sibi_dst/utils/_filepath_generator.py,sha256=ytPSZ9GYOnnSP25zwA-0NjFHupPRZyXwixW
 sibi_dst/utils/_log_utils.py,sha256=AAenyubYUjk77WqiaNkjgkxws3dnAMIdaGl2Ryz_cA4,2245
 sibi_dst/utils/_parquet_saver.py,sha256=-A0o_vucyYe7wlwiby_0_yS-ZfT2GHwImyQHrCIBNwk,9051
 sibi_dst/utils/_storage_manager.py,sha256=KP2HBXnLUMMquqcO30ecfuoU7g1z8RtaV3Dv0TvEXoY,3856
-sibi_dst-0.3.10.dist-info/METADATA,sha256=lrVYU1PPBuHQrEDl_-SURTkE0ip_0xWsJc58AiihHZs,1877
-sibi_dst-0.3.10.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-sibi_dst-0.3.10.dist-info/RECORD,,
+sibi_dst-0.3.11.dist-info/METADATA,sha256=gwl565etE5wLVGk0rqQ7umOyBRtEXpQ_IdCXyEkv2s8,1897
+sibi_dst-0.3.11.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+sibi_dst-0.3.11.dist-info/RECORD,,