meerschaum.utils.dataframe

   1#! /usr/bin/env python3
   2# -*- coding: utf-8 -*-
   3# vim:fenc=utf-8
   4
   5"""
   6Utility functions for working with DataFrames.
   7"""
   8
   9from __future__ import annotations
  10from datetime import datetime
  11from meerschaum.utils.typing import (
  12    Optional, Dict, Any, List, Hashable, Generator,
  13    Iterator, Iterable, Union, Tuple,
  14)
  15
  16
  17def add_missing_cols_to_df(df: 'pd.DataFrame', dtypes: Dict[str, Any]) -> pd.DataFrame:
  18    """
  19    Add columns from the dtypes dictionary as null columns to a new DataFrame.
  20
  21    Parameters
  22    ----------
  23    df: pd.DataFrame
  24        The dataframe we should copy and add null columns.
  25
  26    dtypes:
  27        The data types dictionary which may contain keys not present in `df.columns`.
  28
  29    Returns
  30    -------
  31    A new `DataFrame` with the keys from `dtypes` added as null columns.
  32    If `df.dtypes` is the same as `dtypes`, then return a reference to `df`.
  33    NOTE: This will not ensure that dtypes are enforced!
  34
  35    Examples
  36    --------
  37    >>> import pandas as pd
  38    >>> df = pd.DataFrame([{'a': 1}])
  39    >>> dtypes = {'b': 'Int64'}
  40    >>> add_missing_cols_to_df(df, dtypes)
  41          a  b
  42       0  1  <NA>
  43    >>> add_missing_cols_to_df(df, dtypes).dtypes
  44    a    int64
  45    b    Int64
  46    dtype: object
  47    >>> add_missing_cols_to_df(df, {'a': 'object'}).dtypes
  48    a    int64
  49    dtype: object
  50    >>> 
  51    """
  52    if set(df.columns) == set(dtypes):
  53        return df
  54
  55    import traceback
  56    from meerschaum.utils.packages import import_pandas, attempt_import
  57    from meerschaum.utils.warnings import warn
  58    from meerschaum.utils.dtypes import to_pandas_dtype
  59    pandas = attempt_import('pandas')
  60    
  61    def build_series(dtype: str):
  62        return pandas.Series([], dtype=to_pandas_dtype(dtype))
  63
  64    assign_kwargs = {
  65        str(col): build_series(str(typ))
  66        for col, typ in dtypes.items()
  67        if col not in df.columns
  68    }
  69    return df.assign(**assign_kwargs)
  70
  71
  72def filter_unseen_df(
  73        old_df: 'pd.DataFrame',
  74        new_df: 'pd.DataFrame',
  75        safe_copy: bool = True,
  76        dtypes: Optional[Dict[str, Any]] = None,
  77        debug: bool = False,
  78    ) -> 'pd.DataFrame':
  79    """
  80    Left join two DataFrames to find the newest unseen data.
  81
  82    Parameters
  83    ----------
  84    old_df: 'pd.DataFrame'
  85        The original (target) dataframe. Acts as a filter on the `new_df`.
  86        
  87    new_df: 'pd.DataFrame'
  88        The fetched (source) dataframe. Rows that are contained in `old_df` are removed.
  89
  90    safe_copy: bool, default True
  91        If `True`, create a copy before comparing and modifying the dataframes.
  92        Setting to `False` may mutate the DataFrames.
  93        
  94    dtypes: Optional[Dict[str, Any]], default None
  95        Optionally specify the datatypes of the dataframe.
  96
  97    debug: bool, default False
  98        Verbosity toggle.
  99
 100    Returns
 101    -------
 102    A pandas dataframe of the new, unseen rows in `new_df`.
 103
 104    Examples
 105    --------
 106    ```python
 107    >>> import pandas as pd
 108    >>> df1 = pd.DataFrame({'a': [1,2]})
 109    >>> df2 = pd.DataFrame({'a': [2,3]})
 110    >>> filter_unseen_df(df1, df2)
 111       a
 112    0  3
 113
 114    ```
 115
 116    """
 117    if old_df is None:
 118        return new_df
 119
 120    if safe_copy:
 121        old_df = old_df.copy()
 122        new_df = new_df.copy()
 123
 124    import json
 125    import functools
 126    import traceback
 127    from decimal import Decimal
 128    from meerschaum.utils.warnings import warn
 129    from meerschaum.utils.packages import import_pandas, attempt_import
 130    from meerschaum.utils.dtypes import to_pandas_dtype, are_dtypes_equal, attempt_cast_to_numeric
 131    from meerschaum.utils.debug import dprint
 132    pd = import_pandas(debug=debug)
 133    is_dask = 'dask' in new_df.__module__
 134    if is_dask:
 135        pandas = attempt_import('pandas')
 136        dd = attempt_import('dask.dataframe')
 137        merge = dd.merge
 138        NA = pandas.NA
 139    else:
 140        merge = pd.merge
 141        NA = pd.NA
 142
 143    new_df_dtypes = dict(new_df.dtypes)
 144    old_df_dtypes = dict(old_df.dtypes)
 145
 146    same_cols = set(new_df.columns) == set(old_df.columns)
 147    if not same_cols:
 148        new_df = add_missing_cols_to_df(new_df, old_df_dtypes)
 149        old_df = add_missing_cols_to_df(old_df, new_df_dtypes)
 150
 151    ### Edge case: two empty lists cast to DFs.
 152    elif len(new_df.columns) == 0:
 153        return new_df
 154
 155    try:
 156        ### Order matters when checking equality.
 157        new_df = new_df[old_df.columns]
 158    except Exception as e:
 159        warn(
 160            "Was not able to cast old columns onto new DataFrame. " +
 161            f"Are both DataFrames the same shape? Error:\n{e}",
 162            stacklevel = 3,
 163        )
 164        return new_df[list(new_df_dtypes.keys())]
 165
 166    ### assume the old_df knows what it's doing, even if it's technically wrong.
 167    if dtypes is None:
 168        dtypes = {col: str(typ) for col, typ in old_df.dtypes.items()}
 169
 170    dtypes = {
 171        col: to_pandas_dtype(typ)
 172        for col, typ in dtypes.items()
 173        if col in new_df_dtypes and col in old_df_dtypes
 174    }
 175    for col, typ in new_df_dtypes.items():
 176        if col not in dtypes:
 177            dtypes[col] = typ
 178    
 179    cast_cols = True
 180    try:
 181        new_df = new_df.astype(dtypes)
 182        cast_cols = False
 183    except Exception as e:
 184        warn(
 185            f"Was not able to cast the new DataFrame to the given dtypes.\n{e}"
 186        )
 187
 188    new_numeric_cols_existing = get_numeric_cols(new_df)
 189    old_numeric_cols = get_numeric_cols(old_df)
 190    for col, typ in {k: v for k, v in dtypes.items()}.items():
 191        if not are_dtypes_equal(new_df_dtypes.get(col, 'None'), old_df_dtypes.get(col, 'None')):
 192            new_is_float = are_dtypes_equal(new_df_dtypes.get(col, 'None'), 'float')
 193            new_is_int = are_dtypes_equal(new_df_dtypes.get(col, 'None'), 'int')
 194            new_is_numeric = col in new_numeric_cols_existing
 195            old_is_float = are_dtypes_equal(old_df_dtypes.get(col, 'None'), 'float')
 196            old_is_int = are_dtypes_equal(old_df_dtypes.get(col, 'None'), 'int')
 197            old_is_numeric = col in old_numeric_cols
 198
 199            if (
 200                (new_is_float or new_is_int or new_is_numeric)
 201                and
 202                (old_is_float or old_is_int or old_is_numeric)
 203            ):
 204                dtypes[col] = attempt_cast_to_numeric
 205                cast_cols = True
 206                continue
 207
 208            ### Fallback to object if the types don't match.
 209            warn(
 210                f"Detected different types for '{col}' "
 211                + f"({new_df_dtypes.get(col, None)} vs {old_df_dtypes.get(col, None)}), "
 212                + "falling back to 'object'..."
 213            )
 214            dtypes[col] = 'object'
 215            cast_cols = True
 216
 217    if cast_cols:
 218        for col, dtype in dtypes.items():
 219            if col in new_df.columns:
 220                try:
 221                    new_df[col] = (
 222                        new_df[col].astype(dtype)
 223                        if not callable(dtype)
 224                        else new_df[col].apply(dtype)
 225                    )
 226                except Exception as e:
 227                    warn(f"Was not able to cast column '{col}' to dtype '{dtype}'.\n{e}")
 228
 229    serializer = functools.partial(json.dumps, sort_keys=True, separators=(',', ':'), default=str)
 230    new_json_cols = get_json_cols(new_df)
 231    old_json_cols = get_json_cols(old_df)
 232    json_cols = set(new_json_cols + old_json_cols)
 233    for json_col in old_json_cols:
 234        old_df[json_col] = old_df[json_col].apply(serializer)
 235    for json_col in new_json_cols:
 236        new_df[json_col] = new_df[json_col].apply(serializer)
 237
 238    new_numeric_cols = get_numeric_cols(new_df)
 239    numeric_cols = set(new_numeric_cols + old_numeric_cols)
 240    for numeric_col in old_numeric_cols:
 241        old_df[numeric_col] = old_df[numeric_col].apply(
 242            lambda x: f'{x:f}' if isinstance(x, Decimal) else x
 243        )
 244    for numeric_col in new_numeric_cols:
 245        new_df[numeric_col] = new_df[numeric_col].apply(
 246            lambda x: f'{x:f}' if isinstance(x, Decimal) else x
 247        )
 248
 249    joined_df = merge(
 250        new_df.fillna(NA),
 251        old_df.fillna(NA),
 252        how = 'left',
 253        on = None,
 254        indicator = True,
 255    )
 256    changed_rows_mask = (joined_df['_merge'] == 'left_only')
 257    delta_df = joined_df[list(new_df_dtypes.keys())][changed_rows_mask].reset_index(drop=True)
 258
 259    for json_col in json_cols:
 260        if json_col not in delta_df.columns:
 261            continue
 262        try:
 263            delta_df[json_col] = delta_df[json_col].apply(json.loads)
 264        except Exception as e:
 265            warn(f"Unable to deserialize JSON column '{json_col}':\n{traceback.format_exc()}")
 266
 267    for numeric_col in numeric_cols:
 268        if numeric_col not in delta_df.columns:
 269            continue
 270        try:
 271            delta_df[numeric_col] = delta_df[numeric_col].apply(attempt_cast_to_numeric)
 272        except Exception as e:
 273            warn(f"Unable to parse numeric column '{numeric_col}':\n{traceback.format_exc()}")
 274
 275    return delta_df
 276
 277
 278def parse_df_datetimes(
 279        df: 'pd.DataFrame',
 280        ignore_cols: Optional[Iterable[str]] = None,
 281        chunksize: Optional[int] = None,
 282        dtype_backend: str = 'numpy_nullable',
 283        debug: bool = False,
 284    ) -> 'pd.DataFrame':
 285    """
 286    Parse a pandas DataFrame for datetime columns and cast as datetimes.
 287
 288    Parameters
 289    ----------
 290    df: pd.DataFrame
 291        The pandas DataFrame to parse.
 292
 293    ignore_cols: Optional[Iterable[str]], default None
 294        If provided, do not attempt to coerce these columns as datetimes.
 295
 296    chunksize: Optional[int], default None
 297        If the pandas implementation is `'dask'`, use this chunksize for the distributed dataframe.
 298
 299    dtype_backend: str, default 'numpy_nullable'
 300        If `df` is not a DataFrame and new one needs to be constructed,
 301        use this as the datatypes backend.
 302        Accepted values are 'numpy_nullable' and 'pyarrow'.
 303        
 304    debug: bool, default False
 305        Verbosity toggle.
 306
 307    Returns
 308    -------
 309    A new pandas DataFrame with the determined datetime columns
 310    (usually ISO strings) cast as datetimes.
 311
 312    Examples
 313    --------
 314    ```python
 315    >>> import pandas as pd
 316    >>> df = pd.DataFrame({'a': ['2022-01-01 00:00:00']}) 
 317    >>> df.dtypes
 318    a    object
 319    dtype: object
 320    >>> df = parse_df_datetimes(df)
 321    >>> df.dtypes
 322    a    datetime64[ns]
 323    dtype: object
 324
 325    ```
 326
 327    """
 328    from meerschaum.utils.packages import import_pandas, attempt_import
 329    from meerschaum.utils.debug import dprint
 330    from meerschaum.utils.warnings import warn
 331    import traceback
 332    pd = import_pandas()
 333    pandas = attempt_import('pandas')
 334    pd_name = pd.__name__
 335    using_dask = 'dask' in pd_name
 336    df_is_dask = (hasattr(df, '__module__') and 'dask' in df.__module__)
 337    dask_dataframe = None
 338    if using_dask or df_is_dask:
 339        npartitions = chunksize_to_npartitions(chunksize)
 340        dask_dataframe = attempt_import('dask.dataframe')
 341
 342    ### if df is a dict, build DataFrame
 343    if isinstance(df, pandas.DataFrame):
 344        pdf = df
 345    elif df_is_dask and isinstance(df, dask_dataframe.DataFrame):
 346        pdf = get_first_valid_dask_partition(df)
 347    else:
 348        if debug:
 349            dprint(f"df is of type '{type(df)}'. Building {pd.DataFrame}...")
 350
 351        if using_dask:
 352            if isinstance(df, list):
 353                keys = set()
 354                for doc in df:
 355                    for key in doc:
 356                        keys.add(key)
 357                df = pd.DataFrame.from_dict(
 358                    {
 359                        k: [
 360                            doc.get(k, None)
 361                            for doc in df
 362                        ] for k in keys
 363                    },
 364                    npartitions = npartitions,
 365                )
 366            elif isinstance(df, dict):
 367                df = pd.DataFrame.from_dict(df, npartitions=npartitions)
 368            elif 'pandas.core.frame.DataFrame' in str(type(df)):
 369                df = pd.from_pandas(df, npartitions=npartitions)
 370            else:
 371                raise Exception("Can only parse dictionaries or lists of dictionaries with Dask.")
 372            pandas = attempt_import('pandas')
 373            pdf = get_first_valid_dask_partition(df)
 374
 375        else:
 376            df = pd.DataFrame(df).convert_dtypes(dtype_backend=dtype_backend)
 377            pdf = df
 378
 379    ### skip parsing if DataFrame is empty
 380    if len(pdf) == 0:
 381        if debug:
 382            dprint(f"df is empty. Returning original DataFrame without casting datetime columns...")
 383        return df
 384
 385    ignore_cols = set(
 386        (ignore_cols or []) + [
 387            col
 388            for col, dtype in pdf.dtypes.items() 
 389            if 'datetime' in str(dtype)
 390        ]
 391    )
 392    cols_to_inspect = [col for col in pdf.columns if col not in ignore_cols]
 393
 394    if len(cols_to_inspect) == 0:
 395        if debug:
 396            dprint(f"All columns are ignored, skipping datetime detection...")
 397        return df
 398
 399    ### apply regex to columns to determine which are ISO datetimes
 400    iso_dt_regex = r'\d{4}-\d{2}-\d{2}.\d{2}\:\d{2}\:\d+'
 401    dt_mask = pdf[cols_to_inspect].astype(str).apply(
 402        lambda s: s.str.match(iso_dt_regex).all()
 403    )
 404
 405    ### list of datetime column names
 406    datetime_cols = [col for col in pdf[cols_to_inspect].loc[:, dt_mask]]
 407    if not datetime_cols:
 408        if debug:
 409            dprint("No columns detected as datetimes, returning...")
 410        return df
 411
 412    if debug:
 413        dprint("Converting columns to datetimes: " + str(datetime_cols))
 414
 415    try:
 416        if not using_dask:
 417            df[datetime_cols] = df[datetime_cols].apply(pd.to_datetime, utc=True)
 418        else:
 419            df[datetime_cols] = df[datetime_cols].apply(
 420                pd.to_datetime,
 421                utc = True,
 422                axis = 1,
 423                meta = {
 424                    col: 'datetime64[ns]'
 425                    for col in datetime_cols
 426                }
 427            )
 428    except Exception as e:
 429        warn(
 430            f"Unable to apply `pd.to_datetime` to {items_str(datetime_cols)}:\n"
 431            + f"{traceback.format_exc()}"
 432        )
 433
 434    for dt in datetime_cols:
 435        try:
 436            df[dt] = df[dt].dt.tz_localize(None)
 437        except Exception as e:
 438            warn(f"Unable to convert column '{dt}' to naive datetime:\n{traceback.format_exc()}")
 439
 440    return df
 441
 442
 443def get_unhashable_cols(df: 'pd.DataFrame') -> List[str]:
 444    """
 445    Get the columns which contain unhashable objects from a Pandas DataFrame.
 446
 447    Parameters
 448    ----------
 449    df: pd.DataFrame
 450        The DataFrame which may contain unhashable objects.
 451
 452    Returns
 453    -------
 454    A list of columns.
 455    """
 456    if len(df) == 0:
 457        return []
 458
 459    is_dask = 'dask' in df.__module__
 460    if is_dask:
 461        from meerschaum.utils.packages import attempt_import
 462        pandas = attempt_import('pandas')
 463        df = pandas.DataFrame(get_first_valid_dask_partition(df))
 464    return [
 465        col for col, val in df.iloc[0].items()
 466        if not isinstance(val, Hashable)
 467    ]
 468
 469
 470def get_json_cols(df: 'pd.DataFrame') -> List[str]:
 471    """
 472    Get the columns which contain unhashable objects from a Pandas DataFrame.
 473
 474    Parameters
 475    ----------
 476    df: pd.DataFrame
 477        The DataFrame which may contain unhashable objects.
 478
 479    Returns
 480    -------
 481    A list of columns to be encoded as JSON.
 482    """
 483    is_dask = 'dask' in df.__module__
 484    if is_dask:
 485        df = get_first_valid_dask_partition(df)
 486    
 487    if len(df) == 0:
 488        return []
 489
 490    cols_indices = {
 491        col: df[col].first_valid_index()
 492        for col in df.columns
 493    }
 494    return [
 495        col
 496        for col, ix in cols_indices.items()
 497        if (
 498            ix is not None
 499            and
 500            not isinstance(df.loc[ix][col], Hashable)
 501        )
 502    ]
 503
 504
 505def get_numeric_cols(df: 'pd.DataFrame') -> List[str]:
 506    """
 507    Get the columns which contain `decimal.Decimal` objects from a Pandas DataFrame.
 508
 509    Parameters
 510    ----------
 511    df: pd.DataFrame
 512        The DataFrame which may contain decimal objects.
 513
 514    Returns
 515    -------
 516    A list of columns to treat as numerics.
 517    """
 518    from decimal import Decimal
 519    is_dask = 'dask' in df.__module__
 520    if is_dask:
 521        df = get_first_valid_dask_partition(df)
 522    
 523    if len(df) == 0:
 524        return []
 525
 526    cols_indices = {
 527        col: df[col].first_valid_index()
 528        for col in df.columns
 529    }
 530    return [
 531        col
 532        for col, ix in cols_indices.items()
 533        if (
 534            ix is not None
 535            and
 536            isinstance(df.loc[ix][col], Decimal)
 537        )
 538    ]
 539
 540
 541def enforce_dtypes(
 542        df: 'pd.DataFrame',
 543        dtypes: Dict[str, str],
 544        safe_copy: bool = True,
 545        coerce_numeric: bool = True,
 546        debug: bool = False,
 547    ) -> 'pd.DataFrame':
 548    """
 549    Enforce the `dtypes` dictionary on a DataFrame.
 550
 551    Parameters
 552    ----------
 553    df: pd.DataFrame
 554        The DataFrame on which to enforce dtypes.
 555
 556    dtypes: Dict[str, str]
 557        The data types to attempt to enforce on the DataFrame.
 558
 559    safe_copy: bool, default True
 560        If `True`, create a copy before comparing and modifying the dataframes.
 561        Setting to `False` may mutate the DataFrames.
 562        See `meerschaum.utils.dataframe.filter_unseen_df`.
 563
 564    coerce_numeric: bool, default True
 565        If `True`, convert float and int collisions to numeric.
 566
 567    debug: bool, default False
 568        Verbosity toggle.
 569
 570    Returns
 571    -------
 572    The Pandas DataFrame with the types enforced.
 573    """
 574    import json
 575    import traceback
 576    from decimal import Decimal
 577    from meerschaum.utils.debug import dprint
 578    from meerschaum.utils.warnings import warn
 579    from meerschaum.utils.formatting import pprint
 580    from meerschaum.config.static import STATIC_CONFIG
 581    from meerschaum.utils.packages import import_pandas
 582    from meerschaum.utils.dtypes import (
 583        are_dtypes_equal,
 584        to_pandas_dtype,
 585        is_dtype_numeric,
 586        attempt_cast_to_numeric,
 587    )
 588    if safe_copy:
 589        df = df.copy()
 590    df_dtypes = {c: str(t) for c, t in df.dtypes.items()}
 591    if len(df_dtypes) == 0:
 592        if debug:
 593            dprint("Incoming DataFrame has no columns. Skipping enforcement...")
 594        return df
 595
 596    pipe_pandas_dtypes = {
 597        col: to_pandas_dtype(typ)
 598        for col, typ in dtypes.items()
 599    }
 600    json_cols = [
 601        col
 602        for col, typ in dtypes.items()
 603        if typ == 'json'
 604    ]
 605    numeric_cols = [
 606        col
 607        for col, typ in dtypes.items()
 608        if typ == 'numeric'
 609    ]
 610    df_numeric_cols = get_numeric_cols(df)
 611    if debug:
 612        dprint(f"Desired data types:")
 613        pprint(dtypes)
 614        dprint(f"Data types for incoming DataFrame:")
 615        pprint(df_dtypes)
 616
 617    if json_cols and len(df) > 0:
 618        if debug:
 619            dprint(f"Checking columns for JSON encoding: {json_cols}")
 620        for col in json_cols:
 621            if col in df.columns:
 622                try:
 623                    df[col] = df[col].apply(
 624                        (
 625                            lambda x: (
 626                                json.loads(x)
 627                                if isinstance(x, str)
 628                                else x
 629                            )
 630                        )
 631                    )
 632                except Exception as e:
 633                    if debug:
 634                        dprint(f"Unable to parse column '{col}' as JSON:\n{e}")
 635
 636    if numeric_cols:
 637        if debug:
 638            dprint(f"Checking for numerics: {numeric_cols}")
 639        for col in numeric_cols:
 640            if col in df.columns:
 641                try:
 642                    df[col] = df[col].apply(attempt_cast_to_numeric)
 643                except Exception as e:
 644                    if debug:
 645                        dprint(f"Unable to parse column '{col}' as NUMERIC:\n{e}")
 646
 647    if are_dtypes_equal(df_dtypes, pipe_pandas_dtypes):
 648        if debug:
 649            dprint(f"Data types match. Exiting enforcement...")
 650        return df
 651
 652    common_dtypes = {}
 653    common_diff_dtypes = {}
 654    for col, typ in pipe_pandas_dtypes.items():
 655        if col in df_dtypes:
 656            common_dtypes[col] = typ
 657            if not are_dtypes_equal(typ, df_dtypes[col]):
 658                common_diff_dtypes[col] = df_dtypes[col]
 659
 660    if debug:
 661        dprint(f"Common columns with different dtypes:")
 662        pprint(common_diff_dtypes)
 663
 664    detected_dt_cols = {}
 665    for col, typ in common_diff_dtypes.items():
 666        if 'datetime' in typ and 'datetime' in common_dtypes[col]:
 667            df_dtypes[col] = typ
 668            detected_dt_cols[col] = (common_dtypes[col], common_diff_dtypes[col])
 669    for col in detected_dt_cols:
 670        del common_diff_dtypes[col]
 671
 672    if debug:
 673        dprint(f"Common columns with different dtypes (after dates):")
 674        pprint(common_diff_dtypes)
 675
 676    if are_dtypes_equal(df_dtypes, pipe_pandas_dtypes):
 677        if debug:
 678            dprint(
 679                "The incoming DataFrame has mostly the same types, skipping enforcement."
 680                + f"The only detected difference was in the following datetime columns.\n"
 681                + "    Timezone information may be stripped."
 682            )
 683            pprint(detected_dt_cols)
 684        return df
 685
 686    for col, typ in {k: v for k, v in common_diff_dtypes.items()}.items():
 687        previous_typ = common_dtypes[col]
 688        mixed_numeric_types = (is_dtype_numeric(typ) and is_dtype_numeric(previous_typ))
 689        explicitly_float = are_dtypes_equal(dtypes.get(col, 'object'), 'float')
 690        explicitly_numeric = dtypes.get(col, 'numeric') == 'numeric'
 691        cast_to_numeric = (
 692            explicitly_numeric
 693            or col in df_numeric_cols
 694            or (mixed_numeric_types and not explicitly_float)
 695        ) and coerce_numeric
 696        if cast_to_numeric:
 697            common_dtypes[col] = attempt_cast_to_numeric
 698            common_diff_dtypes[col] = attempt_cast_to_numeric
 699
 700    for d in common_diff_dtypes:
 701        t = common_dtypes[d]
 702        if debug:
 703            dprint(f"Casting column {d} to dtype {t}.")
 704        try:
 705            df[d] = (
 706                df[d].apply(t)
 707                if callable(t)
 708                else df[d].astype(t)
 709            )
 710        except Exception as e:
 711            if debug:
 712                dprint(f"Encountered an error when casting column {d} to type {t}:\n{e}")
 713            if 'int' in str(t).lower():
 714                try:
 715                    df[d] = df[d].astype('float64').astype(t)
 716                except Exception as e:
 717                    if debug:
 718                        dprint(f"Was unable to convert to float then {t}.")
 719    return df
 720
 721
 722def get_datetime_bound_from_df(
 723        df: Union['pd.DataFrame', dict, list],
 724        datetime_column: str,
 725        minimum: bool = True,
 726    ) -> Union[int, 'datetime.datetime', None]:
 727    """
 728    Return the minimum or maximum datetime (or integer) from a DataFrame.
 729
 730    Parameters
 731    ----------
 732    df: pd.DataFrame
 733        The DataFrame, list, or dict which contains the range axis.
 734
 735    datetime_column: str
 736        The name of the datetime (or int) column.
 737
 738    minimum: bool
 739        Whether to return the minimum (default) or maximum value.
 740
 741    Returns
 742    -------
 743    The minimum or maximum datetime value in the dataframe, or `None`.
 744    """
 745    if not datetime_column:
 746        return None
 747
 748    def compare(a, b):
 749        if a is None:
 750            return b
 751        if b is None:
 752            return a
 753        if minimum:
 754            return a if a < b else b
 755        return a if a > b else b
 756
 757    if isinstance(df, list):
 758        if len(df) == 0:
 759            return None
 760        best_yet = df[0].get(datetime_column, None)
 761        for doc in df:
 762            val = doc.get(datetime_column, None)
 763            best_yet = compare(best_yet, val)
 764        return best_yet
 765
 766    if isinstance(df, dict):
 767        if datetime_column not in df:
 768            return None
 769        best_yet = df[datetime_column][0]
 770        for val in df[datetime_column]:
 771            best_yet = compare(best_yet, val)
 772        return best_yet
 773
 774    if 'DataFrame' in str(type(df)):
 775        if datetime_column not in df.columns:
 776            return None
 777        return (
 778            df[datetime_column].min(skipna=True)
 779            if minimum
 780            else df[datetime_column].max(skipna=True)
 781        )
 782
 783    return None
 784
 785
 786def df_is_chunk_generator(df: Any) -> bool:
 787    """
 788    Determine whether to treat `df` as a chunk generator.
 789
 790    Note this should only be used in a context where generators are expected,
 791    as it will return `True` for any iterable.
 792
 793    Parameters
 794    ----------
 795    The DataFrame or chunk generator to evaluate.
 796
 797    Returns
 798    -------
 799    A `bool` indicating whether to treat `df` as a generator.
 800    """
 801    return (
 802        not isinstance(df, (dict, list, str))
 803        and 'DataFrame' not in str(type(df))
 804        and isinstance(df, (Generator, Iterable, Iterator))
 805    )
 806
 807
 808def chunksize_to_npartitions(chunksize: Optional[int]) -> int:
 809    """
 810    Return the Dask `npartitions` value for a given `chunksize`.
 811    """
 812    if chunksize == -1:
 813        from meerschaum.config import get_config
 814        chunksize = get_config('system', 'connectors', 'sql', 'chunksize')
 815    if chunksize is None:
 816        return 1
 817    return -1 * chunksize
 818
 819
 820def df_from_literal(
 821        pipe: Optional['meerschaum.Pipe'] = None,
 822        literal: str = None,
 823        debug: bool = False
 824    ) -> 'pd.DataFrame':
 825    """
 826    Construct a dataframe from a literal value, using the pipe's datetime and value column names.
 827
 828    Parameters
 829    ----------
 830    pipe: Optional['meerschaum.Pipe'], default None
 831        The pipe which will consume the literal value.
 832
 833    Returns
 834    -------
 835    A 1-row pandas DataFrame from with the current UTC timestamp as the datetime columns
 836    and the literal as the value.
 837    """
 838    from meerschaum.utils.packages import import_pandas
 839    from meerschaum.utils.warnings import error, warn
 840    from meerschaum.utils.debug import dprint
 841
 842    if pipe is None or literal is None:
 843        error("Please provide a Pipe and a literal value")
 844    ### this will raise an error if the columns are undefined
 845    dt_name, val_name = pipe.get_columns('datetime', 'value')
 846
 847    val = literal
 848    if isinstance(literal, str):
 849        if debug:
 850            dprint(f"Received literal string: '{literal}'")
 851        import ast
 852        try:
 853            val = ast.literal_eval(literal)
 854        except Exception as e:
 855            warn(
 856                "Failed to parse value from string:\n" + f"{literal}" +
 857                "\n\nWill cast as a string instead."\
 858            )
 859            val = literal
 860
 861    from datetime import datetime, timezone
 862    now = datetime.now(timezone.utc).replace(tzinfo=None)
 863
 864    pd = import_pandas()
 865    return pd.DataFrame({dt_name: [now], val_name: [val]})
 866
 867
 868def get_first_valid_dask_partition(ddf: 'dask.dataframe.DataFrame') -> Union['pd.DataFrame', None]:
 869    """
 870    Return the first valid Dask DataFrame partition (if possible).
 871    """
 872    pdf = None
 873    for partition in ddf.partitions:
 874        try:
 875            pdf = partition.compute()
 876        except Exception as e:
 877            continue
 878        if len(pdf) > 0:
 879            return pdf
 880    return ddf.compute()
 881
 882
 883def query_df(
 884        df: 'pd.DataFrame',
 885        params: Optional[Dict[str, Any]] = None,
 886        begin: Union[datetime, int, None] = None,
 887        end: Union[datetime, int, None] = None,
 888        datetime_column: Optional[str] = None,
 889        select_columns: Optional[List[str]] = None,
 890        omit_columns: Optional[List[str]] = None,
 891        inplace: bool = False,
 892        reset_index: bool = False,
 893        debug: bool = False,
 894    ) -> 'pd.DataFrame':
 895    """
 896    Query the dataframe with the params dictionary.
 897
 898    Parameters
 899    ----------
 900    df: pd.DataFrame
 901        The DataFrame to query against.
 902
 903    params: Optional[Dict[str, Any]], default None
 904        The parameters dictionary to use for the query.
 905
 906    begin: Union[datetime, int, None], default None
 907        If `begin` and `datetime_column` are provided, only return rows with a timestamp
 908        greater than or equal to this value.
 909
 910    end: Union[datetime, int, None], default None
 911        If `begin` and `datetime_column` are provided, only return rows with a timestamp
 912        less than this value.
 913
 914    datetime_column: Optional[str], default None
 915        A `datetime_column` must be provided to use `begin` and `end`.
 916
 917    select_columns: Optional[List[str]], default None
 918        If provided, only return these columns.
 919
 920    omit_columns: Optional[List[str]], default None
 921        If provided, do not include these columns in the result.
 922
 923    inplace: bool, default False
 924        If `True`, modify the DataFrame inplace rather than creating a new DataFrame.
 925
 926    reset_index: bool, default True
 927        If `True`, reset the index in the resulting DataFrame.
 928
 929    Returns
 930    -------
 931    A Pandas DataFrame query result.
 932    """
 933    if not params and not begin and not end:
 934        return df
 935
 936    import json
 937    import meerschaum as mrsm
 938    from meerschaum.utils.debug import dprint
 939    from meerschaum.utils.misc import get_in_ex_params
 940    from meerschaum.utils.warnings import warn
 941
 942    dtypes = {col: str(typ) for col, typ in df.dtypes.items()}
 943
 944    if begin or end:
 945        if not datetime_column or datetime_column not in df.columns:
 946            warn(
 947                f"The datetime column '{datetime_column}' is not present in the Dataframe, "
 948                + "ignoring begin and end...",
 949            )
 950            begin, end = None, None
 951
 952    if debug:
 953        dprint(f"Querying dataframe:\n{params=} {begin=} {end=} {datetime_column=}")
 954
 955    in_ex_params = get_in_ex_params(params)
 956
 957    def serialize(x: Any) -> str:
 958        if isinstance(x, (dict, list, tuple)):
 959            return json.dumps(x, sort_keys=True, separators=(',', ':'), default=str)
 960        if hasattr(x, 'isoformat'):
 961            return x.isoformat()
 962        return str(x)
 963
 964    masks = [
 965        (
 966            (df[datetime_column] >= begin)
 967            if begin is not None and datetime_column
 968            else True
 969        ) & (
 970            (df[datetime_column] < end)
 971            if end is not None and datetime_column
 972            else True
 973        )
 974    ]
 975
 976    masks.extend([
 977        (
 978            (
 979                df[col].apply(serialize).isin(
 980                    [
 981                        serialize(_in_val)
 982                        for _in_val in in_vals
 983                    ]
 984                ) if in_vals else True
 985            ) & (
 986                ~df[col].apply(serialize).isin(
 987                    [
 988                        serialize(_ex_val)
 989                        for _ex_val in ex_vals
 990                    ]
 991                ) if ex_vals else True
 992            )
 993        )
 994        for col, (in_vals, ex_vals) in in_ex_params.items()
 995        if col in df.columns
 996    ])
 997    query_mask = masks[0]
 998    for mask in masks:
 999        query_mask = query_mask & mask
1000
1001    if inplace:
1002        df.where(query_mask, inplace=inplace)
1003        df.dropna(how='all', inplace=inplace)
1004        result_df = df
1005    else:
1006        result_df = df.where(query_mask).dropna(how='all')
1007
1008    if reset_index:
1009        result_df.reset_index(drop=True, inplace=True)
1010
1011    result_df = enforce_dtypes(
1012        result_df,
1013        dtypes,
1014        safe_copy = (not inplace),
1015        debug = debug,
1016        coerce_numeric = False,
1017    )
1018
1019    if select_columns == ['*']:
1020        select_columns = None
1021
1022    if not select_columns and not omit_columns:
1023        return result_df
1024
1025    if select_columns:
1026        for col in list(result_df.columns):
1027            if col not in select_columns:
1028                del result_df[col]
1029        return result_df
1030
1031    if omit_columns:
1032        for col in list(result_df.columns):
1033            if col in omit_columns:
1034                del result_df[col]
1035    if debug:
1036        dprint(f"{dtypes=}")
1037    return result_df