apb_pandas_utils

  1#  coding=utf-8
  2#
  3#  Author: Ernesto Arredondo Martinez (ernestone@gmail.com)
  4#  Created: 
  5#  Copyright (c)
  6"""
  7.. include:: ../README.md
  8"""
  9from __future__ import annotations
 10
 11import re
 12from collections.abc import Iterable
 13from datetime import datetime, time, date
 14from typing import Union, Generator
 15
 16import numpy as np
 17import pandas as pd
 18import requests
 19import sqlalchemy
 20from geopandas import GeoDataFrame
 21from pandas import DataFrame, Timestamp, NaT, CategoricalDtype
 22from requests.adapters import HTTPAdapter
 23from urllib3.util.retry import Retry
 24
 25from apb_extra_utils.postgres_pckg.psql_alchemy import EngPsqlAlchemy
 26from apb_extra_utils.utils_logging import get_base_logger
 27
 28logger = get_base_logger(__name__)
 29
 30EXCLUDED_TYPES_TO_CATEGORIZE = ['datetime', 'category', 'geometry']
 31DEFAULT_MAX_UNIQUE_VALS_COL_CATEGORY = 0.5
 32MAX_DATETIME = datetime(2250, 12, 31, 23, 59, 59)
 33
 34
 35def optimize_df(df: DataFrame | GeoDataFrame, max_perc_unique_vals: float = DEFAULT_MAX_UNIQUE_VALS_COL_CATEGORY,
 36                floats_as_categ: bool = False) -> DataFrame | GeoDataFrame:
 37    """
 38    Retorna el pd.Dataframe optimizado segun columnas que encuentre
 39
 40    Args:
 41        df (Dataframe | GeoDataFrame): Dataframe a optimizar
 42        max_perc_unique_vals (float=DEFAULT_MAX_UNIQUE_VALS_COL_CATEGORY): Màxim percentatge de valors únics respecte total files per a convertir en categoria, expressat entre 0 i 1 (Default 0.5 -> 50%)
 43        floats_as_categ (bool=False): Si True, els floats es converteixen a categoria
 44
 45    Returns:
 46        opt_df (Dataframe | GeoDataFrame): Dataframe optimizado
 47    """
 48    opt_df = df.copy()
 49    df_ints = opt_df.select_dtypes(include=['int64'])
 50    opt_df[df_ints.columns] = df_ints.apply(pd.to_numeric, downcast='signed')
 51    df_floats = opt_df.select_dtypes(include='float')
 52    opt_df[df_floats.columns] = df_floats.apply(pd.to_numeric, downcast='float')
 53
 54    excl_types_cat = EXCLUDED_TYPES_TO_CATEGORIZE
 55    if not floats_as_categ:
 56        excl_types_cat.append('float')
 57
 58    for col in opt_df.select_dtypes(exclude=excl_types_cat).columns:
 59        try:
 60            unic_vals = opt_df[col].unique()
 61        except (pd.errors.DataError, TypeError):
 62            continue
 63
 64        num_unique_values = len(unic_vals)
 65        num_total_values = len(opt_df[col]) - len(opt_df.loc[opt_df[col].isnull()])
 66        if num_total_values > 0 and (num_unique_values / num_total_values) < max_perc_unique_vals:
 67            try:
 68                opt_df[col] = opt_df[col].astype(CategoricalDtype(ordered=True))
 69            except (NotImplementedError, TypeError):
 70                continue
 71
 72    return opt_df
 73
 74
 75def df_filtered_by_prop(df: DataFrame | GeoDataFrame, filter_prop: dict[str, object]) -> DataFrame | GeoDataFrame | None:
 76    """
 77    Filtra el dataframe amb el diccionari passat, on la clau fa referència a la columna i el valor o llistat de valors
 78    separats per comes son els que s’apliquen al filtre. Si la clau/columna no existeix es desestima. Si la clau/columna
 79    comença per alguns d’aquest signes “=, !, -, >, <” s’aplica la corresponent operació de filtre.
 80    En el cas de “!” i “–“ s’aplica la mateixa operació de negat o que no contingui el valor o valors passats.
 81    Els filtres “<“ i “>” no apliquen a camps text i es desestimen. Es poden passar la mateixa columna amb operadors
 82    i valors distints per aplicar filtres diferents
 83
 84    Args:
 85        df (DataFrame | GeoDataFrame): DataFrame a filtrar
 86        filter_prop (dict[str, object]): Propietats de filtrat
 87
 88    Returns:
 89        DataFrame | GeoDataFrame: DataFrame filtrat
 90    """
 91    if df is None or not filter_prop:
 92        return df
 93
 94    idx_names = [idx_col for idx_col in df.index.names if idx_col]
 95    if idx_names:
 96        df = df.reset_index()
 97
 98    def _df_individual_filter(_df_ind: DataFrame, type_col_ind, column: str, value, col_operator: str = '='):
 99        type_column = type_col_ind.categories.dtype if (type_col_name := type_col_ind.name) == 'category' else type_col_ind
100
101        if type_col_name == 'object':
102            if col_operator == '=':
103                _df_ind = _df_ind[_df_ind[column].str.contains(str(value), case=False, na=False)]
104            elif col_operator == '-' or col_operator == '!':
105                _df_ind = _df_ind[~_df_ind[column].str.contains(str(value), case=False, na=False)]
106        else:
107            value = type_column.type(value)
108            if col_operator == '=':
109                _df_ind = _df_ind.loc[_df_ind[column] == value]
110            elif col_operator == '-' or col_operator == '!':
111                _df_ind = _df_ind.loc[_df_ind[column] != value]
112            elif col_operator == '>':
113                _df_ind = _df_ind.loc[_df_ind[column] > value]
114            elif col_operator == '<':
115                _df_ind = _df_ind.loc[_df_ind[column] < value]
116
117        return _df_ind
118
119    col_names = df.columns.values.tolist()
120    for k, v in filter_prop.items():
121        k_operator = "="
122        if k.startswith(('-', '=', '<', '>', '!')):
123            k_operator = k[0:1]
124            k = k[1:]
125        if k.upper() in (col_names + idx_names):
126            k = k.upper()
127        elif k.lower() in (col_names + idx_names):
128            k = k.lower()
129
130        if k in col_names and v is not None:
131            type_col = df.dtypes.get(k)
132            if isinstance(v, list):
133                # es fa amb un bucle i no amb isin perque no val per floats
134                df_list = None
135                for ind_val in v:
136                    df_temp = _df_individual_filter(df, type_col, k, ind_val, k_operator)
137                    if df_list is None:
138                        df_list = df_temp
139                    elif k_operator == '=':
140                        df_list = pd.concat([df_list, df_temp])
141                    if k_operator != '=':
142                        # per als operadors que exclouen s'ha de treballar sobre el df filtrat resultant
143                        df = df_list = df_temp
144                df = df_list
145            else:
146                df = _df_individual_filter(df, type_col, k, v, k_operator)
147
148    if idx_names:
149        df.set_index(idx_names, inplace=True)
150
151    return df
152
153
154def rename_and_drop_columns(df: Union[DataFrame, GeoDataFrame], map_old_new_col_names: dict[str, str],
155                            drop_col: bool = True, strict: bool = False, reordered: bool = False) -> Union[
156    DataFrame, GeoDataFrame]:
157    """
158    Function to rename and remove columns from a dataframe. If the drop_col parameter is True,
159    the columns that are not in the map will be removed. If the strict parameter is True,
160    the names that do not exist in the map as a column will be skipped.
161    Args:
162        df: to remove and rename
163        map_old_new_col_names: the key is the actual name and the value is the new name
164        drop_col: True to remove columns that are not included in the map
165        strict: False to skip names that are not included in the map
166        reordered: True to reorder columns of dataframe
167
168    Returns: modified DataFrame
169
170    """
171    if df is not None and map_old_new_col_names:
172        col_names = df.columns.values.tolist()
173        col_names_to_drop = col_names.copy()
174        final_map = {}
175        for k, v in map_old_new_col_names.items():
176            if k in col_names:
177                final_map[k] = v
178                col_names_to_drop.remove(k)
179        if drop_col:
180            df = df.drop(col_names_to_drop, axis=1)
181        if strict:
182            final_map = map_old_new_col_names
183        df = df.rename(columns=final_map)
184        if reordered:
185            new_cols = list(map_old_new_col_names.values())
186            act_cols = df.columns.tolist()
187            reord_cols = [value for value in new_cols if value in act_cols]
188            df = df[reord_cols]
189        return df
190
191
192def set_null_and_default_values(df: DataFrame | GeoDataFrame) -> DataFrame | GeoDataFrame:
193    """
194    Function to replace NaN values with None in a DataFrame
195    Args:
196        df (DataFrame | GeoDataFrame): DataFrame to replace NaN values with None
197
198    Returns:
199        DataFrame | GeoDataFrame: DataFrame with NaN values replaced with None
200    """
201    df = df.replace({np.nan: None})
202    return df
203
204
205def replace_values_with_null(df: Union[DataFrame | GeoDataFrame], dict_col_values: dict) -> Union[
206    DataFrame | GeoDataFrame]:
207    """
208    Function to replace values with None in a DataFrame
209    Args:
210        df (DataFrame | GeoDataFrame): DataFrame to replace values with None
211        dict_col_values (dict): Dictionary with the column name and the value to replace with None
212
213    Returns:
214        DataFrame | GeoDataFrame: DataFrame with values replaced with None
215    """
216    if df is not None and not df.empty and dict_col_values:
217        for name_col, value in dict_col_values.items():
218            df[name_col] = df[name_col].replace(value, None)
219    return df
220
221
222def convert_to_datetime_col_df(df: DataFrame, cols: list[str],
223                               set_end_day: bool = False, set_nat: bool = False) -> DataFrame:
224    """
225    Force convert date columns to datetime format.
226    If init_date is True, the time is set to 00:00:00 if not to 23:59:59
227
228    Args:
229        df (DataFrame): DataFrame
230        cols (list[str]): Columns to convert
231        set_end_day (bool=False): If False the time is set to 00:00:00 if not to 23:59:59
232        set_nat (bool=False): If True set NaT to MAX_DATETIME
233
234    Returns:
235        DataFrame: DataFrame with datetime columns
236    """
237    if not set_end_day:
238        delta_time = time.min
239    else:
240        delta_time = time(23, 59, 59)
241
242    def _convert_date(value):
243        if type(value) is date:
244            value = datetime.combine(value, time.min)
245
246        if set_nat and (value is NaT or value is None):
247            return MAX_DATETIME
248        elif value is NaT:
249            return value
250        elif ((isinstance(value, Timestamp) or isinstance(value, datetime))
251              and set_end_day and value.time() == time.min):
252            return datetime.combine(value, delta_time)
253        else:
254            return value
255
256    for col in cols:
257        df[col] = df[col].apply(_convert_date)
258
259    return df
260
261
262def df_memory_usage(df: DataFrame | GeoDataFrame) -> float:
263    """
264    Return the memory usage of a DataFrame in MB
265
266    Args:
267        df (DataFrame | GeoDataFrame): DataFrame
268
269    Returns:
270        float: Memory usage in MB
271    """
272    return df.memory_usage(deep=True).sum() / 1024 ** 2
273
274
275def df_from_pg_sql(sql: str, user: str | None = None, psw: str | None = None, srvr_db: str = 'localhost',
276                   port_db: int = 5432, db: str = 'postgres', schemas: str | None = None,
277                   a_logger=None, url_conn_string: str | sqlalchemy.engine.url.URL | None = None) -> DataFrame:
278    """
279    Ejecuta un SQL sobre PostgreSQL y devuelve el resultado como DataFrame.
280
281    Acepta los mismos parametros de conexion que el constructor de ``EngPsqlAlchemy``.
282
283    Args:
284        sql (str): SQL a ejecutar.
285        user (str | None): Usuario PostgreSQL.
286        psw (str | None): Password PostgreSQL.
287        srvr_db (str): Host del servidor.
288        port_db (int): Puerto del servidor.
289        db (str): Nombre de la base de datos.
290        schemas (str | None): Schemas separados por coma para ``search_path``.
291        a_logger (logging.Logger | None): Logger opcional.
292        url_conn_string (str | sqlalchemy.engine.urlURL | None): Connection string completa o URL de SQLAlchemy (ver apb_extra_utils.postgres_pckg.psql_alchemy.url_pg_string_connection). Si se proporciona, se ignoran los parámetros anteriores.
293
294    Returns:
295        DataFrame: Resultado del SQL.
296    """
297    eng = EngPsqlAlchemy.get_cached(
298        user=user,
299        psw=psw,
300        srvr_db=srvr_db,
301        port_db=port_db,
302        db=db,
303        schemas=schemas,
304        a_logger=a_logger,
305        url_conn_string=url_conn_string,
306    )
307
308    return pd.read_sql_query(sql, eng.eng_db)
309
310
311def df_from_pg_table(table: str, filter_sql: str | None = None, user: str | None = None,
312                     psw: str | None = None, srvr_db: str = 'localhost',
313                     port_db: int = 5432, db: str = 'postgres', schemas: str | None = None, a_logger=None,
314                     url_conn_string: str | sqlalchemy.engine.url.URL | None = None) -> DataFrame:
315    """
316    Carga una tabla o vista de PostgreSQL en un DataFrame.
317
318    Args:
319        table (str): Nombre de tabla o vista.
320        filter_sql (str | None): Condición SQL opcional para cláusula WHERE. ATENCIÓN: en PG si hay campos en MAYÚSCULAS hay que indicarlos entre ""
321        user (str | None): Usuario PostgreSQL.
322        psw (str | None): Password PostgreSQL.
323        srvr_db (str): Host del servidor.
324        port_db (int): Puerto del servidor.
325        db (str): Nombre de la base de datos.
326        schemas (str | None): Schemas separados por coma para ``search_path``.
327        a_logger (logging.Logger | None): Logger opcional.
328        url_conn_string (str | sqlalchemy.engine.url.URL | None): Connection string completa o URL de SQLAlchemy (ver apb_extra_utils.postgres_pckg.psql_alchemy.url_pg_string_connection). Si se proporciona, se ignoran los parámetros anteriores.
329
330    Returns:
331        DataFrame: Resultado de ``SELECT *`` sobre la tabla/vista.
332    """
333
334    def _quote_pg_identifier(identifier: str) -> str:
335        # Quote identifiers to support reserved words/mixed case and avoid SQL injection patterns.
336        if not isinstance(identifier, str) or not identifier.strip():
337            raise ValueError("El nombre de tabla/vista debe ser un string no vacio")
338        clean = identifier.strip().replace('"', '""')
339        return f'"{clean}"'
340
341    table_sql = _quote_pg_identifier(table)
342
343    sql = f"SELECT * FROM {table_sql}"
344    if filter_sql and filter_sql.strip():
345        where_sql = filter_sql.strip()
346        if where_sql.lower().startswith('where '):
347            where_sql = where_sql[6:].strip()
348        if where_sql:
349            sql += f" WHERE {where_sql}"
350
351    return df_from_pg_sql(
352        sql=sql,
353        user=user,
354        psw=psw,
355        srvr_db=srvr_db,
356        port_db=port_db,
357        db=db,
358        schemas=schemas,
359        a_logger=a_logger,
360        url_conn_string=url_conn_string,
361    )
362
363
364def extract_operator(input_string):
365    """
366    Extract sql operator from input string
367
368    Args:
369        input_string:
370
371    Returns:
372
373    """
374    special_opers = ('-', '!') # Special operators for negation
375    sql_opers = ('=', '!=', '<>', '>', '<', '>=', '<=') # SQL operators
376
377    match = re.match(r'^(\W+)', input_string)
378    if match:
379        symbols = match.group(1)
380
381        if symbols not in special_opers and symbols not in sql_opers:
382            raise ValueError(f"Operator '{symbols}' not supported")
383
384        return symbols
385
386    return None
387
388
389def sql_from_filter_by_props(**filter_by_props: dict) -> str:
390    """
391    Get SQL from filter by properties
392
393    Args:
394        **filter_by_props: The filter by properties
395
396    Returns:
397        sql (str): The SQL from filter by properties
398    """
399    sql_parts = []
400    for k_fld_oper, value in filter_by_props.items():
401        if k_operator := extract_operator(k_fld_oper):
402            k_fld = k_fld_oper.replace(k_operator, '')
403        else:
404            k_operator = "="
405            k_fld = k_fld_oper
406
407        if isinstance(value, str):
408            value = f"'{value}'"
409        elif isinstance(value, Iterable):
410            value = ', '.join([f"'{v}'" if isinstance(v, str) else str(v) for v in value])
411            value = f"({value})"
412            if k_operator in ('=', '!=', '-'):
413                k_operator = "IN" if k_operator == "=" else "NOT IN"
414            else:
415                raise ValueError(f"Operator '{k_operator}' not supported for iterable values")
416
417        sql_parts.append(f"{k_fld} {k_operator} {value}")
418
419    sql_filter = ' AND '.join(sql_parts)
420
421    return sql_filter
422
423
424def _build_session(max_retries: int = 3) -> requests.Session:
425    """
426    Create a :class:`requests.Session` with a retry strategy and exponential backoff.
427
428    Args:
429        max_retries (int): Number of retries on transient HTTP errors
430            (429, 500, 502, 503, 504). Defaults to ``3``.
431
432    Returns:
433        requests.Session: Configured session with retry logic.
434    """
435    session = requests.Session()
436    retry_strategy = Retry(
437        total=max_retries,
438        backoff_factor=1,
439        status_forcelist=[429, 500, 502, 503, 504],
440        allowed_methods=["GET"],
441        raise_on_status=False,
442    )
443    adapter = HTTPAdapter(max_retries=retry_strategy)
444    session.mount("http://", adapter)
445    session.mount("https://", adapter)
446    return session
447
448
449def _iter_fetch_pages(
450    url_rest_api: str,
451    api_params: dict | None = None,
452    headers: dict | None = None,
453    next_key: str = 'next',
454    timeout: int | tuple[int, int] = (10, 30),
455    max_retries: int = 3,
456    session: requests.Session | None = None,
457) -> Generator[dict | list]:
458    """
459    Internal helper: fetch all pages from a paginated REST API.
460
461    Handles session lifecycle, retry logic and pagination automatically.
462    Each element in the returned list is the raw JSON body of one page.
463
464    Args:
465        url_rest_api (str): The base URL of the API endpoint.
466        api_params (dict, optional): Query parameters for the first request only.
467        headers (dict, optional): HTTP headers added to the session.
468        next_key (str): Key in JSON dict responses that carries the next-page URL.
469            Defaults to ``'next'``.
470        timeout (int | tuple[int, int]): Request timeout ``(connect, read)`` in seconds.
471            Defaults to ``(10, 30)``.
472        max_retries (int): Retries on transient errors. Defaults to ``3``.
473        session (requests.Session, optional): Existing session to reuse.
474            If None, a new session is created and closed after all pages are fetched.
475
476    Yields:
477        dict | list: Raw JSON responses, one element per page.
478
479    Raises:
480        requests.HTTPError: If any HTTP request returns an error status.
481        requests.ConnectionError: If the connection fails after all retries.
482    """
483    own_session = session is None
484    if own_session:
485        session = _build_session(max_retries)
486
487    if headers:
488        session.headers.update(headers)
489
490    url: str | None = url_rest_api
491    params = api_params or {}
492    page = 0
493
494    try:
495        while url:
496            page += 1
497            logger.debug(f"Fetching page {page}: {url}")
498            response = session.get(
499                url,
500                params=params if page == 1 else None,
501                timeout=timeout,
502            )
503            response.raise_for_status()
504            data = response.json()
505            yield data
506
507            # Advance to next page only when response is a dict with a next link
508            url = data.get(next_key) if isinstance(data, dict) else None
509    finally:
510        if own_session:
511            session.close()
512
513
514def df_from_url(
515    url_rest_api: str,
516    api_params: dict | None = None,
517    headers: dict | None = None,
518    results_key: str | None = 'results',
519    next_key: str = 'next',
520    timeout: int | tuple[int, int] = (10, 30),
521    max_retries: int = 3,
522    session: requests.Session | None = None,
523) -> DataFrame | None:
524    """
525    Fetch paginated JSON from a REST API and return a Pandas DataFrame.
526
527    Delegates HTTP handling and pagination to :func:`_fetch_pages`.
528
529    Args:
530        url_rest_api (str): The base URL of the API endpoint.
531        api_params (dict, optional): Query parameters for the initial request.
532        headers (dict, optional): HTTP headers for the request.
533        results_key (str | None, optional): Key in the JSON response that contains
534            the data list. If ``None``, the entire response body is wrapped as a
535            single record. If the key is absent, the first list value found in the
536            dict is used as fallback. Defaults to ``'results'``.
537        next_key (str): Key in the JSON response containing the next-page URL.
538            Defaults to ``'next'``.
539        timeout (int | tuple[int, int]): Request timeout ``(connect, read)`` in seconds.
540            Defaults to ``(10, 30)``.
541        max_retries (int): Retries on transient HTTP errors. Defaults to ``3``.
542        session (requests.Session, optional): Existing session to reuse.
543            If None, a new session is created and closed after use.
544
545    Returns:
546        DataFrame | None: A DataFrame with all collected data, or ``None`` if empty.
547
548    Raises:
549        requests.HTTPError: If any HTTP request returns an error status.
550        requests.ConnectionError: If the connection fails after all retries.
551        ValueError: If the JSON response has an unexpected structure.
552    """
553    all_data: list = []
554
555    for data in _iter_fetch_pages(url_rest_api, api_params, headers, next_key, timeout, max_retries, session):
556        if isinstance(data, list):
557            page_data = data
558        elif isinstance(data, dict):
559            if results_key and results_key in data:
560                page_data = data[results_key]
561            elif results_key is None:
562                page_data = [data]
563            else:
564                # Fallback: primer valor de tipo lista encontrado en el dict
565                page_data = next(
566                    (v for v in data.values() if isinstance(v, list)), [data]
567                )
568        else:
569            raise ValueError(
570                f"Unexpected JSON structure: expected list or dict, got {type(data).__name__}"
571            )
572
573        all_data.extend(page_data)
574        logger.debug(f"Got {len(page_data)} records (total so far: {len(all_data)})")
575
576    return DataFrame(all_data) if all_data else None