meerschaum.utils.dataframe
Utility functions for working with DataFrames.
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 10 11from datetime import datetime, timezone, date 12from collections import defaultdict 13 14import meerschaum as mrsm 15from meerschaum.utils.typing import ( 16 Optional, Dict, Any, List, Hashable, Generator, 17 Iterator, Iterable, Union, TYPE_CHECKING, Tuple, 18) 19 20if TYPE_CHECKING: 21 pd, dask = mrsm.attempt_import('pandas', 'dask') 22 23 24_POLARS_FILTER_MIN_ROWS: int = 100_000 25 26 27def _filter_unseen_df_with_polars( 28 new_df: 'pd.DataFrame', 29 old_df: 'pd.DataFrame', 30) -> Union['pd.DataFrame', None]: 31 """Return an anti-join of ``new_df`` against ``old_df``, or ``None`` to fall back.""" 32 if len(new_df) + len(old_df) < _POLARS_FILTER_MIN_ROWS: 33 return None 34 35 from meerschaum.utils.packages import attempt_import 36 polars = attempt_import('polars', install=False, lazy=False, warn=False) 37 if polars is None: 38 return None 39 if get_uuid_cols(new_df) or get_uuid_cols(old_df): 40 return None 41 42 row_col = '__mrsm_row_id' 43 while row_col in new_df.columns: 44 row_col = '_' + row_col 45 46 try: 47 new_pl = polars.from_pandas(new_df, include_index=False).with_row_index(row_col) 48 old_pl = polars.from_pandas(old_df, include_index=False) 49 row_indices = new_pl.join( 50 old_pl, 51 on=list(new_df.columns), 52 how='anti', 53 nulls_equal=True, 54 ).get_column(row_col).to_list() 55 except Exception: 56 return None 57 58 return new_df.iloc[sorted(row_indices)].reset_index(drop=True) 59 60 61def to_pandas(df: Any) -> Any: 62 """Convert a Polars DataFrame or LazyFrame to Pandas; otherwise return ``df``.""" 63 if df.__class__.__module__.split('.')[0] != 'polars': 64 return df 65 if df.__class__.__name__ == 'LazyFrame': 66 df = df.collect() 67 json_cols = [ 68 col 69 for col, typ in df.schema.items() 70 if getattr(typ, 'ext_name', lambda: None)() == 'arrow.json' 71 ] 72 if json_cols: 73 df = df.with_columns(df[col].ext.storage() for col in json_cols) 74 pandas_df = df.to_pandas(use_pyarrow_extension_array=True) 75 if json_cols: 76 import json 77 for col in json_cols: 78 pandas_df[col] = pandas_df[col].apply( 79 lambda value: json.loads(value) if isinstance(value, str) else None 80 ) 81 return pandas_df 82 83 84def to_polars( 85 df: Any, 86 geometry_cols_types_srids: Optional[Dict[str, Tuple[str, Any]]] = None, 87 json_cols: Optional[List[str]] = None, 88) -> Any: 89 """Convert a Pandas DataFrame to Polars; otherwise return ``df``.""" 90 if df.__class__.__module__.split('.')[0] == 'polars': 91 return df 92 polars = mrsm.attempt_import('polars') 93 geometry_cols_types_srids = ( 94 get_geometry_cols(df, with_types_srids=True) 95 if geometry_cols_types_srids is None and get_geometry_cols(df) 96 else (geometry_cols_types_srids or {}) 97 ) 98 geometry_cols_types_srids = { 99 col: type_srid 100 for col, type_srid in geometry_cols_types_srids.items() 101 if col in df.columns 102 } 103 json_cols = [col for col in (json_cols or []) if col in df.columns] 104 if geometry_cols_types_srids or json_cols: 105 try: 106 import json 107 geometry_cols = list(geometry_cols_types_srids) 108 special_cols = list(dict.fromkeys(geometry_cols + json_cols)) 109 polars_df = to_polars( 110 df.drop(columns=special_cols), 111 geometry_cols_types_srids={}, 112 json_cols=[], 113 ) 114 if geometry_cols: 115 from meerschaum.utils.dtypes import attempt_cast_to_geometry 116 shapely = mrsm.attempt_import('shapely', lazy=False) 117 try: 118 from geoarrow.types.type_pyarrow import register_extension_types 119 register_extension_types() 120 except Exception: 121 pass 122 for col, (_, srid) in geometry_cols_types_srids.items(): 123 crs = str(srid) if srid else None 124 if crs and ':' not in crs: 125 crs = 'EPSG:' + crs 126 metadata = json.dumps( 127 ({'crs': crs, 'crs_type': 'authority_code'} if crs else {}), 128 separators=(',', ':'), 129 ) 130 polars_df = polars_df.with_columns(polars.Series( 131 col, 132 shapely.to_wkb( 133 attempt_cast_to_geometry(df[col]), 134 hex=False, 135 include_srid=True, 136 ).tolist(), 137 dtype=polars.Extension('geoarrow.wkb', polars.Binary, metadata), 138 )) 139 if json_cols: 140 from meerschaum.utils.dtypes import json_serialize_value, value_is_null 141 for col in json_cols: 142 values = [] 143 for value in df[col].tolist(): 144 if value_is_null(value): 145 values.append(None) 146 continue 147 if isinstance(value, str): 148 try: 149 value = json.loads(value) 150 except json.JSONDecodeError: 151 pass 152 values.append(json.dumps( 153 value, 154 default=json_serialize_value, 155 separators=(',', ':'), 156 allow_nan=False, 157 )) 158 polars_df = polars_df.with_columns(polars.Series( 159 col, 160 values, 161 dtype=polars.Extension('arrow.json', polars.String, ''), 162 )) 163 return polars_df.select(list(df.columns)) 164 except Exception: 165 # ponytail: Preserve compatibility if Polars changes its unstable extension API. 166 pass 167 if get_uuid_cols(df): 168 return polars.DataFrame(df.to_dict(orient='list'), strict=False) 169 try: 170 return polars.from_pandas(df, include_index=False) 171 except (TypeError, ValueError): 172 # ponytail: Preserve unsupported Python objects; remove when Polars supports UUIDs. 173 return polars.DataFrame(df.to_dict(orient='list'), strict=False) 174 175 176def add_missing_cols_to_df( 177 df: 'pd.DataFrame', 178 dtypes: Dict[str, Any], 179) -> 'pd.DataFrame': 180 """ 181 Add columns from the dtypes dictionary as null columns to a new DataFrame. 182 183 Parameters 184 ---------- 185 df: pd.DataFrame 186 The dataframe we should copy and add null columns. 187 188 dtypes: 189 The data types dictionary which may contain keys not present in `df.columns`. 190 191 Returns 192 ------- 193 A new `DataFrame` with the keys from `dtypes` added as null columns. 194 If `df.dtypes` is the same as `dtypes`, then return a reference to `df`. 195 NOTE: This will not ensure that dtypes are enforced! 196 197 Examples 198 -------- 199 >>> import pandas as pd 200 >>> df = pd.DataFrame([{'a': 1}]) 201 >>> dtypes = {'b': 'Int64'} 202 >>> add_missing_cols_to_df(df, dtypes) 203 a b 204 0 1 <NA> 205 >>> add_missing_cols_to_df(df, dtypes).dtypes 206 a int64 207 b Int64 208 dtype: object 209 >>> add_missing_cols_to_df(df, {'a': 'object'}).dtypes 210 a int64 211 dtype: object 212 >>> 213 """ 214 if set(df.columns) == set(dtypes): 215 return df 216 217 from meerschaum.utils.packages import attempt_import 218 from meerschaum.utils.dtypes import to_pandas_dtype 219 pandas = attempt_import('pandas') 220 221 def build_series(dtype: str): 222 return pandas.Series([], dtype=to_pandas_dtype(dtype)) 223 224 assign_kwargs = { 225 str(col): build_series(str(typ)) 226 for col, typ in dtypes.items() 227 if col not in df.columns 228 } 229 df_with_cols = df.assign(**assign_kwargs) 230 for col in assign_kwargs: 231 df_with_cols[col] = df_with_cols[col].fillna(pandas.NA) 232 return df_with_cols 233 234 235def filter_unseen_df( 236 old_df: 'pd.DataFrame', 237 new_df: 'pd.DataFrame', 238 safe_copy: bool = True, 239 dtypes: Optional[Dict[str, Any]] = None, 240 include_unchanged_columns: bool = False, 241 coerce_mixed_numerics: bool = True, 242 debug: bool = False, 243) -> 'pd.DataFrame': 244 """ 245 Left join two DataFrames to find the newest unseen data. 246 247 Parameters 248 ---------- 249 old_df: 'pd.DataFrame' 250 The original (target) dataframe. Acts as a filter on the `new_df`. 251 252 new_df: 'pd.DataFrame' 253 The fetched (source) dataframe. Rows that are contained in `old_df` are removed. 254 255 safe_copy: bool, default True 256 If `True`, create a copy before comparing and modifying the dataframes. 257 Setting to `False` may mutate the DataFrames. 258 259 dtypes: Optional[Dict[str, Any]], default None 260 Optionally specify the datatypes of the dataframe. 261 262 include_unchanged_columns: bool, default False 263 If `True`, include columns which haven't changed on rows which have changed. 264 265 coerce_mixed_numerics: bool, default True 266 If `True`, cast mixed integer and float columns between the old and new dataframes into 267 numeric values (`decimal.Decimal`). 268 269 debug: bool, default False 270 Verbosity toggle. 271 272 Returns 273 ------- 274 A pandas dataframe of the new, unseen rows in `new_df`. 275 276 Examples 277 -------- 278 ```python 279 >>> import pandas as pd 280 >>> df1 = pd.DataFrame({'a': [1,2]}) 281 >>> df2 = pd.DataFrame({'a': [2,3]}) 282 >>> filter_unseen_df(df1, df2) 283 a 284 0 3 285 286 ``` 287 288 """ 289 if old_df is None: 290 return new_df 291 292 if safe_copy: 293 old_df = old_df.copy() 294 new_df = new_df.copy() 295 296 import json 297 import functools 298 import traceback 299 from meerschaum.utils.warnings import warn 300 from meerschaum.utils.packages import import_pandas, attempt_import 301 from meerschaum.utils.dtypes import ( 302 to_pandas_dtype, 303 are_dtypes_equal, 304 attempt_cast_to_numeric, 305 attempt_cast_to_uuid, 306 attempt_cast_to_bytes, 307 attempt_cast_to_geometry, 308 coerce_timezone, 309 serialize_decimal, 310 ) 311 from meerschaum.utils.dtypes.sql import get_numeric_precision_scale 312 pd = import_pandas(debug=debug) 313 is_dask = 'dask' in new_df.__module__ 314 if is_dask: 315 pandas = attempt_import('pandas') 316 _ = attempt_import('partd', lazy=False) 317 dd = attempt_import('dask.dataframe') 318 merge = dd.merge 319 NA = pandas.NA 320 else: 321 merge = pd.merge 322 NA = pd.NA 323 324 new_df_dtypes = dict(new_df.dtypes) 325 new_cols = list(new_df_dtypes) 326 old_df_dtypes = dict(old_df.dtypes) 327 328 same_cols = set(new_df.columns) == set(old_df.columns) 329 if not same_cols: 330 new_df = add_missing_cols_to_df(new_df, old_df_dtypes) 331 old_df = add_missing_cols_to_df(old_df, new_df_dtypes) 332 333 new_types_missing_from_old = { 334 col: typ 335 for col, typ in new_df_dtypes.items() 336 if col not in old_df_dtypes 337 } 338 old_types_missing_from_new = { 339 col: typ 340 for col, typ in old_df_dtypes.items() 341 if col not in new_df_dtypes 342 } 343 old_df_dtypes.update(new_types_missing_from_old) 344 new_df_dtypes.update(old_types_missing_from_new) 345 346 ### Edge case: two empty lists cast to DFs. 347 elif len(new_df.columns) == 0: 348 return new_df 349 350 try: 351 ### Order matters when checking equality. 352 new_df = new_df[old_df.columns] 353 354 except Exception as e: 355 warn( 356 "Was not able to cast old columns onto new DataFrame. " + 357 f"Are both DataFrames the same shape? Error:\n{e}", 358 stacklevel=3, 359 ) 360 return new_df[list(new_df_dtypes.keys())] 361 362 ### assume the old_df knows what it's doing, even if it's technically wrong. 363 if dtypes is None: 364 dtypes = {col: str(typ) for col, typ in old_df.dtypes.items()} 365 366 numeric_cols_precisions_scales = { 367 col: get_numeric_precision_scale(None, typ) 368 for col, typ in dtypes.items() 369 if col and str(typ).lower().startswith('numeric') 370 } 371 dtypes = { 372 col: to_pandas_dtype(typ) 373 for col, typ in dtypes.items() 374 if col in new_df_dtypes and col in old_df_dtypes 375 } 376 for col, typ in new_df_dtypes.items(): 377 if col not in dtypes: 378 dtypes[col] = typ 379 380 dt_dtypes = { 381 col: typ 382 for col, typ in dtypes.items() 383 if are_dtypes_equal(typ, 'datetime') 384 } 385 non_dt_dtypes = { 386 col: typ 387 for col, typ in dtypes.items() 388 if col not in dt_dtypes 389 } 390 391 cast_non_dt_cols = True 392 try: 393 new_df = new_df.astype(non_dt_dtypes) 394 cast_non_dt_cols = False 395 except Exception as e: 396 warn( 397 f"Was not able to cast the new DataFrame to the given dtypes.\n{e}" 398 ) 399 400 cast_dt_cols = True 401 try: 402 for col, typ in dt_dtypes.items(): 403 _dtypes_col_dtype = str((dtypes or {}).get(col, 'datetime')) 404 strip_utc = ( 405 _dtypes_col_dtype.startswith('datetime64') 406 and 'utc' not in _dtypes_col_dtype.lower() 407 ) 408 if col in old_df.columns: 409 old_df[col] = coerce_timezone( 410 old_df[col], strip_utc=strip_utc 411 ).astype(typ) 412 if col in new_df.columns: 413 new_df[col] = coerce_timezone( 414 new_df[col], strip_utc=strip_utc 415 ).astype(typ) 416 cast_dt_cols = False 417 except Exception as e: 418 warn(f"Could not cast datetime columns:\n{e}") 419 420 cast_cols = cast_dt_cols or cast_non_dt_cols 421 422 new_numeric_cols_existing = get_numeric_cols(new_df) 423 old_numeric_cols = get_numeric_cols(old_df) 424 for col, typ in {k: v for k, v in dtypes.items()}.items(): 425 if not are_dtypes_equal(new_df_dtypes.get(col, 'None'), old_df_dtypes.get(col, 'None')): 426 new_is_float = are_dtypes_equal(new_df_dtypes.get(col, 'None'), 'float') 427 new_is_int = are_dtypes_equal(new_df_dtypes.get(col, 'None'), 'int') 428 new_is_numeric = col in new_numeric_cols_existing 429 old_is_float = are_dtypes_equal(old_df_dtypes.get(col, 'None'), 'float') 430 old_is_int = are_dtypes_equal(old_df_dtypes.get(col, 'None'), 'int') 431 old_is_numeric = col in old_numeric_cols 432 433 if ( 434 coerce_mixed_numerics 435 and 436 (new_is_float or new_is_int or new_is_numeric) 437 and 438 (old_is_float or old_is_int or old_is_numeric) 439 ): 440 dtypes[col] = attempt_cast_to_numeric 441 cast_cols = True 442 continue 443 444 ### Fallback to object if the types don't match. 445 warn( 446 f"Detected different types for '{col}' " 447 + f"({new_df_dtypes.get(col, None)} vs {old_df_dtypes.get(col, None)}), " 448 + "falling back to 'object'..." 449 ) 450 dtypes[col] = 'object' 451 cast_cols = True 452 453 if cast_cols: 454 for col, dtype in dtypes.items(): 455 for df_to_cast in (new_df, old_df): 456 if col not in df_to_cast.columns: 457 continue 458 try: 459 df_to_cast[col] = ( 460 df_to_cast[col].astype(dtype) 461 if not callable(dtype) 462 else df_to_cast[col].apply(dtype) 463 ) 464 except Exception as e: 465 warn(f"Was not able to cast column '{col}' to dtype '{dtype}'.\n{e}") 466 467 serializer = functools.partial(json.dumps, sort_keys=True, separators=(',', ':'), default=str) 468 new_json_cols = get_json_cols(new_df) 469 old_json_cols = get_json_cols(old_df) 470 json_cols = set(new_json_cols + old_json_cols) 471 for json_col in old_json_cols: 472 old_df[json_col] = old_df[json_col].apply(serializer) 473 for json_col in new_json_cols: 474 new_df[json_col] = new_df[json_col].apply(serializer) 475 476 new_numeric_cols = get_numeric_cols(new_df) 477 numeric_cols = set(new_numeric_cols + old_numeric_cols) 478 for numeric_col in old_numeric_cols: 479 old_df[numeric_col] = old_df[numeric_col].apply(serialize_decimal) 480 for numeric_col in new_numeric_cols: 481 new_df[numeric_col] = new_df[numeric_col].apply(serialize_decimal) 482 483 old_dt_cols = [ 484 col 485 for col, typ in old_df.dtypes.items() 486 if are_dtypes_equal(str(typ), 'datetime') 487 ] 488 for col in old_dt_cols: 489 _dtypes_col_dtype = str((dtypes or {}).get(col, 'datetime')) 490 strip_utc = ( 491 _dtypes_col_dtype.startswith('datetime64') 492 and 'utc' not in _dtypes_col_dtype.lower() 493 ) 494 old_df[col] = coerce_timezone(old_df[col], strip_utc=strip_utc) 495 496 new_dt_cols = [ 497 col 498 for col, typ in new_df.dtypes.items() 499 if are_dtypes_equal(str(typ), 'datetime') 500 ] 501 for col in new_dt_cols: 502 _dtypes_col_dtype = str((dtypes or {}).get(col, 'datetime')) 503 strip_utc = ( 504 _dtypes_col_dtype.startswith('datetime64') 505 and 'utc' not in _dtypes_col_dtype.lower() 506 ) 507 new_df[col] = coerce_timezone(new_df[col], strip_utc=strip_utc) 508 509 old_uuid_cols = get_uuid_cols(old_df) 510 new_uuid_cols = get_uuid_cols(new_df) 511 uuid_cols = set(new_uuid_cols + old_uuid_cols) 512 513 old_bytes_cols = get_bytes_cols(old_df) 514 new_bytes_cols = get_bytes_cols(new_df) 515 bytes_cols = set(new_bytes_cols + old_bytes_cols) 516 517 old_geometry_cols = get_geometry_cols(old_df) 518 new_geometry_cols = get_geometry_cols(new_df) 519 geometry_cols = set(new_geometry_cols + old_geometry_cols) 520 521 na_pattern = r'(?i)^(none|nan|na|nat|natz|<na>)$' 522 def normalize_nulls(df_to_normalize): 523 normalized_df = df_to_normalize.infer_objects() 524 string_cols = normalized_df.select_dtypes( 525 include=['object', 'string', 'category'] 526 ).columns 527 return normalized_df.replace( 528 {col: na_pattern for col in string_cols}, 529 pd.NA, 530 regex=True, 531 ).fillna(NA) 532 533 normalized_new_df = normalize_nulls(new_df) 534 normalized_old_df = normalize_nulls(old_df) 535 delta_df = ( 536 None 537 if is_dask 538 else _filter_unseen_df_with_polars(normalized_new_df, normalized_old_df) 539 ) 540 if delta_df is None: 541 joined_df = merge( 542 normalized_new_df, 543 normalized_old_df, 544 how='left', 545 on=None, 546 indicator=True, 547 ) 548 changed_rows_mask = (joined_df['_merge'] == 'left_only') 549 delta_df = joined_df[new_cols][changed_rows_mask].reset_index(drop=True) 550 else: 551 merge_dtypes = merge( 552 normalized_new_df.head(0), 553 normalized_old_df.head(0), 554 how='left', 555 on=None, 556 indicator=True, 557 ).dtypes 558 delta_df = delta_df.astype({ 559 col: typ 560 for col, typ in merge_dtypes.items() 561 if col in delta_df.columns and str(delta_df.dtypes[col]) != str(typ) 562 }) 563 delta_df = delta_df[new_cols] 564 565 delta_json_cols = get_json_cols(delta_df) 566 for json_col in json_cols: 567 if ( 568 json_col in delta_json_cols 569 or json_col not in delta_df.columns 570 ): 571 continue 572 try: 573 delta_df[json_col] = delta_df[json_col].apply( 574 lambda x: (json.loads(x) if isinstance(x, str) else x) 575 ) 576 except Exception: 577 warn(f"Unable to deserialize JSON column '{json_col}':\n{traceback.format_exc()}") 578 579 delta_numeric_cols = get_numeric_cols(delta_df) 580 for numeric_col in numeric_cols: 581 if ( 582 numeric_col in delta_numeric_cols 583 or numeric_col not in delta_df.columns 584 ): 585 continue 586 try: 587 delta_df[numeric_col] = delta_df[numeric_col].apply( 588 functools.partial( 589 attempt_cast_to_numeric, 590 quantize=True, 591 precision=numeric_cols_precisions_scales.get(numeric_col, (None, None))[0], 592 scale=numeric_cols_precisions_scales.get(numeric_col, (None, None))[1], 593 ) 594 ) 595 except Exception: 596 warn(f"Unable to parse numeric column '{numeric_col}':\n{traceback.format_exc()}") 597 598 delta_uuid_cols = get_uuid_cols(delta_df) 599 for uuid_col in uuid_cols: 600 if ( 601 uuid_col in delta_uuid_cols 602 or uuid_col not in delta_df.columns 603 ): 604 continue 605 try: 606 delta_df[uuid_col] = delta_df[uuid_col].apply(attempt_cast_to_uuid) 607 except Exception: 608 warn(f"Unable to parse numeric column '{uuid_col}':\n{traceback.format_exc()}") 609 610 delta_bytes_cols = get_bytes_cols(delta_df) 611 for bytes_col in bytes_cols: 612 if ( 613 bytes_col in delta_bytes_cols 614 or bytes_col not in delta_df.columns 615 ): 616 continue 617 try: 618 delta_df[bytes_col] = delta_df[bytes_col].apply(attempt_cast_to_bytes) 619 except Exception: 620 warn(f"Unable to parse bytes column '{bytes_col}':\n{traceback.format_exc()}") 621 622 delta_geometry_cols = get_geometry_cols(delta_df) 623 for geometry_col in geometry_cols: 624 if ( 625 geometry_col in delta_geometry_cols 626 or geometry_col not in delta_df.columns 627 ): 628 continue 629 try: 630 delta_df[geometry_col] = attempt_cast_to_geometry(delta_df[geometry_col]) 631 except Exception: 632 warn(f"Unable to parse geometry column '{geometry_col}':\n{traceback.format_exc()}") 633 634 return delta_df 635 636 637def parse_df_datetimes( 638 df: 'pd.DataFrame', 639 ignore_cols: Optional[Iterable[str]] = None, 640 strip_timezone: bool = False, 641 chunksize: Optional[int] = None, 642 dtype_backend: str = 'numpy_nullable', 643 ignore_all: bool = False, 644 precision_unit: Optional[str] = None, 645 coerce_utc: bool = True, 646 debug: bool = False, 647) -> 'pd.DataFrame': 648 """ 649 Parse a pandas DataFrame for datetime columns and cast as datetimes. 650 651 Parameters 652 ---------- 653 df: pd.DataFrame 654 The pandas DataFrame to parse. 655 656 ignore_cols: Optional[Iterable[str]], default None 657 If provided, do not attempt to coerce these columns as datetimes. 658 659 strip_timezone: bool, default False 660 If `True`, remove the UTC `tzinfo` property. 661 662 chunksize: Optional[int], default None 663 If the pandas implementation is `'dask'`, use this chunksize for the distributed dataframe. 664 665 dtype_backend: str, default 'numpy_nullable' 666 If `df` is not a DataFrame and new one needs to be constructed, 667 use this as the datatypes backend. 668 Accepted values are 'numpy_nullable' and 'pyarrow'. 669 670 ignore_all: bool, default False 671 If `True`, do not attempt to cast any columns to datetimes. 672 673 precision_unit: Optional[str], default None 674 If provided, enforce the given precision on the coerced datetime columns. 675 676 coerce_utc: bool, default True 677 Coerce the datetime columns to UTC (see `meerschaum.utils.dtypes.to_datetime()`). 678 679 debug: bool, default False 680 Verbosity toggle. 681 682 Returns 683 ------- 684 A new pandas DataFrame with the determined datetime columns 685 (usually ISO strings) cast as datetimes. 686 687 Examples 688 -------- 689 ```python 690 >>> import pandas as pd 691 >>> df = pd.DataFrame({'a': ['2022-01-01 00:00:00']}) 692 >>> df.dtypes 693 a object 694 dtype: object 695 >>> df2 = parse_df_datetimes(df) 696 >>> df2.dtypes 697 a datetime64[us, UTC] 698 dtype: object 699 700 ``` 701 702 """ 703 from meerschaum.utils.packages import import_pandas, attempt_import 704 from meerschaum.utils.debug import dprint 705 from meerschaum.utils.warnings import warn 706 from meerschaum.utils.misc import items_str 707 from meerschaum.utils.dtypes import to_datetime, MRSM_PD_DTYPES 708 import traceback 709 710 pd = import_pandas() 711 pandas = attempt_import('pandas') 712 pd_name = pd.__name__ 713 using_dask = 'dask' in pd_name 714 df_is_dask = (hasattr(df, '__module__') and 'dask' in df.__module__) 715 dask_dataframe = None 716 if using_dask or df_is_dask: 717 npartitions = chunksize_to_npartitions(chunksize) 718 dask_dataframe = attempt_import('dask.dataframe') 719 720 ### if df is a dict, build DataFrame 721 if isinstance(df, pandas.DataFrame): 722 pdf = df 723 elif df_is_dask and isinstance(df, dask_dataframe.DataFrame): 724 pdf = get_first_valid_dask_partition(df) 725 else: 726 if debug: 727 dprint(f"df is of type '{type(df)}'. Building {pd.DataFrame}...") 728 729 if using_dask: 730 if isinstance(df, list): 731 keys = set() 732 for doc in df: 733 for key in doc: 734 keys.add(key) 735 df = pd.DataFrame.from_dict( 736 { 737 k: [ 738 doc.get(k, None) 739 for doc in df 740 ] for k in keys 741 }, 742 npartitions=npartitions, 743 ) 744 elif isinstance(df, dict): 745 df = pd.DataFrame.from_dict(df, npartitions=npartitions) 746 elif 'pandas.core.frame.DataFrame' in str(type(df)): 747 df = pd.from_pandas(df, npartitions=npartitions) 748 else: 749 raise Exception("Can only parse dictionaries or lists of dictionaries with Dask.") 750 pandas = attempt_import('pandas') 751 pdf = get_first_valid_dask_partition(df) 752 753 else: 754 df = pd.DataFrame(df).convert_dtypes(dtype_backend=dtype_backend) 755 pdf = df 756 757 ### skip parsing if DataFrame is empty 758 if len(pdf) == 0: 759 if debug: 760 dprint("df is empty. Returning original DataFrame without casting datetime columns...") 761 return df 762 763 ignore_cols = set( 764 (ignore_cols or []) + [ 765 col 766 for col, dtype in pdf.dtypes.items() 767 if 'datetime' in str(dtype) 768 ] 769 ) 770 cols_to_inspect = [ 771 col 772 for col in pdf.columns 773 if col not in ignore_cols 774 ] if not ignore_all else [] 775 776 if len(cols_to_inspect) == 0: 777 if debug: 778 dprint("All columns are ignored, skipping datetime detection...") 779 return df.infer_objects().fillna(pandas.NA) 780 781 ### apply regex to columns to determine which are ISO datetimes 782 iso_dt_regex = r'\d{4}-\d{2}-\d{2}.\d{2}\:\d{2}\:\d+' 783 dt_mask = pdf[cols_to_inspect].astype(str).apply( 784 lambda s: s.str.match(iso_dt_regex).all() 785 ) 786 787 ### list of datetime column names 788 datetime_cols = [col for col in pdf[cols_to_inspect].loc[:, dt_mask]] 789 if not datetime_cols: 790 if debug: 791 dprint("No columns detected as datetimes, returning...") 792 return df.infer_objects().fillna(pandas.NA) 793 794 if debug: 795 dprint("Converting columns to datetimes: " + str(datetime_cols)) 796 797 def _parse_to_datetime(x): 798 return to_datetime(x, precision_unit=precision_unit, coerce_utc=coerce_utc) 799 800 try: 801 if not using_dask: 802 df[datetime_cols] = df[datetime_cols].apply(_parse_to_datetime) 803 else: 804 df[datetime_cols] = df[datetime_cols].apply( 805 _parse_to_datetime, 806 utc=True, 807 axis=1, 808 meta={ 809 col: MRSM_PD_DTYPES['datetime'] 810 for col in datetime_cols 811 } 812 ) 813 except Exception: 814 warn( 815 f"Unable to apply `to_datetime()` to {items_str(datetime_cols)}:\n" 816 + f"{traceback.format_exc()}" 817 ) 818 819 if strip_timezone: 820 for dt in datetime_cols: 821 try: 822 df[dt] = df[dt].dt.tz_localize(None) 823 except Exception: 824 warn( 825 f"Unable to convert column '{dt}' to naive datetime:\n" 826 + f"{traceback.format_exc()}" 827 ) 828 829 return df.fillna(pandas.NA) 830 831 832def get_unhashable_cols(df: 'pd.DataFrame') -> List[str]: 833 """ 834 Get the columns which contain unhashable objects from a Pandas DataFrame. 835 836 Parameters 837 ---------- 838 df: pd.DataFrame 839 The DataFrame which may contain unhashable objects. 840 841 Returns 842 ------- 843 A list of columns. 844 """ 845 if df is None: 846 return [] 847 if len(df) == 0: 848 return [] 849 850 is_dask = 'dask' in df.__module__ 851 if is_dask: 852 from meerschaum.utils.packages import attempt_import 853 pandas = attempt_import('pandas') 854 df = pandas.DataFrame(get_first_valid_dask_partition(df)) 855 return [ 856 col for col, val in df.iloc[0].items() 857 if not isinstance(val, Hashable) 858 ] 859 860 861def get_json_cols(df: 'pd.DataFrame') -> List[str]: 862 """ 863 Get the columns which contain unhashable objects from a Pandas DataFrame. 864 865 Parameters 866 ---------- 867 df: pd.DataFrame 868 The DataFrame which may contain unhashable objects. 869 870 Returns 871 ------- 872 A list of columns to be encoded as JSON. 873 """ 874 if df is None: 875 return [] 876 877 is_dask = 'dask' in df.__module__ if hasattr(df, '__module__') else False 878 if is_dask: 879 df = get_first_valid_dask_partition(df) 880 881 if len(df) == 0: 882 return [] 883 884 cols_indices = { 885 col: df[col].first_valid_index() 886 for col in df.columns 887 } 888 return [ 889 col 890 for col, ix in cols_indices.items() 891 if ( 892 ix is not None 893 and isinstance(df.loc[ix][col], (dict, list)) 894 ) 895 ] 896 897 898def get_numeric_cols(df: 'pd.DataFrame') -> List[str]: 899 """ 900 Get the columns which contain `decimal.Decimal` objects from a Pandas DataFrame. 901 902 Parameters 903 ---------- 904 df: pd.DataFrame 905 The DataFrame which may contain decimal objects. 906 907 Returns 908 ------- 909 A list of columns to treat as numerics. 910 """ 911 if df is None: 912 return [] 913 from decimal import Decimal 914 is_dask = 'dask' in df.__module__ 915 if is_dask: 916 df = get_first_valid_dask_partition(df) 917 918 if len(df) == 0: 919 return [] 920 921 cols_indices = { 922 col: df[col].first_valid_index() 923 for col in df.columns 924 } 925 return [ 926 col 927 for col, ix in cols_indices.items() 928 if ( 929 ix is not None 930 and 931 isinstance(df.loc[ix][col], Decimal) 932 ) 933 ] 934 935 936def get_bool_cols(df: 'pd.DataFrame') -> List[str]: 937 """ 938 Get the columns which contain `bool` objects from a Pandas DataFrame. 939 940 Parameters 941 ---------- 942 df: pd.DataFrame 943 The DataFrame which may contain bools. 944 945 Returns 946 ------- 947 A list of columns to treat as bools. 948 """ 949 if df is None: 950 return [] 951 952 is_dask = 'dask' in df.__module__ 953 if is_dask: 954 df = get_first_valid_dask_partition(df) 955 956 if len(df) == 0: 957 return [] 958 959 from meerschaum.utils.dtypes import are_dtypes_equal 960 961 return [ 962 col 963 for col, typ in df.dtypes.items() 964 if are_dtypes_equal(str(typ), 'bool') 965 ] 966 967 968def get_uuid_cols(df: 'pd.DataFrame') -> List[str]: 969 """ 970 Get the columns which contain `uuid.UUID` objects from a Pandas DataFrame. 971 972 Parameters 973 ---------- 974 df: pd.DataFrame 975 The DataFrame which may contain UUID objects. 976 977 Returns 978 ------- 979 A list of columns to treat as UUIDs. 980 """ 981 if df is None: 982 return [] 983 from uuid import UUID 984 is_dask = 'dask' in df.__module__ 985 if is_dask: 986 df = get_first_valid_dask_partition(df) 987 988 if len(df) == 0: 989 return [] 990 991 cols_indices = { 992 col: df[col].first_valid_index() 993 for col in df.columns 994 } 995 return [ 996 col 997 for col, ix in cols_indices.items() 998 if ( 999 ix is not None 1000 and 1001 isinstance(df.loc[ix][col], UUID) 1002 ) 1003 ] 1004 1005 1006def get_datetime_cols( 1007 df: 'pd.DataFrame', 1008 timezone_aware: bool = True, 1009 timezone_naive: bool = True, 1010 with_tz_precision: bool = False, 1011) -> Union[List[str], Dict[str, Tuple[Union[str, None], str]]]: 1012 """ 1013 Get the columns which contain `datetime` or `Timestamp` objects from a Pandas DataFrame. 1014 1015 Parameters 1016 ---------- 1017 df: pd.DataFrame 1018 The DataFrame which may contain `datetime` or `Timestamp` objects. 1019 1020 timezone_aware: bool, default True 1021 If `True`, include timezone-aware datetime columns. 1022 1023 timezone_naive: bool, default True 1024 If `True`, include timezone-naive datetime columns. 1025 1026 with_tz_precision: bool, default False 1027 If `True`, return a dictionary mapping column names to tuples in the form 1028 `(timezone, precision)`. 1029 1030 Returns 1031 ------- 1032 A list of columns to treat as datetimes, or a dictionary of columns to tz+precision tuples 1033 (if `with_tz_precision` is `True`). 1034 """ 1035 if not timezone_aware and not timezone_naive: 1036 raise ValueError("`timezone_aware` and `timezone_naive` cannot both be `False`.") 1037 1038 if df is None: 1039 return [] if not with_tz_precision else {} 1040 1041 from meerschaum.utils.dtypes import are_dtypes_equal, MRSM_PRECISION_UNITS_ALIASES 1042 is_dask = 'dask' in df.__module__ 1043 if is_dask: 1044 df = get_first_valid_dask_partition(df) 1045 1046 def get_tz_precision_from_dtype(dtype: str) -> Tuple[Union[str, None], str]: 1047 """ 1048 Extract the tz + precision tuple from a dtype string. 1049 """ 1050 dtype = str(dtype).removesuffix('[pyarrow]') 1051 meta_str = dtype.split('[', maxsplit=1)[-1].rstrip(']').replace(' ', '') 1052 tz = ( 1053 None 1054 if ',' not in meta_str 1055 else meta_str.split(',', maxsplit=1)[-1].removeprefix('tz=') 1056 ) 1057 precision_abbreviation = ( 1058 meta_str 1059 if ',' not in meta_str 1060 else meta_str.split(',')[0] 1061 ) 1062 precision = MRSM_PRECISION_UNITS_ALIASES[precision_abbreviation] 1063 return tz, precision 1064 1065 def get_tz_precision_from_datetime(dt: datetime) -> Tuple[Union[str, None], str]: 1066 """ 1067 Return the tz + precision tuple from a Python datetime object. 1068 """ 1069 return dt.tzname(), 'microsecond' 1070 1071 known_dt_cols_types = { 1072 col: str(typ) 1073 for col, typ in df.dtypes.items() 1074 if are_dtypes_equal('datetime', str(typ)) 1075 } 1076 1077 known_dt_cols_tuples = { 1078 col: get_tz_precision_from_dtype(typ) 1079 for col, typ in known_dt_cols_types.items() 1080 } 1081 1082 if len(df) == 0: 1083 return ( 1084 list(known_dt_cols_types) 1085 if not with_tz_precision 1086 else known_dt_cols_tuples 1087 ) 1088 1089 cols_indices = { 1090 col: df[col].first_valid_index() 1091 for col in df.columns 1092 if col not in known_dt_cols_types 1093 } 1094 pydt_cols_tuples = { 1095 col: get_tz_precision_from_datetime(sample_val) 1096 for col, ix in cols_indices.items() 1097 if ( 1098 ix is not None 1099 and 1100 isinstance((sample_val := df.loc[ix][col]), datetime) 1101 ) 1102 } 1103 1104 dt_cols_tuples = { 1105 **known_dt_cols_tuples, 1106 **pydt_cols_tuples 1107 } 1108 1109 all_dt_cols_tuples = { 1110 col: dt_cols_tuples[col] 1111 for col in df.columns 1112 if col in dt_cols_tuples 1113 } 1114 if timezone_aware and timezone_naive: 1115 return ( 1116 list(all_dt_cols_tuples) 1117 if not with_tz_precision 1118 else all_dt_cols_tuples 1119 ) 1120 1121 known_timezone_aware_dt_cols = [ 1122 col 1123 for col in known_dt_cols_types 1124 if getattr(df[col], 'tz', None) is not None 1125 ] 1126 timezone_aware_pydt_cols_tuples = { 1127 col: (tz, precision) 1128 for col, (tz, precision) in pydt_cols_tuples.items() 1129 if df.loc[cols_indices[col]][col].tzinfo is not None 1130 } 1131 timezone_aware_dt_cols_set = set( 1132 known_timezone_aware_dt_cols + list(timezone_aware_pydt_cols_tuples) 1133 ) 1134 timezone_aware_cols_tuples = { 1135 col: (tz, precision) 1136 for col, (tz, precision) in all_dt_cols_tuples.items() 1137 if col in timezone_aware_dt_cols_set 1138 } 1139 timezone_naive_cols_tuples = { 1140 col: (tz, precision) 1141 for col, (tz, precision) in all_dt_cols_tuples.items() 1142 if col not in timezone_aware_dt_cols_set 1143 } 1144 1145 if timezone_aware: 1146 return ( 1147 list(timezone_aware_cols_tuples) 1148 if not with_tz_precision 1149 else timezone_aware_cols_tuples 1150 ) 1151 1152 return ( 1153 list(timezone_naive_cols_tuples) 1154 if not with_tz_precision 1155 else timezone_naive_cols_tuples 1156 ) 1157 1158 1159def get_datetime_cols_types(df: 'pd.DataFrame') -> Dict[str, str]: 1160 """ 1161 Return a dictionary mapping datetime columns to specific types strings. 1162 1163 Parameters 1164 ---------- 1165 df: pd.DataFrame 1166 The DataFrame which may contain datetime columns. 1167 1168 Returns 1169 ------- 1170 A dictionary mapping the datetime columns' names to dtype strings 1171 (containing timezone and precision metadata). 1172 1173 Examples 1174 -------- 1175 >>> from datetime import datetime, timezone 1176 >>> import pandas as pd 1177 >>> df = pd.DataFrame({'dt_tz_aware': [datetime(2025, 1, 1, tzinfo=timezone.utc)]}) 1178 >>> get_datetime_cols_types(df) 1179 {'dt_tz_aware': 'datetime64[us, UTC]'} 1180 >>> df = pd.DataFrame({'distant_dt': [datetime(1, 1, 1)]}) 1181 >>> get_datetime_cols_types(df) 1182 {'distant_dt': 'datetime64[us]'} 1183 >>> df = pd.DataFrame({'dt_second': datetime(2025, 1, 1)}) 1184 >>> df['dt_second'] = df['dt_second'].astype('datetime64[s]') 1185 >>> get_datetime_cols_types(df) 1186 {'dt_second': 'datetime64[s]'} 1187 """ 1188 from meerschaum.utils.dtypes import MRSM_PRECISION_UNITS_ABBREVIATIONS 1189 dt_cols_tuples = get_datetime_cols(df, with_tz_precision=True) 1190 if not dt_cols_tuples: 1191 return {} 1192 1193 return { 1194 col: ( 1195 f"datetime64[{MRSM_PRECISION_UNITS_ABBREVIATIONS[precision]}]" 1196 if tz is None 1197 else f"datetime64[{MRSM_PRECISION_UNITS_ABBREVIATIONS[precision]}, {tz}]" 1198 ) 1199 for col, (tz, precision) in dt_cols_tuples.items() 1200 } 1201 1202 1203def get_date_cols(df: 'pd.DataFrame') -> List[str]: 1204 """ 1205 Get the `date` columns from a Pandas DataFrame. 1206 1207 Parameters 1208 ---------- 1209 df: pd.DataFrame 1210 The DataFrame which may contain dates. 1211 1212 Returns 1213 ------- 1214 A list of columns to treat as dates. 1215 """ 1216 from meerschaum.utils.dtypes import are_dtypes_equal 1217 if df is None: 1218 return [] 1219 1220 is_dask = 'dask' in df.__module__ 1221 if is_dask: 1222 df = get_first_valid_dask_partition(df) 1223 1224 known_date_cols = [ 1225 col 1226 for col, typ in df.dtypes.items() 1227 if are_dtypes_equal(typ, 'date') 1228 ] 1229 1230 if len(df) == 0: 1231 return known_date_cols 1232 1233 cols_indices = { 1234 col: df[col].first_valid_index() 1235 for col in df.columns 1236 if col not in known_date_cols 1237 } 1238 object_date_cols = [ 1239 col 1240 for col, ix in cols_indices.items() 1241 if ( 1242 ix is not None 1243 and isinstance(df.loc[ix][col], date) 1244 ) 1245 ] 1246 1247 all_date_cols = set(known_date_cols + object_date_cols) 1248 1249 return [ 1250 col 1251 for col in df.columns 1252 if col in all_date_cols 1253 ] 1254 1255 1256def get_bytes_cols(df: 'pd.DataFrame') -> List[str]: 1257 """ 1258 Get the columns which contain bytes strings from a Pandas DataFrame. 1259 1260 Parameters 1261 ---------- 1262 df: pd.DataFrame 1263 The DataFrame which may contain bytes strings. 1264 1265 Returns 1266 ------- 1267 A list of columns to treat as bytes. 1268 """ 1269 if df is None: 1270 return [] 1271 1272 is_dask = 'dask' in df.__module__ 1273 if is_dask: 1274 df = get_first_valid_dask_partition(df) 1275 1276 known_bytes_cols = [ 1277 col 1278 for col, typ in df.dtypes.items() 1279 if str(typ) in ('binary[pyarrow]', 'large_binary[pyarrow]') 1280 ] 1281 1282 if len(df) == 0: 1283 return known_bytes_cols 1284 1285 cols_indices = { 1286 col: df[col].first_valid_index() 1287 for col in df.columns 1288 if col not in known_bytes_cols 1289 } 1290 object_bytes_cols = [ 1291 col 1292 for col, ix in cols_indices.items() 1293 if ( 1294 ix is not None 1295 and isinstance(df.loc[ix][col], bytes) 1296 ) 1297 ] 1298 1299 all_bytes_cols = set(known_bytes_cols + object_bytes_cols) 1300 1301 return [ 1302 col 1303 for col in df.columns 1304 if col in all_bytes_cols 1305 ] 1306 1307 1308def get_geometry_cols( 1309 df: 'pd.DataFrame', 1310 with_types_srids: bool = False, 1311) -> Union[List[str], Dict[str, Any]]: 1312 """ 1313 Get the columns which contain shapely objects from a Pandas DataFrame. 1314 1315 Parameters 1316 ---------- 1317 df: pd.DataFrame 1318 The DataFrame which may contain bytes strings. 1319 1320 with_types_srids: bool, default False 1321 If `True`, return a dictionary mapping columns to geometry types and SRIDs. 1322 1323 Returns 1324 ------- 1325 A list of columns to treat as `geometry`. 1326 If `with_types_srids`, return a dictionary mapping columns to tuples in the form (type, SRID). 1327 """ 1328 if df is None: 1329 return [] if not with_types_srids else {} 1330 1331 is_dask = 'dask' in df.__module__ 1332 if is_dask: 1333 df = get_first_valid_dask_partition(df) 1334 1335 if len(df) == 0: 1336 return [] if not with_types_srids else {} 1337 1338 cols_indices = { 1339 col: df[col].first_valid_index() 1340 for col in df.columns 1341 } 1342 geo_cols = [ 1343 col 1344 for col, ix in cols_indices.items() 1345 if ( 1346 ix is not None 1347 and 1348 'shapely' in str(type(df.loc[ix][col])) 1349 ) 1350 ] 1351 if not with_types_srids: 1352 return geo_cols 1353 1354 gpd = mrsm.attempt_import('geopandas', lazy=False) 1355 geo_cols_types_srids = {} 1356 for col in geo_cols: 1357 try: 1358 sample_geo_series = gpd.GeoSeries(df[col], crs=None) 1359 geometry_types = { 1360 geom.geom_type 1361 for geom in sample_geo_series 1362 if hasattr(geom, 'geom_type') 1363 } 1364 geometry_has_z = any(getattr(geom, 'has_z', False) for geom in sample_geo_series) 1365 srid = ( 1366 ( 1367 sample_geo_series.crs.sub_crs_list[0].to_epsg() 1368 if sample_geo_series.crs.is_compound 1369 else sample_geo_series.crs.to_epsg() 1370 ) 1371 if sample_geo_series.crs 1372 else 0 1373 ) 1374 geometry_type = list(geometry_types)[0] if len(geometry_types) == 1 else 'geometry' 1375 if geometry_type != 'geometry' and geometry_has_z: 1376 geometry_type = geometry_type + 'Z' 1377 except Exception: 1378 srid = 0 1379 geometry_type = 'geometry' 1380 geo_cols_types_srids[col] = (geometry_type, srid) 1381 1382 return geo_cols_types_srids 1383 1384 1385def get_geometry_cols_types(df: 'pd.DataFrame') -> Dict[str, str]: 1386 """ 1387 Return a dtypes dictionary mapping columns to specific geometry types (type, srid). 1388 """ 1389 geometry_cols_types_srids = get_geometry_cols(df, with_types_srids=True) 1390 new_cols_types = {} 1391 for col, (geometry_type, srid) in geometry_cols_types_srids.items(): 1392 new_dtype = "geometry" 1393 modifier = "" 1394 if not srid and geometry_type.lower() == 'geometry': 1395 new_cols_types[col] = new_dtype 1396 continue 1397 1398 modifier = "[" 1399 if geometry_type.lower() != 'geometry': 1400 modifier += f"{geometry_type}" 1401 1402 if srid: 1403 if modifier != '[': 1404 modifier += ", " 1405 modifier += f"{srid}" 1406 modifier += "]" 1407 new_cols_types[col] = f"{new_dtype}{modifier}" 1408 return new_cols_types 1409 1410 1411def get_special_cols(df: 'pd.DataFrame') -> Dict[str, str]: 1412 """ 1413 Return a dtypes dictionary mapping special columns to their dtypes. 1414 """ 1415 return { 1416 **{col: 'json' for col in get_json_cols(df)}, 1417 **{col: 'uuid' for col in get_uuid_cols(df)}, 1418 **{col: 'bytes' for col in get_bytes_cols(df)}, 1419 **{col: 'bool' for col in get_bool_cols(df)}, 1420 **{col: 'numeric' for col in get_numeric_cols(df)}, 1421 **{col: 'date' for col in get_date_cols(df)}, 1422 **get_datetime_cols_types(df), 1423 **get_geometry_cols_types(df), 1424 } 1425 1426 1427def _enforce_dtypes_with_polars( 1428 df: Any, 1429 dtypes: Dict[str, str], 1430 strip_timezone: bool = False, 1431) -> Any: 1432 """Return a Polars frame with enforced Arrow-native dtypes, or ``None`` to fall back.""" 1433 if 'dask' in getattr(df, '__module__', ''): 1434 return None 1435 1436 from meerschaum.utils.packages import attempt_import 1437 from meerschaum.utils.dtypes import MRSM_ALIAS_DTYPES 1438 from meerschaum.utils.dtypes.sql import get_numeric_precision_scale 1439 polars = attempt_import('polars', install=False, lazy=False, warn=False) 1440 if polars is None: 1441 return None 1442 1443 normalized_dtypes = { 1444 col: MRSM_ALIAS_DTYPES.get(str(typ), str(typ)).lower() 1445 for col, typ in dtypes.items() 1446 } 1447 for typ in normalized_dtypes.values(): 1448 if typ in ('uuid', 'object') or typ.startswith(('geometry', 'geography')): 1449 return None 1450 if typ.startswith('numeric') and None in get_numeric_precision_scale(None, typ): 1451 return None 1452 1453 try: 1454 json_cols = [col for col, typ in normalized_dtypes.items() if typ == 'json'] 1455 polars_df = to_polars(df, json_cols=json_cols) 1456 if polars_df.__class__.__name__ == 'LazyFrame': 1457 polars_df = polars_df.collect() 1458 expressions = [] 1459 integer_dtypes = { 1460 'int': polars.Int64, 1461 'int8': polars.Int8, 1462 'int16': polars.Int16, 1463 'int32': polars.Int32, 1464 'int64': polars.Int64, 1465 'uint8': polars.UInt8, 1466 'uint16': polars.UInt16, 1467 'uint32': polars.UInt32, 1468 'uint64': polars.UInt64, 1469 } 1470 float_dtypes = { 1471 'float': polars.Float64, 1472 'float32': polars.Float32, 1473 'float64': polars.Float64, 1474 'double': polars.Float64, 1475 } 1476 for col, typ in normalized_dtypes.items(): 1477 if col not in polars_df.columns or typ == 'object': 1478 continue 1479 expr = polars.col(col) 1480 source_dtype = polars_df.schema[col] 1481 if typ in integer_dtypes: 1482 expr = expr.cast(integer_dtypes[typ], strict=True) 1483 elif typ in float_dtypes: 1484 expr = expr.cast(float_dtypes[typ], strict=True) 1485 elif typ in ('str', 'string'): 1486 expr = expr.cast(polars.String, strict=True) 1487 elif typ == 'bool': 1488 expr = ( 1489 expr.cast(polars.String).str.to_lowercase().replace_strict( 1490 {'true': True, 'false': False, '1': True, '0': False}, 1491 return_dtype=polars.Boolean, 1492 ) 1493 if source_dtype == polars.String 1494 else expr.cast(polars.Boolean, strict=True) 1495 ) 1496 elif typ == 'bytes': 1497 expr = ( 1498 expr.str.decode('base64', strict=True) 1499 if source_dtype == polars.String 1500 else expr.cast(polars.Binary, strict=True) 1501 ) 1502 elif typ == 'date': 1503 expr = ( 1504 expr.str.to_date(strict=True) 1505 if source_dtype == polars.String 1506 else expr.cast(polars.Date, strict=True) 1507 ) 1508 elif typ.startswith('datetime'): 1509 precision = typ.split('[', maxsplit=1)[-1].split(',', maxsplit=1)[0] 1510 precision = precision if precision in ('ns', 'us', 'ms') else 'us' 1511 timezone = None if strip_timezone else 'UTC' 1512 expr = ( 1513 expr.str.to_datetime( 1514 time_unit=precision, 1515 time_zone=timezone, 1516 strict=True, 1517 ) 1518 if source_dtype == polars.String 1519 else expr.cast(polars.Datetime(precision, timezone), strict=True) 1520 ) 1521 elif typ.startswith('numeric'): 1522 precision, scale = get_numeric_precision_scale(None, typ) 1523 expr = expr.cast(polars.Decimal(precision, scale), strict=True) 1524 elif typ == 'json': 1525 if getattr(source_dtype, 'ext_name', lambda: None)() != 'arrow.json': 1526 return None 1527 continue 1528 else: 1529 return None 1530 expressions.append(expr.alias(col)) 1531 return polars_df.with_columns(expressions) 1532 except Exception: 1533 return None 1534 1535 1536def enforce_dtypes( 1537 df: 'pd.DataFrame', 1538 dtypes: Dict[str, str], 1539 explicit_dtypes: Optional[Dict[str, str]] = None, 1540 safe_copy: bool = True, 1541 coerce_numeric: bool = False, 1542 coerce_timezone: bool = True, 1543 strip_timezone: bool = False, 1544 as_polars: bool = False, 1545 debug: bool = False, 1546) -> 'pd.DataFrame': 1547 """ 1548 Enforce the `dtypes` dictionary on a DataFrame. 1549 1550 Parameters 1551 ---------- 1552 df: pd.DataFrame 1553 The DataFrame on which to enforce dtypes. 1554 1555 dtypes: Dict[str, str] 1556 The data types to attempt to enforce on the DataFrame. 1557 1558 explicit_dtypes: Optional[Dict[str, str]], default None 1559 If provided, automatic dtype coersion will respect explicitly configured 1560 dtypes (`int`, `float`, `numeric`). 1561 1562 safe_copy: bool, default True 1563 If `True`, create a copy before comparing and modifying the dataframes. 1564 Setting to `False` may mutate the DataFrames. 1565 See `meerschaum.utils.dataframe.filter_unseen_df`. 1566 1567 coerce_numeric: bool, default False 1568 If `True`, convert float and int collisions to numeric. 1569 1570 coerce_timezone: bool, default True 1571 If `True`, convert datetimes to UTC. 1572 1573 strip_timezone: bool, default False 1574 If `coerce_timezone` and `strip_timezone` are `True`, 1575 remove timezone information from datetimes. 1576 1577 as_polars: bool, default False 1578 If `True`, return a Polars DataFrame. Supported Arrow-native dtypes are enforced 1579 with Polars; unsupported schemas use the established Pandas path before conversion. 1580 1581 debug: bool, default False 1582 Verbosity toggle. 1583 1584 Returns 1585 ------- 1586 The Pandas DataFrame with the types enforced. 1587 """ 1588 import json 1589 import functools 1590 from meerschaum.utils.debug import dprint 1591 from meerschaum.utils.formatting import pprint 1592 from meerschaum.utils.dtypes import ( 1593 MRSM_ALIAS_DTYPES, 1594 are_dtypes_equal, 1595 to_pandas_dtype, 1596 is_dtype_numeric, 1597 attempt_cast_to_numeric, 1598 attempt_cast_to_uuid, 1599 attempt_cast_to_bytes, 1600 attempt_cast_to_geometry, 1601 coerce_timezone as _coerce_timezone, 1602 get_geometry_type_srid, 1603 ) 1604 from meerschaum.utils.dtypes.sql import get_numeric_precision_scale 1605 pandas = mrsm.attempt_import('pandas') 1606 is_dask = 'dask' in df.__module__ 1607 normalized_dtypes = { 1608 col: MRSM_ALIAS_DTYPES.get(str(typ), str(typ)).lower() 1609 for col, typ in dtypes.items() 1610 } 1611 declared_json_cols = [col for col, typ in normalized_dtypes.items() if typ == 'json'] 1612 fallback_dtypes = { 1613 col: dtypes[col] 1614 for col, typ in normalized_dtypes.items() 1615 if ( 1616 typ in ('uuid', 'object') 1617 or typ.startswith(('geometry', 'geography')) 1618 or (typ.startswith('numeric') and None in get_numeric_precision_scale(None, typ)) 1619 ) 1620 } 1621 native_dtypes = { 1622 col: typ 1623 for col, typ in dtypes.items() 1624 if col not in fallback_dtypes 1625 } 1626 df_cols = df.collect_schema().names() if hasattr(df, 'collect_schema') else df.columns 1627 untyped_cols = [col for col in df_cols if col not in normalized_dtypes] 1628 polars_df = ( 1629 _enforce_dtypes_with_polars( 1630 df, 1631 native_dtypes, 1632 strip_timezone=(strip_timezone if coerce_timezone else False), 1633 ) 1634 if native_dtypes and not fallback_dtypes and not untyped_cols 1635 else None 1636 ) 1637 if polars_df is not None: 1638 return polars_df if as_polars else to_pandas(polars_df) 1639 1640 if native_dtypes and (fallback_dtypes or untyped_cols): 1641 df = to_pandas(df) 1642 native_cols = [col for col in native_dtypes if col in df.columns] 1643 native_df = _enforce_dtypes_with_polars( 1644 df[native_cols], 1645 native_dtypes, 1646 strip_timezone=(strip_timezone if coerce_timezone else False), 1647 ) if native_cols else None 1648 if native_df is not None: 1649 if safe_copy: 1650 df = df.copy() 1651 native_pd_df = to_pandas(native_df) 1652 for col in native_cols: 1653 df[col] = native_pd_df[col].array 1654 dtypes = fallback_dtypes 1655 safe_copy = False 1656 1657 df = to_pandas(df) 1658 if safe_copy: 1659 df = df.copy() 1660 if len(df.columns) == 0: 1661 if debug: 1662 dprint("Incoming DataFrame has no columns. Skipping enforcement...") 1663 return to_polars(df) if as_polars else df 1664 1665 explicit_dtypes = explicit_dtypes or {} 1666 pipe_pandas_dtypes = { 1667 col: to_pandas_dtype(typ) 1668 for col, typ in dtypes.items() 1669 } 1670 json_cols = [ 1671 col 1672 for col, typ in dtypes.items() 1673 if typ == 'json' 1674 ] 1675 numeric_cols = [ 1676 col 1677 for col, typ in dtypes.items() 1678 if typ.startswith('numeric') 1679 ] 1680 geometry_cols_types_srids = { 1681 col: get_geometry_type_srid(typ, default_srid=0) 1682 for col, typ in dtypes.items() 1683 if typ.startswith('geometry') or typ.startswith('geography') 1684 } 1685 uuid_cols = [ 1686 col 1687 for col, typ in dtypes.items() 1688 if typ == 'uuid' 1689 ] 1690 bytes_cols = [ 1691 col 1692 for col, typ in dtypes.items() 1693 if typ == 'bytes' 1694 ] 1695 datetime_cols = [ 1696 col 1697 for col, typ in dtypes.items() 1698 if are_dtypes_equal(typ, 'datetime') 1699 ] 1700 df_numeric_cols = get_numeric_cols(df) 1701 if debug: 1702 dprint("Desired data types:") 1703 pprint(dtypes) 1704 dprint("Data types for incoming DataFrame:") 1705 pprint({_col: str(_typ) for _col, _typ in df.dtypes.items()}) 1706 1707 if json_cols and len(df) > 0: 1708 if debug: 1709 dprint(f"Checking columns for JSON encoding: {json_cols}") 1710 for col in json_cols: 1711 if col in df.columns: 1712 try: 1713 df[col] = df[col].apply( 1714 ( 1715 lambda x: ( 1716 json.loads(x) 1717 if isinstance(x, str) 1718 else x 1719 ) 1720 ) 1721 ) 1722 except Exception as e: 1723 if debug: 1724 dprint(f"Unable to parse column '{col}' as JSON:\n{e}") 1725 1726 if numeric_cols: 1727 if debug: 1728 dprint(f"Checking for numerics: {numeric_cols}") 1729 for col in numeric_cols: 1730 precision, scale = get_numeric_precision_scale(None, dtypes.get(col, '')) 1731 if col in df.columns: 1732 try: 1733 df[col] = df[col].apply( 1734 functools.partial( 1735 attempt_cast_to_numeric, 1736 quantize=True, 1737 precision=precision, 1738 scale=scale, 1739 ) 1740 ) 1741 except Exception as e: 1742 if debug: 1743 dprint(f"Unable to parse column '{col}' as NUMERIC:\n{e}") 1744 1745 if uuid_cols: 1746 if debug: 1747 dprint(f"Checking for UUIDs: {uuid_cols}") 1748 for col in uuid_cols: 1749 if col in df.columns: 1750 try: 1751 df[col] = df[col].apply(attempt_cast_to_uuid) 1752 except Exception as e: 1753 if debug: 1754 dprint(f"Unable to parse column '{col}' as UUID:\n{e}") 1755 1756 if bytes_cols: 1757 if debug: 1758 dprint(f"Checking for bytes: {bytes_cols}") 1759 for col in bytes_cols: 1760 if col in df.columns: 1761 try: 1762 df[col] = df[col].apply(attempt_cast_to_bytes) 1763 except Exception as e: 1764 if debug: 1765 dprint(f"Unable to parse column '{col}' as bytes:\n{e}") 1766 1767 if datetime_cols and coerce_timezone: 1768 if debug: 1769 dprint(f"Checking for datetime conversion: {datetime_cols}") 1770 for col in datetime_cols: 1771 if col in df.columns: 1772 if not strip_timezone and 'utc' in str(df.dtypes[col]).lower(): 1773 if debug: 1774 dprint(f"Skip UTC coersion for column '{col}' ({str(df[col].dtype)}).") 1775 continue 1776 if strip_timezone and ',' not in str(df.dtypes[col]): 1777 if debug: 1778 dprint( 1779 f"Skip UTC coersion (stripped) for column '{col}' " 1780 f"({str(df[col].dtype)})." 1781 ) 1782 continue 1783 1784 if debug: 1785 dprint( 1786 f"Data type for column '{col}' before timezone coersion: " 1787 f"{str(df[col].dtype)}" 1788 ) 1789 1790 df[col] = _coerce_timezone(df[col], strip_utc=strip_timezone) 1791 if debug: 1792 dprint( 1793 f"Data type for column '{col}' after timezone coersion: " 1794 f"{str(df[col].dtype)}" 1795 ) 1796 1797 if geometry_cols_types_srids: 1798 geopandas = mrsm.attempt_import('geopandas') 1799 if debug: 1800 dprint(f"Checking for geometry: {list(geometry_cols_types_srids)}") 1801 parsed_geom_cols = [] 1802 for col in geometry_cols_types_srids: 1803 if col not in df.columns: 1804 continue 1805 try: 1806 df[col] = attempt_cast_to_geometry(df[col]) 1807 parsed_geom_cols.append(col) 1808 except Exception as e: 1809 import traceback 1810 traceback.print_exc() 1811 if debug: 1812 dprint(f"Unable to parse column '{col}' as geometry:\n{e}") 1813 1814 if parsed_geom_cols: 1815 if debug: 1816 dprint(f"Converting to GeoDataFrame (geometry column: '{parsed_geom_cols[0]}')...") 1817 try: 1818 _, default_srid = geometry_cols_types_srids[parsed_geom_cols[0]] 1819 df = geopandas.GeoDataFrame(df, geometry=parsed_geom_cols[0], crs=default_srid) 1820 for col, (_, srid) in geometry_cols_types_srids.items(): 1821 if srid: 1822 if debug: 1823 dprint(f"Setting '{col}' to SRID '{srid}'...") 1824 _ = df[col].set_crs(srid) 1825 if parsed_geom_cols[0] not in df.columns: 1826 df.rename_geometry(parsed_geom_cols[0], inplace=True) 1827 except (ValueError, TypeError): 1828 import traceback 1829 dprint(f"Failed to cast to GeoDataFrame:\n{traceback.format_exc()}") 1830 1831 df_dtypes = {c: str(t) for c, t in df.dtypes.items()} 1832 if are_dtypes_equal(df_dtypes, pipe_pandas_dtypes): 1833 if debug: 1834 dprint("Data types match. Exiting enforcement...") 1835 return ( 1836 to_polars( 1837 df, 1838 geometry_cols_types_srids=geometry_cols_types_srids, 1839 json_cols=declared_json_cols, 1840 ) 1841 if as_polars 1842 else df 1843 ) 1844 1845 common_dtypes = {} 1846 common_diff_dtypes = {} 1847 for col, typ in pipe_pandas_dtypes.items(): 1848 if col in df_dtypes: 1849 common_dtypes[col] = typ 1850 if not are_dtypes_equal(typ, df_dtypes[col]): 1851 common_diff_dtypes[col] = df_dtypes[col] 1852 1853 if debug: 1854 dprint("Common columns with different dtypes:") 1855 pprint(common_diff_dtypes) 1856 1857 detected_dt_cols = {} 1858 for col, typ in common_diff_dtypes.items(): 1859 if 'datetime' in typ and 'datetime' in common_dtypes[col]: 1860 df_dtypes[col] = typ 1861 detected_dt_cols[col] = (common_dtypes[col], common_diff_dtypes[col]) 1862 for col in detected_dt_cols: 1863 del common_diff_dtypes[col] 1864 1865 if debug: 1866 dprint("Common columns with different dtypes (after dates):") 1867 pprint(common_diff_dtypes) 1868 1869 if are_dtypes_equal(df_dtypes, pipe_pandas_dtypes): 1870 if debug: 1871 dprint( 1872 "The incoming DataFrame has mostly the same types, skipping enforcement." 1873 + "The only detected difference was in the following datetime columns." 1874 ) 1875 pprint(detected_dt_cols) 1876 return ( 1877 to_polars( 1878 df, 1879 geometry_cols_types_srids=geometry_cols_types_srids, 1880 json_cols=declared_json_cols, 1881 ) 1882 if as_polars 1883 else df 1884 ) 1885 1886 for col, typ in {k: v for k, v in common_diff_dtypes.items()}.items(): 1887 previous_typ = common_dtypes[col] 1888 mixed_numeric_types = (is_dtype_numeric(typ) and is_dtype_numeric(previous_typ)) 1889 explicitly_float = are_dtypes_equal(explicit_dtypes.get(col, 'object'), 'float') 1890 explicitly_int = are_dtypes_equal(explicit_dtypes.get(col, 'object'), 'int') 1891 explicitly_numeric = explicit_dtypes.get(col, 'object').startswith('numeric') 1892 all_nan = ( 1893 df[col].isnull().all() 1894 if mixed_numeric_types and coerce_numeric and not (explicitly_float or explicitly_int) 1895 else None 1896 ) 1897 cast_to_numeric = explicitly_numeric or ( 1898 ( 1899 col in df_numeric_cols 1900 or ( 1901 mixed_numeric_types 1902 and not (explicitly_float or explicitly_int) 1903 and not all_nan 1904 and coerce_numeric 1905 ) 1906 ) 1907 ) 1908 1909 if debug and (explicitly_numeric or df_numeric_cols or mixed_numeric_types): 1910 from meerschaum.utils.formatting import make_header 1911 msg = ( 1912 make_header(f"Coercing column '{col}' to numeric:", left_pad=0) 1913 + "\n" 1914 + f" Previous type: {previous_typ}\n" 1915 + f" Current type: {typ if col not in df_numeric_cols else 'Decimal'}" 1916 + ("\n Column is explicitly numeric." if explicitly_numeric else "") 1917 ) if cast_to_numeric else ( 1918 f"Will not coerce column '{col}' to numeric.\n" 1919 f" Numeric columns in dataframe: {df_numeric_cols}\n" 1920 f" Mixed numeric types: {mixed_numeric_types}\n" 1921 f" Explicitly float: {explicitly_float}\n" 1922 f" Explicitly int: {explicitly_int}\n" 1923 f" All NaN: {all_nan}\n" 1924 f" Coerce numeric: {coerce_numeric}" 1925 ) 1926 dprint(msg) 1927 1928 if cast_to_numeric: 1929 common_dtypes[col] = attempt_cast_to_numeric 1930 common_diff_dtypes[col] = attempt_cast_to_numeric 1931 1932 for d in common_diff_dtypes: 1933 t = common_dtypes[d] 1934 if debug: 1935 dprint(f"Casting column {d} to dtype {t}.") 1936 try: 1937 df[d] = ( 1938 df[d].apply(t) 1939 if callable(t) 1940 else df[d].astype(t) 1941 ) 1942 except Exception as e: 1943 if debug: 1944 dprint(f"Encountered an error when casting column {d} to type {t}:\n{e}\ndf:\n{df}") 1945 if 'int' in str(t).lower(): 1946 try: 1947 df[d] = df[d].astype('float64').astype(t) 1948 except Exception: 1949 if debug: 1950 dprint(f"Was unable to convert to float then {t}.") 1951 return ( 1952 to_polars( 1953 df, 1954 geometry_cols_types_srids=geometry_cols_types_srids, 1955 json_cols=declared_json_cols, 1956 ) 1957 if as_polars 1958 else df 1959 ) 1960 1961 1962def get_datetime_bound_from_df( 1963 df: Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]], 1964 datetime_column: str, 1965 minimum: bool = True, 1966) -> Union[int, datetime, None]: 1967 """ 1968 Return the minimum or maximum datetime (or integer) from a DataFrame. 1969 1970 Parameters 1971 ---------- 1972 df: Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]] 1973 The DataFrame, list, or dict which contains the range axis. 1974 1975 datetime_column: str 1976 The name of the datetime (or int) column. 1977 1978 minimum: bool 1979 Whether to return the minimum (default) or maximum value. 1980 1981 Returns 1982 ------- 1983 The minimum or maximum datetime value in the dataframe, or `None`. 1984 """ 1985 from meerschaum.utils.dtypes import to_datetime, value_is_null 1986 1987 if df is None: 1988 return None 1989 if not datetime_column: 1990 return None 1991 1992 def compare(a, b): 1993 if a is None: 1994 return b 1995 if b is None: 1996 return a 1997 if minimum: 1998 return a if a < b else b 1999 return a if a > b else b 2000 2001 if isinstance(df, list): 2002 if len(df) == 0: 2003 return None 2004 best_yet = df[0].get(datetime_column, None) 2005 for doc in df: 2006 val = doc.get(datetime_column, None) 2007 best_yet = compare(best_yet, val) 2008 return best_yet 2009 2010 if isinstance(df, dict): 2011 if datetime_column not in df: 2012 return None 2013 best_yet = df[datetime_column][0] 2014 for val in df[datetime_column]: 2015 best_yet = compare(best_yet, val) 2016 return best_yet 2017 2018 if 'DataFrame' in str(type(df)): 2019 from meerschaum.utils.dtypes import are_dtypes_equal 2020 pandas = mrsm.attempt_import('pandas') 2021 is_dask = 'dask' in df.__module__ 2022 2023 if datetime_column not in df.columns: 2024 return None 2025 2026 try: 2027 dt_val = ( 2028 df[datetime_column].min(skipna=True) 2029 if minimum 2030 else df[datetime_column].max(skipna=True) 2031 ) 2032 except Exception: 2033 dt_val = pandas.NA 2034 if is_dask and dt_val is not None and dt_val is not pandas.NA: 2035 dt_val = dt_val.compute() 2036 2037 return ( 2038 to_datetime(dt_val, as_pydatetime=True) 2039 if are_dtypes_equal(str(type(dt_val)), 'datetime') 2040 else (dt_val if not value_is_null(dt_val) else None) 2041 ) 2042 2043 return None 2044 2045 2046def get_unique_index_values( 2047 df: Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]], 2048 indices: List[str], 2049) -> Dict[str, List[Any]]: 2050 """ 2051 Return a dictionary of the unique index values in a DataFrame. 2052 2053 Parameters 2054 ---------- 2055 df: Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]] 2056 The dataframe (or list or dict) which contains index values. 2057 2058 indices: List[str] 2059 The list of index columns. 2060 2061 Returns 2062 ------- 2063 A dictionary mapping indices to unique values. 2064 """ 2065 if df is None: 2066 return {} 2067 if 'dataframe' in str(type(df)).lower(): 2068 pandas = mrsm.attempt_import('pandas') 2069 return { 2070 col: list({ 2071 (val if val is not pandas.NA else None) 2072 for val in df[col].unique() 2073 }) 2074 for col in indices 2075 if col in df.columns 2076 } 2077 2078 unique_indices = defaultdict(lambda: set()) 2079 if isinstance(df, list): 2080 for doc in df: 2081 for index in indices: 2082 if index in doc: 2083 unique_indices[index].add(doc[index]) 2084 2085 elif isinstance(df, dict): 2086 for index in indices: 2087 if index in df: 2088 unique_indices[index] = unique_indices[index].union(set(df[index])) 2089 2090 return {key: list(val) for key, val in unique_indices.items()} 2091 2092 2093def df_is_chunk_generator(df: Any) -> bool: 2094 """ 2095 Determine whether to treat `df` as a chunk generator. 2096 2097 Note this should only be used in a context where generators are expected, 2098 as it will return `True` for any iterable. 2099 2100 Parameters 2101 ---------- 2102 The DataFrame or chunk generator to evaluate. 2103 2104 Returns 2105 ------- 2106 A `bool` indicating whether to treat `df` as a generator. 2107 """ 2108 return ( 2109 not isinstance(df, (dict, list, str)) 2110 and 'DataFrame' not in str(type(df)) 2111 and isinstance(df, (Generator, Iterable, Iterator)) 2112 ) 2113 2114 2115def chunksize_to_npartitions(chunksize: Optional[int]) -> int: 2116 """ 2117 Return the Dask `npartitions` value for a given `chunksize`. 2118 """ 2119 if chunksize == -1: 2120 from meerschaum.config import get_config 2121 chunksize = get_config('system', 'connectors', 'sql', 'chunksize') 2122 if chunksize is None: 2123 return 1 2124 return -1 * chunksize 2125 2126 2127def df_from_literal( 2128 pipe: Optional[mrsm.Pipe] = None, 2129 literal: Optional[str] = None, 2130 debug: bool = False 2131) -> 'pd.DataFrame': 2132 """ 2133 Construct a dataframe from a literal value, using the pipe's datetime and value column names. 2134 2135 Parameters 2136 ---------- 2137 pipe: Optional['meerschaum.Pipe'], default None 2138 The pipe which will consume the literal value. 2139 2140 Returns 2141 ------- 2142 A 1-row pandas DataFrame from with the current UTC timestamp as the datetime columns 2143 and the literal as the value. 2144 """ 2145 from meerschaum.utils.packages import import_pandas 2146 from meerschaum.utils.warnings import error, warn 2147 from meerschaum.utils.debug import dprint 2148 from meerschaum.utils.dtypes import get_current_timestamp 2149 2150 if pipe is None or literal is None: 2151 error("Please provide a Pipe and a literal value") 2152 2153 dt_col = pipe.columns.get( 2154 'datetime', 2155 mrsm.get_config('pipes', 'autotime', 'column_name_if_datetime_missing') 2156 ) 2157 val_col = pipe.get_val_column(debug=debug) 2158 2159 val = literal 2160 if isinstance(literal, str): 2161 if debug: 2162 dprint(f"Received literal string: '{literal}'") 2163 import ast 2164 try: 2165 val = ast.literal_eval(literal) 2166 except Exception: 2167 warn( 2168 "Failed to parse value from string:\n" + f"{literal}" + 2169 "\n\nWill cast as a string instead."\ 2170 ) 2171 val = literal 2172 2173 now = get_current_timestamp(pipe.precision) 2174 pd = import_pandas() 2175 return pd.DataFrame({dt_col: [now], val_col: [val]}) 2176 2177 2178def get_first_valid_dask_partition(ddf: 'dask.dataframe.DataFrame') -> Union['pd.DataFrame', None]: 2179 """ 2180 Return the first valid Dask DataFrame partition (if possible). 2181 """ 2182 pdf = None 2183 for partition in ddf.partitions: 2184 try: 2185 pdf = partition.compute() 2186 except Exception: 2187 continue 2188 if len(pdf) > 0: 2189 return pdf 2190 _ = mrsm.attempt_import('partd', lazy=False) 2191 return ddf.compute() 2192 2193 2194def query_df( 2195 df: 'pd.DataFrame', 2196 params: Optional[Dict[str, Any]] = None, 2197 begin: Union[datetime, int, None] = None, 2198 end: Union[datetime, int, None] = None, 2199 datetime_column: Optional[str] = None, 2200 select_columns: Optional[List[str]] = None, 2201 omit_columns: Optional[List[str]] = None, 2202 inplace: bool = False, 2203 reset_index: bool = False, 2204 coerce_types: bool = False, 2205 debug: bool = False, 2206) -> 'pd.DataFrame': 2207 """ 2208 Query the dataframe with the params dictionary. 2209 2210 Parameters 2211 ---------- 2212 df: pd.DataFrame 2213 The DataFrame to query against. 2214 2215 params: Optional[Dict[str, Any]], default None 2216 The parameters dictionary to use for the query. 2217 2218 begin: Union[datetime, int, None], default None 2219 If `begin` and `datetime_column` are provided, only return rows with a timestamp 2220 greater than or equal to this value. 2221 2222 end: Union[datetime, int, None], default None 2223 If `begin` and `datetime_column` are provided, only return rows with a timestamp 2224 less than this value. 2225 2226 datetime_column: Optional[str], default None 2227 A `datetime_column` must be provided to use `begin` and `end`. 2228 2229 select_columns: Optional[List[str]], default None 2230 If provided, only return these columns. 2231 2232 omit_columns: Optional[List[str]], default None 2233 If provided, do not include these columns in the result. 2234 2235 inplace: bool, default False 2236 If `True`, modify the DataFrame inplace rather than creating a new DataFrame. 2237 2238 reset_index: bool, default False 2239 If `True`, reset the index in the resulting DataFrame. 2240 2241 coerce_types: bool, default False 2242 If `True`, cast the dataframe and parameters as strings before querying. 2243 2244 Returns 2245 ------- 2246 A Pandas DataFrame query result. 2247 """ 2248 2249 def _process_select_columns(_df): 2250 if not select_columns: 2251 return 2252 for col in list(_df.columns): 2253 if col not in select_columns: 2254 del _df[col] 2255 2256 def _process_omit_columns(_df): 2257 if not omit_columns: 2258 return 2259 for col in list(_df.columns): 2260 if col in omit_columns: 2261 del _df[col] 2262 2263 if not params and not begin and not end: 2264 if not inplace: 2265 df = df.copy() 2266 _process_select_columns(df) 2267 _process_omit_columns(df) 2268 return df 2269 2270 from meerschaum.utils.debug import dprint 2271 from meerschaum.utils.misc import get_in_ex_params 2272 from meerschaum.utils.warnings import warn 2273 from meerschaum.utils.dtypes import are_dtypes_equal, value_is_null 2274 dateutil_parser = mrsm.attempt_import('dateutil.parser') 2275 pandas = mrsm.attempt_import('pandas') 2276 NA = pandas.NA 2277 2278 if params: 2279 proto_in_ex_params = get_in_ex_params(params) 2280 for key, (proto_in_vals, proto_ex_vals) in proto_in_ex_params.items(): 2281 if proto_ex_vals: 2282 coerce_types = True 2283 break 2284 params = params.copy() 2285 for key, val in {k: v for k, v in params.items()}.items(): 2286 if isinstance(val, (list, tuple, set)) or hasattr(val, 'astype'): 2287 if None in val: 2288 val = [item for item in val if item is not None] + [NA] 2289 params[key] = val 2290 if coerce_types: 2291 params[key] = [str(x) for x in val] 2292 else: 2293 if value_is_null(val): 2294 val = NA 2295 params[key] = NA 2296 if coerce_types: 2297 params[key] = str(val) 2298 2299 dtypes = {col: str(typ) for col, typ in df.dtypes.items()} 2300 2301 if inplace: 2302 df.fillna(NA, inplace=True) 2303 else: 2304 df = df.infer_objects().fillna(NA) 2305 2306 if isinstance(begin, str): 2307 begin = dateutil_parser.parse(begin) 2308 if isinstance(end, str): 2309 end = dateutil_parser.parse(end) 2310 2311 if begin is not None or end is not None: 2312 if not datetime_column or datetime_column not in df.columns: 2313 warn( 2314 f"The datetime column '{datetime_column}' is not present in the Dataframe, " 2315 + "ignoring begin and end...", 2316 ) 2317 begin, end = None, None 2318 2319 if debug: 2320 dprint(f"Querying dataframe:\n{params=} {begin=} {end=} {datetime_column=}") 2321 2322 if datetime_column and (begin is not None or end is not None): 2323 if debug: 2324 dprint("Checking for datetime column compatability.") 2325 2326 from meerschaum.utils.dtypes import coerce_timezone 2327 df_is_dt = are_dtypes_equal(str(df.dtypes[datetime_column]), 'datetime') 2328 begin_is_int = are_dtypes_equal(str(type(begin)), 'int') 2329 end_is_int = are_dtypes_equal(str(type(end)), 'int') 2330 2331 if df_is_dt: 2332 df_tz = ( 2333 getattr(df[datetime_column].dt, 'tz', None) 2334 if hasattr(df[datetime_column], 'dt') 2335 else None 2336 ) 2337 2338 if begin_is_int: 2339 begin = datetime.fromtimestamp(int(begin), timezone.utc).replace(tzinfo=None) 2340 if debug: 2341 dprint(f"`begin` will be cast to '{begin}'.") 2342 if end_is_int: 2343 end = datetime.fromtimestamp(int(end), timezone.utc).replace(tzinfo=None) 2344 if debug: 2345 dprint(f"`end` will be cast to '{end}'.") 2346 2347 begin = coerce_timezone(begin, strip_utc=(df_tz is None)) if begin is not None else None 2348 end = coerce_timezone(end, strip_utc=(df_tz is None)) if begin is not None else None 2349 2350 in_ex_params = get_in_ex_params(params) 2351 2352 masks = [ 2353 ( 2354 (df[datetime_column] >= begin) 2355 if begin is not None and datetime_column 2356 else True 2357 ) & ( 2358 (df[datetime_column] < end) 2359 if end is not None and datetime_column 2360 else True 2361 ) 2362 ] 2363 2364 masks.extend([ 2365 ( 2366 ( 2367 (df[col] if not coerce_types else df[col].astype(str)).isin(in_vals) 2368 if in_vals 2369 else True 2370 ) & ( 2371 ~(df[col] if not coerce_types else df[col].astype(str)).isin(ex_vals) 2372 if ex_vals 2373 else True 2374 ) 2375 ) 2376 for col, (in_vals, ex_vals) in in_ex_params.items() 2377 if col in df.columns 2378 ]) 2379 query_mask = masks[0] 2380 for mask in masks[1:]: 2381 query_mask = query_mask & mask 2382 2383 original_cols = df.columns 2384 2385 ### NOTE: We must cast bool columns to `boolean[pyarrow]` 2386 ### to allow for `<NA>` values. 2387 bool_cols = [ 2388 col 2389 for col, typ in df.dtypes.items() 2390 if are_dtypes_equal(str(typ), 'bool') 2391 ] 2392 for col in bool_cols: 2393 df[col] = df[col].astype('boolean[pyarrow]') 2394 2395 if not isinstance(query_mask, bool): 2396 df['__mrsm_mask'] = ( 2397 query_mask.astype('boolean[pyarrow]') 2398 if hasattr(query_mask, 'astype') 2399 else query_mask 2400 ) 2401 2402 if inplace: 2403 df.where(query_mask, other=NA, inplace=True) 2404 df.dropna(how='all', inplace=True) 2405 result_df = df 2406 else: 2407 result_df = df.where(query_mask, other=NA) 2408 result_df.dropna(how='all', inplace=True) 2409 2410 else: 2411 result_df = df 2412 2413 if '__mrsm_mask' in df.columns: 2414 del df['__mrsm_mask'] 2415 if '__mrsm_mask' in result_df.columns: 2416 del result_df['__mrsm_mask'] 2417 2418 if reset_index: 2419 result_df.reset_index(drop=True, inplace=True) 2420 2421 result_df = enforce_dtypes( 2422 result_df, 2423 dtypes, 2424 safe_copy=False, 2425 debug=debug, 2426 coerce_numeric=False, 2427 coerce_timezone=False, 2428 ) 2429 2430 if select_columns == ['*']: 2431 select_columns = None 2432 2433 if not select_columns and not omit_columns: 2434 return result_df[original_cols] 2435 2436 _process_select_columns(result_df) 2437 _process_omit_columns(result_df) 2438 2439 return result_df 2440 2441 2442def to_json( 2443 df: 'pd.DataFrame', 2444 safe_copy: bool = True, 2445 orient: str = 'records', 2446 date_format: str = 'iso', 2447 date_unit: str = 'us', 2448 double_precision: int = 15, 2449 geometry_format: str = 'geojson', 2450 **kwargs: Any 2451) -> str: 2452 """ 2453 Serialize the given dataframe as a JSON string. 2454 2455 Parameters 2456 ---------- 2457 df: pd.DataFrame 2458 The DataFrame to be serialized. 2459 2460 safe_copy: bool, default True 2461 If `False`, modify the DataFrame inplace. 2462 2463 date_format: str, default 'iso' 2464 The default format for timestamps. 2465 2466 date_unit: str, default 'us' 2467 The precision of the timestamps. 2468 2469 double_precision: int, default 15 2470 The number of decimal places to use when encoding floating point values (maximum 15). 2471 2472 geometry_format: str, default 'geojson' 2473 The serialization format for geometry data. 2474 Accepted values are `geojson`, `wkb_hex`, and `wkt`. 2475 2476 Returns 2477 ------- 2478 A JSON string. 2479 """ 2480 import warnings 2481 import functools 2482 from meerschaum.utils.packages import import_pandas 2483 from meerschaum.utils.dtypes import ( 2484 serialize_bytes, 2485 serialize_decimal, 2486 serialize_geometry, 2487 serialize_date, 2488 ) 2489 pd = import_pandas() 2490 uuid_cols = get_uuid_cols(df) 2491 bytes_cols = get_bytes_cols(df) 2492 numeric_cols = get_numeric_cols(df) 2493 date_cols = get_date_cols(df) 2494 geometry_cols = get_geometry_cols(df) 2495 geometry_cols_srids = { 2496 col: int((getattr(df[col].crs, 'srs', '') or '').split(':', maxsplit=1)[-1] or '0') 2497 for col in geometry_cols 2498 } if 'geodataframe' in str(type(df)).lower() else {} 2499 if safe_copy and bool(uuid_cols or bytes_cols or geometry_cols or numeric_cols): 2500 df = df.copy() 2501 if 'geodataframe' in str(type(df)).lower(): 2502 geometry_data = { 2503 col: df[col] 2504 for col in geometry_cols 2505 } 2506 df = pd.DataFrame({ 2507 col: df[col] 2508 for col in df.columns 2509 if col not in geometry_cols 2510 }) 2511 for col in geometry_cols: 2512 df[col] = pd.Series(ob for ob in geometry_data[col]) 2513 for col in uuid_cols: 2514 df[col] = df[col].astype(str) 2515 for col in bytes_cols: 2516 df[col] = df[col].apply(serialize_bytes) 2517 for col in numeric_cols: 2518 df[col] = df[col].apply(serialize_decimal) 2519 for col in date_cols: 2520 df[col] = df[col].apply(serialize_date) 2521 with warnings.catch_warnings(): 2522 warnings.simplefilter("ignore") 2523 for col in geometry_cols: 2524 srid = geometry_cols_srids.get(col, None) or None 2525 df[col] = pd.Series( 2526 serialize_geometry(val, geometry_format=geometry_format, srid=srid) 2527 for val in df[col] 2528 ) 2529 return df.infer_objects().fillna(pd.NA).to_json( 2530 date_format=date_format, 2531 date_unit=date_unit, 2532 double_precision=double_precision, 2533 orient=orient, 2534 **kwargs 2535 ) 2536 2537 2538def to_simple_lines(df: 'pd.DataFrame') -> str: 2539 """ 2540 Serialize a Pandas Dataframe as lines of simple dictionaries. 2541 2542 Parameters 2543 ---------- 2544 df: pd.DataFrame 2545 The dataframe to serialize into simple lines text. 2546 2547 Returns 2548 ------- 2549 A string of simple line dictionaries joined by newlines. 2550 """ 2551 from meerschaum.utils.misc import to_simple_dict 2552 if df is None or len(df) == 0: 2553 return '' 2554 2555 docs = df.to_dict(orient='records') 2556 return '\n'.join(to_simple_dict(doc) for doc in docs) 2557 2558 2559def parse_simple_lines(data: str) -> 'pd.DataFrame': 2560 """ 2561 Parse simple lines text into a DataFrame. 2562 2563 Parameters 2564 ---------- 2565 data: str 2566 The simple lines text to parse into a DataFrame. 2567 2568 Returns 2569 ------- 2570 A dataframe containing the rows serialized in `data`. 2571 """ 2572 from meerschaum.utils.misc import string_to_dict 2573 from meerschaum.utils.packages import import_pandas 2574 pd = import_pandas() 2575 lines = data.splitlines() 2576 try: 2577 docs = [string_to_dict(line) for line in lines] 2578 df = pd.DataFrame(docs) 2579 except Exception: 2580 df = None 2581 2582 if df is None: 2583 raise ValueError("Cannot parse simple lines into a dataframe.") 2584 2585 return df
def
to_pandas(df: Any) -> Any:
62def to_pandas(df: Any) -> Any: 63 """Convert a Polars DataFrame or LazyFrame to Pandas; otherwise return ``df``.""" 64 if df.__class__.__module__.split('.')[0] != 'polars': 65 return df 66 if df.__class__.__name__ == 'LazyFrame': 67 df = df.collect() 68 json_cols = [ 69 col 70 for col, typ in df.schema.items() 71 if getattr(typ, 'ext_name', lambda: None)() == 'arrow.json' 72 ] 73 if json_cols: 74 df = df.with_columns(df[col].ext.storage() for col in json_cols) 75 pandas_df = df.to_pandas(use_pyarrow_extension_array=True) 76 if json_cols: 77 import json 78 for col in json_cols: 79 pandas_df[col] = pandas_df[col].apply( 80 lambda value: json.loads(value) if isinstance(value, str) else None 81 ) 82 return pandas_df
Convert a Polars DataFrame or LazyFrame to Pandas; otherwise return df.
def
to_polars( df: Any, geometry_cols_types_srids: Optional[Dict[str, Tuple[str, Any]]] = None, json_cols: Optional[List[str]] = None) -> Any:
85def to_polars( 86 df: Any, 87 geometry_cols_types_srids: Optional[Dict[str, Tuple[str, Any]]] = None, 88 json_cols: Optional[List[str]] = None, 89) -> Any: 90 """Convert a Pandas DataFrame to Polars; otherwise return ``df``.""" 91 if df.__class__.__module__.split('.')[0] == 'polars': 92 return df 93 polars = mrsm.attempt_import('polars') 94 geometry_cols_types_srids = ( 95 get_geometry_cols(df, with_types_srids=True) 96 if geometry_cols_types_srids is None and get_geometry_cols(df) 97 else (geometry_cols_types_srids or {}) 98 ) 99 geometry_cols_types_srids = { 100 col: type_srid 101 for col, type_srid in geometry_cols_types_srids.items() 102 if col in df.columns 103 } 104 json_cols = [col for col in (json_cols or []) if col in df.columns] 105 if geometry_cols_types_srids or json_cols: 106 try: 107 import json 108 geometry_cols = list(geometry_cols_types_srids) 109 special_cols = list(dict.fromkeys(geometry_cols + json_cols)) 110 polars_df = to_polars( 111 df.drop(columns=special_cols), 112 geometry_cols_types_srids={}, 113 json_cols=[], 114 ) 115 if geometry_cols: 116 from meerschaum.utils.dtypes import attempt_cast_to_geometry 117 shapely = mrsm.attempt_import('shapely', lazy=False) 118 try: 119 from geoarrow.types.type_pyarrow import register_extension_types 120 register_extension_types() 121 except Exception: 122 pass 123 for col, (_, srid) in geometry_cols_types_srids.items(): 124 crs = str(srid) if srid else None 125 if crs and ':' not in crs: 126 crs = 'EPSG:' + crs 127 metadata = json.dumps( 128 ({'crs': crs, 'crs_type': 'authority_code'} if crs else {}), 129 separators=(',', ':'), 130 ) 131 polars_df = polars_df.with_columns(polars.Series( 132 col, 133 shapely.to_wkb( 134 attempt_cast_to_geometry(df[col]), 135 hex=False, 136 include_srid=True, 137 ).tolist(), 138 dtype=polars.Extension('geoarrow.wkb', polars.Binary, metadata), 139 )) 140 if json_cols: 141 from meerschaum.utils.dtypes import json_serialize_value, value_is_null 142 for col in json_cols: 143 values = [] 144 for value in df[col].tolist(): 145 if value_is_null(value): 146 values.append(None) 147 continue 148 if isinstance(value, str): 149 try: 150 value = json.loads(value) 151 except json.JSONDecodeError: 152 pass 153 values.append(json.dumps( 154 value, 155 default=json_serialize_value, 156 separators=(',', ':'), 157 allow_nan=False, 158 )) 159 polars_df = polars_df.with_columns(polars.Series( 160 col, 161 values, 162 dtype=polars.Extension('arrow.json', polars.String, ''), 163 )) 164 return polars_df.select(list(df.columns)) 165 except Exception: 166 # ponytail: Preserve compatibility if Polars changes its unstable extension API. 167 pass 168 if get_uuid_cols(df): 169 return polars.DataFrame(df.to_dict(orient='list'), strict=False) 170 try: 171 return polars.from_pandas(df, include_index=False) 172 except (TypeError, ValueError): 173 # ponytail: Preserve unsupported Python objects; remove when Polars supports UUIDs. 174 return polars.DataFrame(df.to_dict(orient='list'), strict=False)
Convert a Pandas DataFrame to Polars; otherwise return df.
def
add_missing_cols_to_df(df: pandas.DataFrame, dtypes: Dict[str, Any]) -> pandas.DataFrame:
177def add_missing_cols_to_df( 178 df: 'pd.DataFrame', 179 dtypes: Dict[str, Any], 180) -> 'pd.DataFrame': 181 """ 182 Add columns from the dtypes dictionary as null columns to a new DataFrame. 183 184 Parameters 185 ---------- 186 df: pd.DataFrame 187 The dataframe we should copy and add null columns. 188 189 dtypes: 190 The data types dictionary which may contain keys not present in `df.columns`. 191 192 Returns 193 ------- 194 A new `DataFrame` with the keys from `dtypes` added as null columns. 195 If `df.dtypes` is the same as `dtypes`, then return a reference to `df`. 196 NOTE: This will not ensure that dtypes are enforced! 197 198 Examples 199 -------- 200 >>> import pandas as pd 201 >>> df = pd.DataFrame([{'a': 1}]) 202 >>> dtypes = {'b': 'Int64'} 203 >>> add_missing_cols_to_df(df, dtypes) 204 a b 205 0 1 <NA> 206 >>> add_missing_cols_to_df(df, dtypes).dtypes 207 a int64 208 b Int64 209 dtype: object 210 >>> add_missing_cols_to_df(df, {'a': 'object'}).dtypes 211 a int64 212 dtype: object 213 >>> 214 """ 215 if set(df.columns) == set(dtypes): 216 return df 217 218 from meerschaum.utils.packages import attempt_import 219 from meerschaum.utils.dtypes import to_pandas_dtype 220 pandas = attempt_import('pandas') 221 222 def build_series(dtype: str): 223 return pandas.Series([], dtype=to_pandas_dtype(dtype)) 224 225 assign_kwargs = { 226 str(col): build_series(str(typ)) 227 for col, typ in dtypes.items() 228 if col not in df.columns 229 } 230 df_with_cols = df.assign(**assign_kwargs) 231 for col in assign_kwargs: 232 df_with_cols[col] = df_with_cols[col].fillna(pandas.NA) 233 return df_with_cols
Add columns from the dtypes dictionary as null columns to a new DataFrame.
Parameters
- df (pd.DataFrame): The dataframe we should copy and add null columns.
- dtypes:: The data types dictionary which may contain keys not present in
df.columns.
Returns
- A new
DataFramewith the keys fromdtypesadded as null columns. - If
df.dtypesis the same asdtypes, then return a reference todf. - NOTE (This will not ensure that dtypes are enforced!):
Examples
>>> import pandas as pd
>>> df = pd.DataFrame([{'a': 1}])
>>> dtypes = {'b': 'Int64'}
>>> add_missing_cols_to_df(df, dtypes)
a b
0 1 <NA>
>>> add_missing_cols_to_df(df, dtypes)meerschaum.utils.dtypes
a int64
b Int64
dtype: object
>>> add_missing_cols_to_df(df, {'a': 'object'})meerschaum.utils.dtypes
a int64
dtype: object
>>>
def
filter_unseen_df( old_df: pandas.DataFrame, new_df: pandas.DataFrame, safe_copy: bool = True, dtypes: Optional[Dict[str, Any]] = None, include_unchanged_columns: bool = False, coerce_mixed_numerics: bool = True, debug: bool = False) -> pandas.DataFrame:
236def filter_unseen_df( 237 old_df: 'pd.DataFrame', 238 new_df: 'pd.DataFrame', 239 safe_copy: bool = True, 240 dtypes: Optional[Dict[str, Any]] = None, 241 include_unchanged_columns: bool = False, 242 coerce_mixed_numerics: bool = True, 243 debug: bool = False, 244) -> 'pd.DataFrame': 245 """ 246 Left join two DataFrames to find the newest unseen data. 247 248 Parameters 249 ---------- 250 old_df: 'pd.DataFrame' 251 The original (target) dataframe. Acts as a filter on the `new_df`. 252 253 new_df: 'pd.DataFrame' 254 The fetched (source) dataframe. Rows that are contained in `old_df` are removed. 255 256 safe_copy: bool, default True 257 If `True`, create a copy before comparing and modifying the dataframes. 258 Setting to `False` may mutate the DataFrames. 259 260 dtypes: Optional[Dict[str, Any]], default None 261 Optionally specify the datatypes of the dataframe. 262 263 include_unchanged_columns: bool, default False 264 If `True`, include columns which haven't changed on rows which have changed. 265 266 coerce_mixed_numerics: bool, default True 267 If `True`, cast mixed integer and float columns between the old and new dataframes into 268 numeric values (`decimal.Decimal`). 269 270 debug: bool, default False 271 Verbosity toggle. 272 273 Returns 274 ------- 275 A pandas dataframe of the new, unseen rows in `new_df`. 276 277 Examples 278 -------- 279 ```python 280 >>> import pandas as pd 281 >>> df1 = pd.DataFrame({'a': [1,2]}) 282 >>> df2 = pd.DataFrame({'a': [2,3]}) 283 >>> filter_unseen_df(df1, df2) 284 a 285 0 3 286 287 ``` 288 289 """ 290 if old_df is None: 291 return new_df 292 293 if safe_copy: 294 old_df = old_df.copy() 295 new_df = new_df.copy() 296 297 import json 298 import functools 299 import traceback 300 from meerschaum.utils.warnings import warn 301 from meerschaum.utils.packages import import_pandas, attempt_import 302 from meerschaum.utils.dtypes import ( 303 to_pandas_dtype, 304 are_dtypes_equal, 305 attempt_cast_to_numeric, 306 attempt_cast_to_uuid, 307 attempt_cast_to_bytes, 308 attempt_cast_to_geometry, 309 coerce_timezone, 310 serialize_decimal, 311 ) 312 from meerschaum.utils.dtypes.sql import get_numeric_precision_scale 313 pd = import_pandas(debug=debug) 314 is_dask = 'dask' in new_df.__module__ 315 if is_dask: 316 pandas = attempt_import('pandas') 317 _ = attempt_import('partd', lazy=False) 318 dd = attempt_import('dask.dataframe') 319 merge = dd.merge 320 NA = pandas.NA 321 else: 322 merge = pd.merge 323 NA = pd.NA 324 325 new_df_dtypes = dict(new_df.dtypes) 326 new_cols = list(new_df_dtypes) 327 old_df_dtypes = dict(old_df.dtypes) 328 329 same_cols = set(new_df.columns) == set(old_df.columns) 330 if not same_cols: 331 new_df = add_missing_cols_to_df(new_df, old_df_dtypes) 332 old_df = add_missing_cols_to_df(old_df, new_df_dtypes) 333 334 new_types_missing_from_old = { 335 col: typ 336 for col, typ in new_df_dtypes.items() 337 if col not in old_df_dtypes 338 } 339 old_types_missing_from_new = { 340 col: typ 341 for col, typ in old_df_dtypes.items() 342 if col not in new_df_dtypes 343 } 344 old_df_dtypes.update(new_types_missing_from_old) 345 new_df_dtypes.update(old_types_missing_from_new) 346 347 ### Edge case: two empty lists cast to DFs. 348 elif len(new_df.columns) == 0: 349 return new_df 350 351 try: 352 ### Order matters when checking equality. 353 new_df = new_df[old_df.columns] 354 355 except Exception as e: 356 warn( 357 "Was not able to cast old columns onto new DataFrame. " + 358 f"Are both DataFrames the same shape? Error:\n{e}", 359 stacklevel=3, 360 ) 361 return new_df[list(new_df_dtypes.keys())] 362 363 ### assume the old_df knows what it's doing, even if it's technically wrong. 364 if dtypes is None: 365 dtypes = {col: str(typ) for col, typ in old_df.dtypes.items()} 366 367 numeric_cols_precisions_scales = { 368 col: get_numeric_precision_scale(None, typ) 369 for col, typ in dtypes.items() 370 if col and str(typ).lower().startswith('numeric') 371 } 372 dtypes = { 373 col: to_pandas_dtype(typ) 374 for col, typ in dtypes.items() 375 if col in new_df_dtypes and col in old_df_dtypes 376 } 377 for col, typ in new_df_dtypes.items(): 378 if col not in dtypes: 379 dtypes[col] = typ 380 381 dt_dtypes = { 382 col: typ 383 for col, typ in dtypes.items() 384 if are_dtypes_equal(typ, 'datetime') 385 } 386 non_dt_dtypes = { 387 col: typ 388 for col, typ in dtypes.items() 389 if col not in dt_dtypes 390 } 391 392 cast_non_dt_cols = True 393 try: 394 new_df = new_df.astype(non_dt_dtypes) 395 cast_non_dt_cols = False 396 except Exception as e: 397 warn( 398 f"Was not able to cast the new DataFrame to the given dtypes.\n{e}" 399 ) 400 401 cast_dt_cols = True 402 try: 403 for col, typ in dt_dtypes.items(): 404 _dtypes_col_dtype = str((dtypes or {}).get(col, 'datetime')) 405 strip_utc = ( 406 _dtypes_col_dtype.startswith('datetime64') 407 and 'utc' not in _dtypes_col_dtype.lower() 408 ) 409 if col in old_df.columns: 410 old_df[col] = coerce_timezone( 411 old_df[col], strip_utc=strip_utc 412 ).astype(typ) 413 if col in new_df.columns: 414 new_df[col] = coerce_timezone( 415 new_df[col], strip_utc=strip_utc 416 ).astype(typ) 417 cast_dt_cols = False 418 except Exception as e: 419 warn(f"Could not cast datetime columns:\n{e}") 420 421 cast_cols = cast_dt_cols or cast_non_dt_cols 422 423 new_numeric_cols_existing = get_numeric_cols(new_df) 424 old_numeric_cols = get_numeric_cols(old_df) 425 for col, typ in {k: v for k, v in dtypes.items()}.items(): 426 if not are_dtypes_equal(new_df_dtypes.get(col, 'None'), old_df_dtypes.get(col, 'None')): 427 new_is_float = are_dtypes_equal(new_df_dtypes.get(col, 'None'), 'float') 428 new_is_int = are_dtypes_equal(new_df_dtypes.get(col, 'None'), 'int') 429 new_is_numeric = col in new_numeric_cols_existing 430 old_is_float = are_dtypes_equal(old_df_dtypes.get(col, 'None'), 'float') 431 old_is_int = are_dtypes_equal(old_df_dtypes.get(col, 'None'), 'int') 432 old_is_numeric = col in old_numeric_cols 433 434 if ( 435 coerce_mixed_numerics 436 and 437 (new_is_float or new_is_int or new_is_numeric) 438 and 439 (old_is_float or old_is_int or old_is_numeric) 440 ): 441 dtypes[col] = attempt_cast_to_numeric 442 cast_cols = True 443 continue 444 445 ### Fallback to object if the types don't match. 446 warn( 447 f"Detected different types for '{col}' " 448 + f"({new_df_dtypes.get(col, None)} vs {old_df_dtypes.get(col, None)}), " 449 + "falling back to 'object'..." 450 ) 451 dtypes[col] = 'object' 452 cast_cols = True 453 454 if cast_cols: 455 for col, dtype in dtypes.items(): 456 for df_to_cast in (new_df, old_df): 457 if col not in df_to_cast.columns: 458 continue 459 try: 460 df_to_cast[col] = ( 461 df_to_cast[col].astype(dtype) 462 if not callable(dtype) 463 else df_to_cast[col].apply(dtype) 464 ) 465 except Exception as e: 466 warn(f"Was not able to cast column '{col}' to dtype '{dtype}'.\n{e}") 467 468 serializer = functools.partial(json.dumps, sort_keys=True, separators=(',', ':'), default=str) 469 new_json_cols = get_json_cols(new_df) 470 old_json_cols = get_json_cols(old_df) 471 json_cols = set(new_json_cols + old_json_cols) 472 for json_col in old_json_cols: 473 old_df[json_col] = old_df[json_col].apply(serializer) 474 for json_col in new_json_cols: 475 new_df[json_col] = new_df[json_col].apply(serializer) 476 477 new_numeric_cols = get_numeric_cols(new_df) 478 numeric_cols = set(new_numeric_cols + old_numeric_cols) 479 for numeric_col in old_numeric_cols: 480 old_df[numeric_col] = old_df[numeric_col].apply(serialize_decimal) 481 for numeric_col in new_numeric_cols: 482 new_df[numeric_col] = new_df[numeric_col].apply(serialize_decimal) 483 484 old_dt_cols = [ 485 col 486 for col, typ in old_df.dtypes.items() 487 if are_dtypes_equal(str(typ), 'datetime') 488 ] 489 for col in old_dt_cols: 490 _dtypes_col_dtype = str((dtypes or {}).get(col, 'datetime')) 491 strip_utc = ( 492 _dtypes_col_dtype.startswith('datetime64') 493 and 'utc' not in _dtypes_col_dtype.lower() 494 ) 495 old_df[col] = coerce_timezone(old_df[col], strip_utc=strip_utc) 496 497 new_dt_cols = [ 498 col 499 for col, typ in new_df.dtypes.items() 500 if are_dtypes_equal(str(typ), 'datetime') 501 ] 502 for col in new_dt_cols: 503 _dtypes_col_dtype = str((dtypes or {}).get(col, 'datetime')) 504 strip_utc = ( 505 _dtypes_col_dtype.startswith('datetime64') 506 and 'utc' not in _dtypes_col_dtype.lower() 507 ) 508 new_df[col] = coerce_timezone(new_df[col], strip_utc=strip_utc) 509 510 old_uuid_cols = get_uuid_cols(old_df) 511 new_uuid_cols = get_uuid_cols(new_df) 512 uuid_cols = set(new_uuid_cols + old_uuid_cols) 513 514 old_bytes_cols = get_bytes_cols(old_df) 515 new_bytes_cols = get_bytes_cols(new_df) 516 bytes_cols = set(new_bytes_cols + old_bytes_cols) 517 518 old_geometry_cols = get_geometry_cols(old_df) 519 new_geometry_cols = get_geometry_cols(new_df) 520 geometry_cols = set(new_geometry_cols + old_geometry_cols) 521 522 na_pattern = r'(?i)^(none|nan|na|nat|natz|<na>)$' 523 def normalize_nulls(df_to_normalize): 524 normalized_df = df_to_normalize.infer_objects() 525 string_cols = normalized_df.select_dtypes( 526 include=['object', 'string', 'category'] 527 ).columns 528 return normalized_df.replace( 529 {col: na_pattern for col in string_cols}, 530 pd.NA, 531 regex=True, 532 ).fillna(NA) 533 534 normalized_new_df = normalize_nulls(new_df) 535 normalized_old_df = normalize_nulls(old_df) 536 delta_df = ( 537 None 538 if is_dask 539 else _filter_unseen_df_with_polars(normalized_new_df, normalized_old_df) 540 ) 541 if delta_df is None: 542 joined_df = merge( 543 normalized_new_df, 544 normalized_old_df, 545 how='left', 546 on=None, 547 indicator=True, 548 ) 549 changed_rows_mask = (joined_df['_merge'] == 'left_only') 550 delta_df = joined_df[new_cols][changed_rows_mask].reset_index(drop=True) 551 else: 552 merge_dtypes = merge( 553 normalized_new_df.head(0), 554 normalized_old_df.head(0), 555 how='left', 556 on=None, 557 indicator=True, 558 ).dtypes 559 delta_df = delta_df.astype({ 560 col: typ 561 for col, typ in merge_dtypes.items() 562 if col in delta_df.columns and str(delta_df.dtypes[col]) != str(typ) 563 }) 564 delta_df = delta_df[new_cols] 565 566 delta_json_cols = get_json_cols(delta_df) 567 for json_col in json_cols: 568 if ( 569 json_col in delta_json_cols 570 or json_col not in delta_df.columns 571 ): 572 continue 573 try: 574 delta_df[json_col] = delta_df[json_col].apply( 575 lambda x: (json.loads(x) if isinstance(x, str) else x) 576 ) 577 except Exception: 578 warn(f"Unable to deserialize JSON column '{json_col}':\n{traceback.format_exc()}") 579 580 delta_numeric_cols = get_numeric_cols(delta_df) 581 for numeric_col in numeric_cols: 582 if ( 583 numeric_col in delta_numeric_cols 584 or numeric_col not in delta_df.columns 585 ): 586 continue 587 try: 588 delta_df[numeric_col] = delta_df[numeric_col].apply( 589 functools.partial( 590 attempt_cast_to_numeric, 591 quantize=True, 592 precision=numeric_cols_precisions_scales.get(numeric_col, (None, None))[0], 593 scale=numeric_cols_precisions_scales.get(numeric_col, (None, None))[1], 594 ) 595 ) 596 except Exception: 597 warn(f"Unable to parse numeric column '{numeric_col}':\n{traceback.format_exc()}") 598 599 delta_uuid_cols = get_uuid_cols(delta_df) 600 for uuid_col in uuid_cols: 601 if ( 602 uuid_col in delta_uuid_cols 603 or uuid_col not in delta_df.columns 604 ): 605 continue 606 try: 607 delta_df[uuid_col] = delta_df[uuid_col].apply(attempt_cast_to_uuid) 608 except Exception: 609 warn(f"Unable to parse numeric column '{uuid_col}':\n{traceback.format_exc()}") 610 611 delta_bytes_cols = get_bytes_cols(delta_df) 612 for bytes_col in bytes_cols: 613 if ( 614 bytes_col in delta_bytes_cols 615 or bytes_col not in delta_df.columns 616 ): 617 continue 618 try: 619 delta_df[bytes_col] = delta_df[bytes_col].apply(attempt_cast_to_bytes) 620 except Exception: 621 warn(f"Unable to parse bytes column '{bytes_col}':\n{traceback.format_exc()}") 622 623 delta_geometry_cols = get_geometry_cols(delta_df) 624 for geometry_col in geometry_cols: 625 if ( 626 geometry_col in delta_geometry_cols 627 or geometry_col not in delta_df.columns 628 ): 629 continue 630 try: 631 delta_df[geometry_col] = attempt_cast_to_geometry(delta_df[geometry_col]) 632 except Exception: 633 warn(f"Unable to parse geometry column '{geometry_col}':\n{traceback.format_exc()}") 634 635 return delta_df
Left join two DataFrames to find the newest unseen data.
Parameters
- old_df ('pd.DataFrame'):
The original (target) dataframe. Acts as a filter on the
new_df. - new_df ('pd.DataFrame'):
The fetched (source) dataframe. Rows that are contained in
old_dfare removed. - safe_copy (bool, default True):
If
True, create a copy before comparing and modifying the dataframes. Setting toFalsemay mutate the DataFrames. - dtypes (Optional[Dict[str, Any]], default None): Optionally specify the datatypes of the dataframe.
- include_unchanged_columns (bool, default False):
If
True, include columns which haven't changed on rows which have changed. - coerce_mixed_numerics (bool, default True):
If
True, cast mixed integer and float columns between the old and new dataframes into numeric values (decimal.Decimal). - debug (bool, default False): Verbosity toggle.
Returns
- A pandas dataframe of the new, unseen rows in
new_df.
Examples
>>> import pandas as pd
>>> df1 = pd.DataFrame({'a': [1,2]})
>>> df2 = pd.DataFrame({'a': [2,3]})
>>> filter_unseen_df(df1, df2)
a
0 3
def
parse_df_datetimes( df: pandas.DataFrame, ignore_cols: Optional[Iterable[str]] = None, strip_timezone: bool = False, chunksize: Optional[int] = None, dtype_backend: str = 'numpy_nullable', ignore_all: bool = False, precision_unit: Optional[str] = None, coerce_utc: bool = True, debug: bool = False) -> pandas.DataFrame:
638def parse_df_datetimes( 639 df: 'pd.DataFrame', 640 ignore_cols: Optional[Iterable[str]] = None, 641 strip_timezone: bool = False, 642 chunksize: Optional[int] = None, 643 dtype_backend: str = 'numpy_nullable', 644 ignore_all: bool = False, 645 precision_unit: Optional[str] = None, 646 coerce_utc: bool = True, 647 debug: bool = False, 648) -> 'pd.DataFrame': 649 """ 650 Parse a pandas DataFrame for datetime columns and cast as datetimes. 651 652 Parameters 653 ---------- 654 df: pd.DataFrame 655 The pandas DataFrame to parse. 656 657 ignore_cols: Optional[Iterable[str]], default None 658 If provided, do not attempt to coerce these columns as datetimes. 659 660 strip_timezone: bool, default False 661 If `True`, remove the UTC `tzinfo` property. 662 663 chunksize: Optional[int], default None 664 If the pandas implementation is `'dask'`, use this chunksize for the distributed dataframe. 665 666 dtype_backend: str, default 'numpy_nullable' 667 If `df` is not a DataFrame and new one needs to be constructed, 668 use this as the datatypes backend. 669 Accepted values are 'numpy_nullable' and 'pyarrow'. 670 671 ignore_all: bool, default False 672 If `True`, do not attempt to cast any columns to datetimes. 673 674 precision_unit: Optional[str], default None 675 If provided, enforce the given precision on the coerced datetime columns. 676 677 coerce_utc: bool, default True 678 Coerce the datetime columns to UTC (see `meerschaum.utils.dtypes.to_datetime()`). 679 680 debug: bool, default False 681 Verbosity toggle. 682 683 Returns 684 ------- 685 A new pandas DataFrame with the determined datetime columns 686 (usually ISO strings) cast as datetimes. 687 688 Examples 689 -------- 690 ```python 691 >>> import pandas as pd 692 >>> df = pd.DataFrame({'a': ['2022-01-01 00:00:00']}) 693 >>> df.dtypes 694 a object 695 dtype: object 696 >>> df2 = parse_df_datetimes(df) 697 >>> df2.dtypes 698 a datetime64[us, UTC] 699 dtype: object 700 701 ``` 702 703 """ 704 from meerschaum.utils.packages import import_pandas, attempt_import 705 from meerschaum.utils.debug import dprint 706 from meerschaum.utils.warnings import warn 707 from meerschaum.utils.misc import items_str 708 from meerschaum.utils.dtypes import to_datetime, MRSM_PD_DTYPES 709 import traceback 710 711 pd = import_pandas() 712 pandas = attempt_import('pandas') 713 pd_name = pd.__name__ 714 using_dask = 'dask' in pd_name 715 df_is_dask = (hasattr(df, '__module__') and 'dask' in df.__module__) 716 dask_dataframe = None 717 if using_dask or df_is_dask: 718 npartitions = chunksize_to_npartitions(chunksize) 719 dask_dataframe = attempt_import('dask.dataframe') 720 721 ### if df is a dict, build DataFrame 722 if isinstance(df, pandas.DataFrame): 723 pdf = df 724 elif df_is_dask and isinstance(df, dask_dataframe.DataFrame): 725 pdf = get_first_valid_dask_partition(df) 726 else: 727 if debug: 728 dprint(f"df is of type '{type(df)}'. Building {pd.DataFrame}...") 729 730 if using_dask: 731 if isinstance(df, list): 732 keys = set() 733 for doc in df: 734 for key in doc: 735 keys.add(key) 736 df = pd.DataFrame.from_dict( 737 { 738 k: [ 739 doc.get(k, None) 740 for doc in df 741 ] for k in keys 742 }, 743 npartitions=npartitions, 744 ) 745 elif isinstance(df, dict): 746 df = pd.DataFrame.from_dict(df, npartitions=npartitions) 747 elif 'pandas.core.frame.DataFrame' in str(type(df)): 748 df = pd.from_pandas(df, npartitions=npartitions) 749 else: 750 raise Exception("Can only parse dictionaries or lists of dictionaries with Dask.") 751 pandas = attempt_import('pandas') 752 pdf = get_first_valid_dask_partition(df) 753 754 else: 755 df = pd.DataFrame(df).convert_dtypes(dtype_backend=dtype_backend) 756 pdf = df 757 758 ### skip parsing if DataFrame is empty 759 if len(pdf) == 0: 760 if debug: 761 dprint("df is empty. Returning original DataFrame without casting datetime columns...") 762 return df 763 764 ignore_cols = set( 765 (ignore_cols or []) + [ 766 col 767 for col, dtype in pdf.dtypes.items() 768 if 'datetime' in str(dtype) 769 ] 770 ) 771 cols_to_inspect = [ 772 col 773 for col in pdf.columns 774 if col not in ignore_cols 775 ] if not ignore_all else [] 776 777 if len(cols_to_inspect) == 0: 778 if debug: 779 dprint("All columns are ignored, skipping datetime detection...") 780 return df.infer_objects().fillna(pandas.NA) 781 782 ### apply regex to columns to determine which are ISO datetimes 783 iso_dt_regex = r'\d{4}-\d{2}-\d{2}.\d{2}\:\d{2}\:\d+' 784 dt_mask = pdf[cols_to_inspect].astype(str).apply( 785 lambda s: s.str.match(iso_dt_regex).all() 786 ) 787 788 ### list of datetime column names 789 datetime_cols = [col for col in pdf[cols_to_inspect].loc[:, dt_mask]] 790 if not datetime_cols: 791 if debug: 792 dprint("No columns detected as datetimes, returning...") 793 return df.infer_objects().fillna(pandas.NA) 794 795 if debug: 796 dprint("Converting columns to datetimes: " + str(datetime_cols)) 797 798 def _parse_to_datetime(x): 799 return to_datetime(x, precision_unit=precision_unit, coerce_utc=coerce_utc) 800 801 try: 802 if not using_dask: 803 df[datetime_cols] = df[datetime_cols].apply(_parse_to_datetime) 804 else: 805 df[datetime_cols] = df[datetime_cols].apply( 806 _parse_to_datetime, 807 utc=True, 808 axis=1, 809 meta={ 810 col: MRSM_PD_DTYPES['datetime'] 811 for col in datetime_cols 812 } 813 ) 814 except Exception: 815 warn( 816 f"Unable to apply `to_datetime()` to {items_str(datetime_cols)}:\n" 817 + f"{traceback.format_exc()}" 818 ) 819 820 if strip_timezone: 821 for dt in datetime_cols: 822 try: 823 df[dt] = df[dt].dt.tz_localize(None) 824 except Exception: 825 warn( 826 f"Unable to convert column '{dt}' to naive datetime:\n" 827 + f"{traceback.format_exc()}" 828 ) 829 830 return df.fillna(pandas.NA)
Parse a pandas DataFrame for datetime columns and cast as datetimes.
Parameters
- df (pd.DataFrame): The pandas DataFrame to parse.
- ignore_cols (Optional[Iterable[str]], default None): If provided, do not attempt to coerce these columns as datetimes.
- strip_timezone (bool, default False):
If
True, remove the UTCtzinfoproperty. - chunksize (Optional[int], default None):
If the pandas implementation is
'dask', use this chunksize for the distributed dataframe. - dtype_backend (str, default 'numpy_nullable'):
If
dfis not a DataFrame and new one needs to be constructed, use this as the datatypes backend. Accepted values are 'numpy_nullable' and 'pyarrow'. - ignore_all (bool, default False):
If
True, do not attempt to cast any columns to datetimes. - precision_unit (Optional[str], default None): If provided, enforce the given precision on the coerced datetime columns.
- coerce_utc (bool, default True):
Coerce the datetime columns to UTC (see
meerschaum.utils.dtypes.to_datetime()). - debug (bool, default False): Verbosity toggle.
Returns
- A new pandas DataFrame with the determined datetime columns
- (usually ISO strings) cast as datetimes.
Examples
>>> import pandas as pd
>>> df = pd.DataFrame({'a': ['2022-01-01 00:00:00']})
>>> df.dtypes
a object
dtype: object
>>> df2 = parse_df_datetimes(df)
>>> df2.dtypes
a datetime64[us, UTC]
dtype: object
def
get_unhashable_cols(df: pandas.DataFrame) -> List[str]:
833def get_unhashable_cols(df: 'pd.DataFrame') -> List[str]: 834 """ 835 Get the columns which contain unhashable objects from a Pandas DataFrame. 836 837 Parameters 838 ---------- 839 df: pd.DataFrame 840 The DataFrame which may contain unhashable objects. 841 842 Returns 843 ------- 844 A list of columns. 845 """ 846 if df is None: 847 return [] 848 if len(df) == 0: 849 return [] 850 851 is_dask = 'dask' in df.__module__ 852 if is_dask: 853 from meerschaum.utils.packages import attempt_import 854 pandas = attempt_import('pandas') 855 df = pandas.DataFrame(get_first_valid_dask_partition(df)) 856 return [ 857 col for col, val in df.iloc[0].items() 858 if not isinstance(val, Hashable) 859 ]
Get the columns which contain unhashable objects from a Pandas DataFrame.
Parameters
- df (pd.DataFrame): The DataFrame which may contain unhashable objects.
Returns
- A list of columns.
def
get_json_cols(df: pandas.DataFrame) -> List[str]:
862def get_json_cols(df: 'pd.DataFrame') -> List[str]: 863 """ 864 Get the columns which contain unhashable objects from a Pandas DataFrame. 865 866 Parameters 867 ---------- 868 df: pd.DataFrame 869 The DataFrame which may contain unhashable objects. 870 871 Returns 872 ------- 873 A list of columns to be encoded as JSON. 874 """ 875 if df is None: 876 return [] 877 878 is_dask = 'dask' in df.__module__ if hasattr(df, '__module__') else False 879 if is_dask: 880 df = get_first_valid_dask_partition(df) 881 882 if len(df) == 0: 883 return [] 884 885 cols_indices = { 886 col: df[col].first_valid_index() 887 for col in df.columns 888 } 889 return [ 890 col 891 for col, ix in cols_indices.items() 892 if ( 893 ix is not None 894 and isinstance(df.loc[ix][col], (dict, list)) 895 ) 896 ]
Get the columns which contain unhashable objects from a Pandas DataFrame.
Parameters
- df (pd.DataFrame): The DataFrame which may contain unhashable objects.
Returns
- A list of columns to be encoded as JSON.
def
get_numeric_cols(df: pandas.DataFrame) -> List[str]:
899def get_numeric_cols(df: 'pd.DataFrame') -> List[str]: 900 """ 901 Get the columns which contain `decimal.Decimal` objects from a Pandas DataFrame. 902 903 Parameters 904 ---------- 905 df: pd.DataFrame 906 The DataFrame which may contain decimal objects. 907 908 Returns 909 ------- 910 A list of columns to treat as numerics. 911 """ 912 if df is None: 913 return [] 914 from decimal import Decimal 915 is_dask = 'dask' in df.__module__ 916 if is_dask: 917 df = get_first_valid_dask_partition(df) 918 919 if len(df) == 0: 920 return [] 921 922 cols_indices = { 923 col: df[col].first_valid_index() 924 for col in df.columns 925 } 926 return [ 927 col 928 for col, ix in cols_indices.items() 929 if ( 930 ix is not None 931 and 932 isinstance(df.loc[ix][col], Decimal) 933 ) 934 ]
Get the columns which contain decimal.Decimal objects from a Pandas DataFrame.
Parameters
- df (pd.DataFrame): The DataFrame which may contain decimal objects.
Returns
- A list of columns to treat as numerics.
def
get_bool_cols(df: pandas.DataFrame) -> List[str]:
937def get_bool_cols(df: 'pd.DataFrame') -> List[str]: 938 """ 939 Get the columns which contain `bool` objects from a Pandas DataFrame. 940 941 Parameters 942 ---------- 943 df: pd.DataFrame 944 The DataFrame which may contain bools. 945 946 Returns 947 ------- 948 A list of columns to treat as bools. 949 """ 950 if df is None: 951 return [] 952 953 is_dask = 'dask' in df.__module__ 954 if is_dask: 955 df = get_first_valid_dask_partition(df) 956 957 if len(df) == 0: 958 return [] 959 960 from meerschaum.utils.dtypes import are_dtypes_equal 961 962 return [ 963 col 964 for col, typ in df.dtypes.items() 965 if are_dtypes_equal(str(typ), 'bool') 966 ]
Get the columns which contain bool objects from a Pandas DataFrame.
Parameters
- df (pd.DataFrame): The DataFrame which may contain bools.
Returns
- A list of columns to treat as bools.
def
get_uuid_cols(df: pandas.DataFrame) -> List[str]:
969def get_uuid_cols(df: 'pd.DataFrame') -> List[str]: 970 """ 971 Get the columns which contain `uuid.UUID` objects from a Pandas DataFrame. 972 973 Parameters 974 ---------- 975 df: pd.DataFrame 976 The DataFrame which may contain UUID objects. 977 978 Returns 979 ------- 980 A list of columns to treat as UUIDs. 981 """ 982 if df is None: 983 return [] 984 from uuid import UUID 985 is_dask = 'dask' in df.__module__ 986 if is_dask: 987 df = get_first_valid_dask_partition(df) 988 989 if len(df) == 0: 990 return [] 991 992 cols_indices = { 993 col: df[col].first_valid_index() 994 for col in df.columns 995 } 996 return [ 997 col 998 for col, ix in cols_indices.items() 999 if ( 1000 ix is not None 1001 and 1002 isinstance(df.loc[ix][col], UUID) 1003 ) 1004 ]
Get the columns which contain uuid.UUID objects from a Pandas DataFrame.
Parameters
- df (pd.DataFrame): The DataFrame which may contain UUID objects.
Returns
- A list of columns to treat as UUIDs.
def
get_datetime_cols( df: pandas.DataFrame, timezone_aware: bool = True, timezone_naive: bool = True, with_tz_precision: bool = False) -> Union[List[str], Dict[str, Tuple[Optional[str], str]]]:
1007def get_datetime_cols( 1008 df: 'pd.DataFrame', 1009 timezone_aware: bool = True, 1010 timezone_naive: bool = True, 1011 with_tz_precision: bool = False, 1012) -> Union[List[str], Dict[str, Tuple[Union[str, None], str]]]: 1013 """ 1014 Get the columns which contain `datetime` or `Timestamp` objects from a Pandas DataFrame. 1015 1016 Parameters 1017 ---------- 1018 df: pd.DataFrame 1019 The DataFrame which may contain `datetime` or `Timestamp` objects. 1020 1021 timezone_aware: bool, default True 1022 If `True`, include timezone-aware datetime columns. 1023 1024 timezone_naive: bool, default True 1025 If `True`, include timezone-naive datetime columns. 1026 1027 with_tz_precision: bool, default False 1028 If `True`, return a dictionary mapping column names to tuples in the form 1029 `(timezone, precision)`. 1030 1031 Returns 1032 ------- 1033 A list of columns to treat as datetimes, or a dictionary of columns to tz+precision tuples 1034 (if `with_tz_precision` is `True`). 1035 """ 1036 if not timezone_aware and not timezone_naive: 1037 raise ValueError("`timezone_aware` and `timezone_naive` cannot both be `False`.") 1038 1039 if df is None: 1040 return [] if not with_tz_precision else {} 1041 1042 from meerschaum.utils.dtypes import are_dtypes_equal, MRSM_PRECISION_UNITS_ALIASES 1043 is_dask = 'dask' in df.__module__ 1044 if is_dask: 1045 df = get_first_valid_dask_partition(df) 1046 1047 def get_tz_precision_from_dtype(dtype: str) -> Tuple[Union[str, None], str]: 1048 """ 1049 Extract the tz + precision tuple from a dtype string. 1050 """ 1051 dtype = str(dtype).removesuffix('[pyarrow]') 1052 meta_str = dtype.split('[', maxsplit=1)[-1].rstrip(']').replace(' ', '') 1053 tz = ( 1054 None 1055 if ',' not in meta_str 1056 else meta_str.split(',', maxsplit=1)[-1].removeprefix('tz=') 1057 ) 1058 precision_abbreviation = ( 1059 meta_str 1060 if ',' not in meta_str 1061 else meta_str.split(',')[0] 1062 ) 1063 precision = MRSM_PRECISION_UNITS_ALIASES[precision_abbreviation] 1064 return tz, precision 1065 1066 def get_tz_precision_from_datetime(dt: datetime) -> Tuple[Union[str, None], str]: 1067 """ 1068 Return the tz + precision tuple from a Python datetime object. 1069 """ 1070 return dt.tzname(), 'microsecond' 1071 1072 known_dt_cols_types = { 1073 col: str(typ) 1074 for col, typ in df.dtypes.items() 1075 if are_dtypes_equal('datetime', str(typ)) 1076 } 1077 1078 known_dt_cols_tuples = { 1079 col: get_tz_precision_from_dtype(typ) 1080 for col, typ in known_dt_cols_types.items() 1081 } 1082 1083 if len(df) == 0: 1084 return ( 1085 list(known_dt_cols_types) 1086 if not with_tz_precision 1087 else known_dt_cols_tuples 1088 ) 1089 1090 cols_indices = { 1091 col: df[col].first_valid_index() 1092 for col in df.columns 1093 if col not in known_dt_cols_types 1094 } 1095 pydt_cols_tuples = { 1096 col: get_tz_precision_from_datetime(sample_val) 1097 for col, ix in cols_indices.items() 1098 if ( 1099 ix is not None 1100 and 1101 isinstance((sample_val := df.loc[ix][col]), datetime) 1102 ) 1103 } 1104 1105 dt_cols_tuples = { 1106 **known_dt_cols_tuples, 1107 **pydt_cols_tuples 1108 } 1109 1110 all_dt_cols_tuples = { 1111 col: dt_cols_tuples[col] 1112 for col in df.columns 1113 if col in dt_cols_tuples 1114 } 1115 if timezone_aware and timezone_naive: 1116 return ( 1117 list(all_dt_cols_tuples) 1118 if not with_tz_precision 1119 else all_dt_cols_tuples 1120 ) 1121 1122 known_timezone_aware_dt_cols = [ 1123 col 1124 for col in known_dt_cols_types 1125 if getattr(df[col], 'tz', None) is not None 1126 ] 1127 timezone_aware_pydt_cols_tuples = { 1128 col: (tz, precision) 1129 for col, (tz, precision) in pydt_cols_tuples.items() 1130 if df.loc[cols_indices[col]][col].tzinfo is not None 1131 } 1132 timezone_aware_dt_cols_set = set( 1133 known_timezone_aware_dt_cols + list(timezone_aware_pydt_cols_tuples) 1134 ) 1135 timezone_aware_cols_tuples = { 1136 col: (tz, precision) 1137 for col, (tz, precision) in all_dt_cols_tuples.items() 1138 if col in timezone_aware_dt_cols_set 1139 } 1140 timezone_naive_cols_tuples = { 1141 col: (tz, precision) 1142 for col, (tz, precision) in all_dt_cols_tuples.items() 1143 if col not in timezone_aware_dt_cols_set 1144 } 1145 1146 if timezone_aware: 1147 return ( 1148 list(timezone_aware_cols_tuples) 1149 if not with_tz_precision 1150 else timezone_aware_cols_tuples 1151 ) 1152 1153 return ( 1154 list(timezone_naive_cols_tuples) 1155 if not with_tz_precision 1156 else timezone_naive_cols_tuples 1157 )
Get the columns which contain datetime or Timestamp objects from a Pandas DataFrame.
Parameters
- df (pd.DataFrame):
The DataFrame which may contain
datetimeorTimestampobjects. - timezone_aware (bool, default True):
If
True, include timezone-aware datetime columns. - timezone_naive (bool, default True):
If
True, include timezone-naive datetime columns. - with_tz_precision (bool, default False):
If
True, return a dictionary mapping column names to tuples in the form(timezone, precision).
Returns
- A list of columns to treat as datetimes, or a dictionary of columns to tz+precision tuples
- (if
with_tz_precisionisTrue).
def
get_datetime_cols_types(df: pandas.DataFrame) -> Dict[str, str]:
1160def get_datetime_cols_types(df: 'pd.DataFrame') -> Dict[str, str]: 1161 """ 1162 Return a dictionary mapping datetime columns to specific types strings. 1163 1164 Parameters 1165 ---------- 1166 df: pd.DataFrame 1167 The DataFrame which may contain datetime columns. 1168 1169 Returns 1170 ------- 1171 A dictionary mapping the datetime columns' names to dtype strings 1172 (containing timezone and precision metadata). 1173 1174 Examples 1175 -------- 1176 >>> from datetime import datetime, timezone 1177 >>> import pandas as pd 1178 >>> df = pd.DataFrame({'dt_tz_aware': [datetime(2025, 1, 1, tzinfo=timezone.utc)]}) 1179 >>> get_datetime_cols_types(df) 1180 {'dt_tz_aware': 'datetime64[us, UTC]'} 1181 >>> df = pd.DataFrame({'distant_dt': [datetime(1, 1, 1)]}) 1182 >>> get_datetime_cols_types(df) 1183 {'distant_dt': 'datetime64[us]'} 1184 >>> df = pd.DataFrame({'dt_second': datetime(2025, 1, 1)}) 1185 >>> df['dt_second'] = df['dt_second'].astype('datetime64[s]') 1186 >>> get_datetime_cols_types(df) 1187 {'dt_second': 'datetime64[s]'} 1188 """ 1189 from meerschaum.utils.dtypes import MRSM_PRECISION_UNITS_ABBREVIATIONS 1190 dt_cols_tuples = get_datetime_cols(df, with_tz_precision=True) 1191 if not dt_cols_tuples: 1192 return {} 1193 1194 return { 1195 col: ( 1196 f"datetime64[{MRSM_PRECISION_UNITS_ABBREVIATIONS[precision]}]" 1197 if tz is None 1198 else f"datetime64[{MRSM_PRECISION_UNITS_ABBREVIATIONS[precision]}, {tz}]" 1199 ) 1200 for col, (tz, precision) in dt_cols_tuples.items() 1201 }
Return a dictionary mapping datetime columns to specific types strings.
Parameters
- df (pd.DataFrame): The DataFrame which may contain datetime columns.
Returns
- A dictionary mapping the datetime columns' names to dtype strings
- (containing timezone and precision metadata).
Examples
>>> from datetime import datetime, timezone
>>> import pandas as pd
>>> df = pd.DataFrame({'dt_tz_aware': [datetime(2025, 1, 1, tzinfo=timezone.utc)]})
>>> get_datetime_cols_types(df)
{'dt_tz_aware': 'datetime64[us, UTC]'}
>>> df = pd.DataFrame({'distant_dt': [datetime(1, 1, 1)]})
>>> get_datetime_cols_types(df)
{'distant_dt': 'datetime64[us]'}
>>> df = pd.DataFrame({'dt_second': datetime(2025, 1, 1)})
>>> df['dt_second'] = df['dt_second'].astype('datetime64[s]')
>>> get_datetime_cols_types(df)
{'dt_second': 'datetime64[s]'}
def
get_date_cols(df: pandas.DataFrame) -> List[str]:
1204def get_date_cols(df: 'pd.DataFrame') -> List[str]: 1205 """ 1206 Get the `date` columns from a Pandas DataFrame. 1207 1208 Parameters 1209 ---------- 1210 df: pd.DataFrame 1211 The DataFrame which may contain dates. 1212 1213 Returns 1214 ------- 1215 A list of columns to treat as dates. 1216 """ 1217 from meerschaum.utils.dtypes import are_dtypes_equal 1218 if df is None: 1219 return [] 1220 1221 is_dask = 'dask' in df.__module__ 1222 if is_dask: 1223 df = get_first_valid_dask_partition(df) 1224 1225 known_date_cols = [ 1226 col 1227 for col, typ in df.dtypes.items() 1228 if are_dtypes_equal(typ, 'date') 1229 ] 1230 1231 if len(df) == 0: 1232 return known_date_cols 1233 1234 cols_indices = { 1235 col: df[col].first_valid_index() 1236 for col in df.columns 1237 if col not in known_date_cols 1238 } 1239 object_date_cols = [ 1240 col 1241 for col, ix in cols_indices.items() 1242 if ( 1243 ix is not None 1244 and isinstance(df.loc[ix][col], date) 1245 ) 1246 ] 1247 1248 all_date_cols = set(known_date_cols + object_date_cols) 1249 1250 return [ 1251 col 1252 for col in df.columns 1253 if col in all_date_cols 1254 ]
Get the date columns from a Pandas DataFrame.
Parameters
- df (pd.DataFrame): The DataFrame which may contain dates.
Returns
- A list of columns to treat as dates.
def
get_bytes_cols(df: pandas.DataFrame) -> List[str]:
1257def get_bytes_cols(df: 'pd.DataFrame') -> List[str]: 1258 """ 1259 Get the columns which contain bytes strings from a Pandas DataFrame. 1260 1261 Parameters 1262 ---------- 1263 df: pd.DataFrame 1264 The DataFrame which may contain bytes strings. 1265 1266 Returns 1267 ------- 1268 A list of columns to treat as bytes. 1269 """ 1270 if df is None: 1271 return [] 1272 1273 is_dask = 'dask' in df.__module__ 1274 if is_dask: 1275 df = get_first_valid_dask_partition(df) 1276 1277 known_bytes_cols = [ 1278 col 1279 for col, typ in df.dtypes.items() 1280 if str(typ) in ('binary[pyarrow]', 'large_binary[pyarrow]') 1281 ] 1282 1283 if len(df) == 0: 1284 return known_bytes_cols 1285 1286 cols_indices = { 1287 col: df[col].first_valid_index() 1288 for col in df.columns 1289 if col not in known_bytes_cols 1290 } 1291 object_bytes_cols = [ 1292 col 1293 for col, ix in cols_indices.items() 1294 if ( 1295 ix is not None 1296 and isinstance(df.loc[ix][col], bytes) 1297 ) 1298 ] 1299 1300 all_bytes_cols = set(known_bytes_cols + object_bytes_cols) 1301 1302 return [ 1303 col 1304 for col in df.columns 1305 if col in all_bytes_cols 1306 ]
Get the columns which contain bytes strings from a Pandas DataFrame.
Parameters
- df (pd.DataFrame): The DataFrame which may contain bytes strings.
Returns
- A list of columns to treat as bytes.
def
get_geometry_cols( df: pandas.DataFrame, with_types_srids: bool = False) -> Union[List[str], Dict[str, Any]]:
1309def get_geometry_cols( 1310 df: 'pd.DataFrame', 1311 with_types_srids: bool = False, 1312) -> Union[List[str], Dict[str, Any]]: 1313 """ 1314 Get the columns which contain shapely objects from a Pandas DataFrame. 1315 1316 Parameters 1317 ---------- 1318 df: pd.DataFrame 1319 The DataFrame which may contain bytes strings. 1320 1321 with_types_srids: bool, default False 1322 If `True`, return a dictionary mapping columns to geometry types and SRIDs. 1323 1324 Returns 1325 ------- 1326 A list of columns to treat as `geometry`. 1327 If `with_types_srids`, return a dictionary mapping columns to tuples in the form (type, SRID). 1328 """ 1329 if df is None: 1330 return [] if not with_types_srids else {} 1331 1332 is_dask = 'dask' in df.__module__ 1333 if is_dask: 1334 df = get_first_valid_dask_partition(df) 1335 1336 if len(df) == 0: 1337 return [] if not with_types_srids else {} 1338 1339 cols_indices = { 1340 col: df[col].first_valid_index() 1341 for col in df.columns 1342 } 1343 geo_cols = [ 1344 col 1345 for col, ix in cols_indices.items() 1346 if ( 1347 ix is not None 1348 and 1349 'shapely' in str(type(df.loc[ix][col])) 1350 ) 1351 ] 1352 if not with_types_srids: 1353 return geo_cols 1354 1355 gpd = mrsm.attempt_import('geopandas', lazy=False) 1356 geo_cols_types_srids = {} 1357 for col in geo_cols: 1358 try: 1359 sample_geo_series = gpd.GeoSeries(df[col], crs=None) 1360 geometry_types = { 1361 geom.geom_type 1362 for geom in sample_geo_series 1363 if hasattr(geom, 'geom_type') 1364 } 1365 geometry_has_z = any(getattr(geom, 'has_z', False) for geom in sample_geo_series) 1366 srid = ( 1367 ( 1368 sample_geo_series.crs.sub_crs_list[0].to_epsg() 1369 if sample_geo_series.crs.is_compound 1370 else sample_geo_series.crs.to_epsg() 1371 ) 1372 if sample_geo_series.crs 1373 else 0 1374 ) 1375 geometry_type = list(geometry_types)[0] if len(geometry_types) == 1 else 'geometry' 1376 if geometry_type != 'geometry' and geometry_has_z: 1377 geometry_type = geometry_type + 'Z' 1378 except Exception: 1379 srid = 0 1380 geometry_type = 'geometry' 1381 geo_cols_types_srids[col] = (geometry_type, srid) 1382 1383 return geo_cols_types_srids
Get the columns which contain shapely objects from a Pandas DataFrame.
Parameters
- df (pd.DataFrame): The DataFrame which may contain bytes strings.
- with_types_srids (bool, default False):
If
True, return a dictionary mapping columns to geometry types and SRIDs.
Returns
- A list of columns to treat as
geometry. - If
with_types_srids, return a dictionary mapping columns to tuples in the form (type, SRID).
def
get_geometry_cols_types(df: pandas.DataFrame) -> Dict[str, str]:
1386def get_geometry_cols_types(df: 'pd.DataFrame') -> Dict[str, str]: 1387 """ 1388 Return a dtypes dictionary mapping columns to specific geometry types (type, srid). 1389 """ 1390 geometry_cols_types_srids = get_geometry_cols(df, with_types_srids=True) 1391 new_cols_types = {} 1392 for col, (geometry_type, srid) in geometry_cols_types_srids.items(): 1393 new_dtype = "geometry" 1394 modifier = "" 1395 if not srid and geometry_type.lower() == 'geometry': 1396 new_cols_types[col] = new_dtype 1397 continue 1398 1399 modifier = "[" 1400 if geometry_type.lower() != 'geometry': 1401 modifier += f"{geometry_type}" 1402 1403 if srid: 1404 if modifier != '[': 1405 modifier += ", " 1406 modifier += f"{srid}" 1407 modifier += "]" 1408 new_cols_types[col] = f"{new_dtype}{modifier}" 1409 return new_cols_types
Return a dtypes dictionary mapping columns to specific geometry types (type, srid).
def
get_special_cols(df: pandas.DataFrame) -> Dict[str, str]:
1412def get_special_cols(df: 'pd.DataFrame') -> Dict[str, str]: 1413 """ 1414 Return a dtypes dictionary mapping special columns to their dtypes. 1415 """ 1416 return { 1417 **{col: 'json' for col in get_json_cols(df)}, 1418 **{col: 'uuid' for col in get_uuid_cols(df)}, 1419 **{col: 'bytes' for col in get_bytes_cols(df)}, 1420 **{col: 'bool' for col in get_bool_cols(df)}, 1421 **{col: 'numeric' for col in get_numeric_cols(df)}, 1422 **{col: 'date' for col in get_date_cols(df)}, 1423 **get_datetime_cols_types(df), 1424 **get_geometry_cols_types(df), 1425 }
Return a dtypes dictionary mapping special columns to their dtypes.
def
enforce_dtypes( df: pandas.DataFrame, dtypes: Dict[str, str], explicit_dtypes: Optional[Dict[str, str]] = None, safe_copy: bool = True, coerce_numeric: bool = False, coerce_timezone: bool = True, strip_timezone: bool = False, as_polars: bool = False, debug: bool = False) -> pandas.DataFrame:
1537def enforce_dtypes( 1538 df: 'pd.DataFrame', 1539 dtypes: Dict[str, str], 1540 explicit_dtypes: Optional[Dict[str, str]] = None, 1541 safe_copy: bool = True, 1542 coerce_numeric: bool = False, 1543 coerce_timezone: bool = True, 1544 strip_timezone: bool = False, 1545 as_polars: bool = False, 1546 debug: bool = False, 1547) -> 'pd.DataFrame': 1548 """ 1549 Enforce the `dtypes` dictionary on a DataFrame. 1550 1551 Parameters 1552 ---------- 1553 df: pd.DataFrame 1554 The DataFrame on which to enforce dtypes. 1555 1556 dtypes: Dict[str, str] 1557 The data types to attempt to enforce on the DataFrame. 1558 1559 explicit_dtypes: Optional[Dict[str, str]], default None 1560 If provided, automatic dtype coersion will respect explicitly configured 1561 dtypes (`int`, `float`, `numeric`). 1562 1563 safe_copy: bool, default True 1564 If `True`, create a copy before comparing and modifying the dataframes. 1565 Setting to `False` may mutate the DataFrames. 1566 See `meerschaum.utils.dataframe.filter_unseen_df`. 1567 1568 coerce_numeric: bool, default False 1569 If `True`, convert float and int collisions to numeric. 1570 1571 coerce_timezone: bool, default True 1572 If `True`, convert datetimes to UTC. 1573 1574 strip_timezone: bool, default False 1575 If `coerce_timezone` and `strip_timezone` are `True`, 1576 remove timezone information from datetimes. 1577 1578 as_polars: bool, default False 1579 If `True`, return a Polars DataFrame. Supported Arrow-native dtypes are enforced 1580 with Polars; unsupported schemas use the established Pandas path before conversion. 1581 1582 debug: bool, default False 1583 Verbosity toggle. 1584 1585 Returns 1586 ------- 1587 The Pandas DataFrame with the types enforced. 1588 """ 1589 import json 1590 import functools 1591 from meerschaum.utils.debug import dprint 1592 from meerschaum.utils.formatting import pprint 1593 from meerschaum.utils.dtypes import ( 1594 MRSM_ALIAS_DTYPES, 1595 are_dtypes_equal, 1596 to_pandas_dtype, 1597 is_dtype_numeric, 1598 attempt_cast_to_numeric, 1599 attempt_cast_to_uuid, 1600 attempt_cast_to_bytes, 1601 attempt_cast_to_geometry, 1602 coerce_timezone as _coerce_timezone, 1603 get_geometry_type_srid, 1604 ) 1605 from meerschaum.utils.dtypes.sql import get_numeric_precision_scale 1606 pandas = mrsm.attempt_import('pandas') 1607 is_dask = 'dask' in df.__module__ 1608 normalized_dtypes = { 1609 col: MRSM_ALIAS_DTYPES.get(str(typ), str(typ)).lower() 1610 for col, typ in dtypes.items() 1611 } 1612 declared_json_cols = [col for col, typ in normalized_dtypes.items() if typ == 'json'] 1613 fallback_dtypes = { 1614 col: dtypes[col] 1615 for col, typ in normalized_dtypes.items() 1616 if ( 1617 typ in ('uuid', 'object') 1618 or typ.startswith(('geometry', 'geography')) 1619 or (typ.startswith('numeric') and None in get_numeric_precision_scale(None, typ)) 1620 ) 1621 } 1622 native_dtypes = { 1623 col: typ 1624 for col, typ in dtypes.items() 1625 if col not in fallback_dtypes 1626 } 1627 df_cols = df.collect_schema().names() if hasattr(df, 'collect_schema') else df.columns 1628 untyped_cols = [col for col in df_cols if col not in normalized_dtypes] 1629 polars_df = ( 1630 _enforce_dtypes_with_polars( 1631 df, 1632 native_dtypes, 1633 strip_timezone=(strip_timezone if coerce_timezone else False), 1634 ) 1635 if native_dtypes and not fallback_dtypes and not untyped_cols 1636 else None 1637 ) 1638 if polars_df is not None: 1639 return polars_df if as_polars else to_pandas(polars_df) 1640 1641 if native_dtypes and (fallback_dtypes or untyped_cols): 1642 df = to_pandas(df) 1643 native_cols = [col for col in native_dtypes if col in df.columns] 1644 native_df = _enforce_dtypes_with_polars( 1645 df[native_cols], 1646 native_dtypes, 1647 strip_timezone=(strip_timezone if coerce_timezone else False), 1648 ) if native_cols else None 1649 if native_df is not None: 1650 if safe_copy: 1651 df = df.copy() 1652 native_pd_df = to_pandas(native_df) 1653 for col in native_cols: 1654 df[col] = native_pd_df[col].array 1655 dtypes = fallback_dtypes 1656 safe_copy = False 1657 1658 df = to_pandas(df) 1659 if safe_copy: 1660 df = df.copy() 1661 if len(df.columns) == 0: 1662 if debug: 1663 dprint("Incoming DataFrame has no columns. Skipping enforcement...") 1664 return to_polars(df) if as_polars else df 1665 1666 explicit_dtypes = explicit_dtypes or {} 1667 pipe_pandas_dtypes = { 1668 col: to_pandas_dtype(typ) 1669 for col, typ in dtypes.items() 1670 } 1671 json_cols = [ 1672 col 1673 for col, typ in dtypes.items() 1674 if typ == 'json' 1675 ] 1676 numeric_cols = [ 1677 col 1678 for col, typ in dtypes.items() 1679 if typ.startswith('numeric') 1680 ] 1681 geometry_cols_types_srids = { 1682 col: get_geometry_type_srid(typ, default_srid=0) 1683 for col, typ in dtypes.items() 1684 if typ.startswith('geometry') or typ.startswith('geography') 1685 } 1686 uuid_cols = [ 1687 col 1688 for col, typ in dtypes.items() 1689 if typ == 'uuid' 1690 ] 1691 bytes_cols = [ 1692 col 1693 for col, typ in dtypes.items() 1694 if typ == 'bytes' 1695 ] 1696 datetime_cols = [ 1697 col 1698 for col, typ in dtypes.items() 1699 if are_dtypes_equal(typ, 'datetime') 1700 ] 1701 df_numeric_cols = get_numeric_cols(df) 1702 if debug: 1703 dprint("Desired data types:") 1704 pprint(dtypes) 1705 dprint("Data types for incoming DataFrame:") 1706 pprint({_col: str(_typ) for _col, _typ in df.dtypes.items()}) 1707 1708 if json_cols and len(df) > 0: 1709 if debug: 1710 dprint(f"Checking columns for JSON encoding: {json_cols}") 1711 for col in json_cols: 1712 if col in df.columns: 1713 try: 1714 df[col] = df[col].apply( 1715 ( 1716 lambda x: ( 1717 json.loads(x) 1718 if isinstance(x, str) 1719 else x 1720 ) 1721 ) 1722 ) 1723 except Exception as e: 1724 if debug: 1725 dprint(f"Unable to parse column '{col}' as JSON:\n{e}") 1726 1727 if numeric_cols: 1728 if debug: 1729 dprint(f"Checking for numerics: {numeric_cols}") 1730 for col in numeric_cols: 1731 precision, scale = get_numeric_precision_scale(None, dtypes.get(col, '')) 1732 if col in df.columns: 1733 try: 1734 df[col] = df[col].apply( 1735 functools.partial( 1736 attempt_cast_to_numeric, 1737 quantize=True, 1738 precision=precision, 1739 scale=scale, 1740 ) 1741 ) 1742 except Exception as e: 1743 if debug: 1744 dprint(f"Unable to parse column '{col}' as NUMERIC:\n{e}") 1745 1746 if uuid_cols: 1747 if debug: 1748 dprint(f"Checking for UUIDs: {uuid_cols}") 1749 for col in uuid_cols: 1750 if col in df.columns: 1751 try: 1752 df[col] = df[col].apply(attempt_cast_to_uuid) 1753 except Exception as e: 1754 if debug: 1755 dprint(f"Unable to parse column '{col}' as UUID:\n{e}") 1756 1757 if bytes_cols: 1758 if debug: 1759 dprint(f"Checking for bytes: {bytes_cols}") 1760 for col in bytes_cols: 1761 if col in df.columns: 1762 try: 1763 df[col] = df[col].apply(attempt_cast_to_bytes) 1764 except Exception as e: 1765 if debug: 1766 dprint(f"Unable to parse column '{col}' as bytes:\n{e}") 1767 1768 if datetime_cols and coerce_timezone: 1769 if debug: 1770 dprint(f"Checking for datetime conversion: {datetime_cols}") 1771 for col in datetime_cols: 1772 if col in df.columns: 1773 if not strip_timezone and 'utc' in str(df.dtypes[col]).lower(): 1774 if debug: 1775 dprint(f"Skip UTC coersion for column '{col}' ({str(df[col].dtype)}).") 1776 continue 1777 if strip_timezone and ',' not in str(df.dtypes[col]): 1778 if debug: 1779 dprint( 1780 f"Skip UTC coersion (stripped) for column '{col}' " 1781 f"({str(df[col].dtype)})." 1782 ) 1783 continue 1784 1785 if debug: 1786 dprint( 1787 f"Data type for column '{col}' before timezone coersion: " 1788 f"{str(df[col].dtype)}" 1789 ) 1790 1791 df[col] = _coerce_timezone(df[col], strip_utc=strip_timezone) 1792 if debug: 1793 dprint( 1794 f"Data type for column '{col}' after timezone coersion: " 1795 f"{str(df[col].dtype)}" 1796 ) 1797 1798 if geometry_cols_types_srids: 1799 geopandas = mrsm.attempt_import('geopandas') 1800 if debug: 1801 dprint(f"Checking for geometry: {list(geometry_cols_types_srids)}") 1802 parsed_geom_cols = [] 1803 for col in geometry_cols_types_srids: 1804 if col not in df.columns: 1805 continue 1806 try: 1807 df[col] = attempt_cast_to_geometry(df[col]) 1808 parsed_geom_cols.append(col) 1809 except Exception as e: 1810 import traceback 1811 traceback.print_exc() 1812 if debug: 1813 dprint(f"Unable to parse column '{col}' as geometry:\n{e}") 1814 1815 if parsed_geom_cols: 1816 if debug: 1817 dprint(f"Converting to GeoDataFrame (geometry column: '{parsed_geom_cols[0]}')...") 1818 try: 1819 _, default_srid = geometry_cols_types_srids[parsed_geom_cols[0]] 1820 df = geopandas.GeoDataFrame(df, geometry=parsed_geom_cols[0], crs=default_srid) 1821 for col, (_, srid) in geometry_cols_types_srids.items(): 1822 if srid: 1823 if debug: 1824 dprint(f"Setting '{col}' to SRID '{srid}'...") 1825 _ = df[col].set_crs(srid) 1826 if parsed_geom_cols[0] not in df.columns: 1827 df.rename_geometry(parsed_geom_cols[0], inplace=True) 1828 except (ValueError, TypeError): 1829 import traceback 1830 dprint(f"Failed to cast to GeoDataFrame:\n{traceback.format_exc()}") 1831 1832 df_dtypes = {c: str(t) for c, t in df.dtypes.items()} 1833 if are_dtypes_equal(df_dtypes, pipe_pandas_dtypes): 1834 if debug: 1835 dprint("Data types match. Exiting enforcement...") 1836 return ( 1837 to_polars( 1838 df, 1839 geometry_cols_types_srids=geometry_cols_types_srids, 1840 json_cols=declared_json_cols, 1841 ) 1842 if as_polars 1843 else df 1844 ) 1845 1846 common_dtypes = {} 1847 common_diff_dtypes = {} 1848 for col, typ in pipe_pandas_dtypes.items(): 1849 if col in df_dtypes: 1850 common_dtypes[col] = typ 1851 if not are_dtypes_equal(typ, df_dtypes[col]): 1852 common_diff_dtypes[col] = df_dtypes[col] 1853 1854 if debug: 1855 dprint("Common columns with different dtypes:") 1856 pprint(common_diff_dtypes) 1857 1858 detected_dt_cols = {} 1859 for col, typ in common_diff_dtypes.items(): 1860 if 'datetime' in typ and 'datetime' in common_dtypes[col]: 1861 df_dtypes[col] = typ 1862 detected_dt_cols[col] = (common_dtypes[col], common_diff_dtypes[col]) 1863 for col in detected_dt_cols: 1864 del common_diff_dtypes[col] 1865 1866 if debug: 1867 dprint("Common columns with different dtypes (after dates):") 1868 pprint(common_diff_dtypes) 1869 1870 if are_dtypes_equal(df_dtypes, pipe_pandas_dtypes): 1871 if debug: 1872 dprint( 1873 "The incoming DataFrame has mostly the same types, skipping enforcement." 1874 + "The only detected difference was in the following datetime columns." 1875 ) 1876 pprint(detected_dt_cols) 1877 return ( 1878 to_polars( 1879 df, 1880 geometry_cols_types_srids=geometry_cols_types_srids, 1881 json_cols=declared_json_cols, 1882 ) 1883 if as_polars 1884 else df 1885 ) 1886 1887 for col, typ in {k: v for k, v in common_diff_dtypes.items()}.items(): 1888 previous_typ = common_dtypes[col] 1889 mixed_numeric_types = (is_dtype_numeric(typ) and is_dtype_numeric(previous_typ)) 1890 explicitly_float = are_dtypes_equal(explicit_dtypes.get(col, 'object'), 'float') 1891 explicitly_int = are_dtypes_equal(explicit_dtypes.get(col, 'object'), 'int') 1892 explicitly_numeric = explicit_dtypes.get(col, 'object').startswith('numeric') 1893 all_nan = ( 1894 df[col].isnull().all() 1895 if mixed_numeric_types and coerce_numeric and not (explicitly_float or explicitly_int) 1896 else None 1897 ) 1898 cast_to_numeric = explicitly_numeric or ( 1899 ( 1900 col in df_numeric_cols 1901 or ( 1902 mixed_numeric_types 1903 and not (explicitly_float or explicitly_int) 1904 and not all_nan 1905 and coerce_numeric 1906 ) 1907 ) 1908 ) 1909 1910 if debug and (explicitly_numeric or df_numeric_cols or mixed_numeric_types): 1911 from meerschaum.utils.formatting import make_header 1912 msg = ( 1913 make_header(f"Coercing column '{col}' to numeric:", left_pad=0) 1914 + "\n" 1915 + f" Previous type: {previous_typ}\n" 1916 + f" Current type: {typ if col not in df_numeric_cols else 'Decimal'}" 1917 + ("\n Column is explicitly numeric." if explicitly_numeric else "") 1918 ) if cast_to_numeric else ( 1919 f"Will not coerce column '{col}' to numeric.\n" 1920 f" Numeric columns in dataframe: {df_numeric_cols}\n" 1921 f" Mixed numeric types: {mixed_numeric_types}\n" 1922 f" Explicitly float: {explicitly_float}\n" 1923 f" Explicitly int: {explicitly_int}\n" 1924 f" All NaN: {all_nan}\n" 1925 f" Coerce numeric: {coerce_numeric}" 1926 ) 1927 dprint(msg) 1928 1929 if cast_to_numeric: 1930 common_dtypes[col] = attempt_cast_to_numeric 1931 common_diff_dtypes[col] = attempt_cast_to_numeric 1932 1933 for d in common_diff_dtypes: 1934 t = common_dtypes[d] 1935 if debug: 1936 dprint(f"Casting column {d} to dtype {t}.") 1937 try: 1938 df[d] = ( 1939 df[d].apply(t) 1940 if callable(t) 1941 else df[d].astype(t) 1942 ) 1943 except Exception as e: 1944 if debug: 1945 dprint(f"Encountered an error when casting column {d} to type {t}:\n{e}\ndf:\n{df}") 1946 if 'int' in str(t).lower(): 1947 try: 1948 df[d] = df[d].astype('float64').astype(t) 1949 except Exception: 1950 if debug: 1951 dprint(f"Was unable to convert to float then {t}.") 1952 return ( 1953 to_polars( 1954 df, 1955 geometry_cols_types_srids=geometry_cols_types_srids, 1956 json_cols=declared_json_cols, 1957 ) 1958 if as_polars 1959 else df 1960 )
Enforce the dtypes dictionary on a DataFrame.
Parameters
- df (pd.DataFrame): The DataFrame on which to enforce dtypes.
- dtypes (Dict[str, str]): The data types to attempt to enforce on the DataFrame.
- explicit_dtypes (Optional[Dict[str, str]], default None):
If provided, automatic dtype coersion will respect explicitly configured
dtypes (
int,float,numeric). - safe_copy (bool, default True):
If
True, create a copy before comparing and modifying the dataframes. Setting toFalsemay mutate the DataFrames. Seemeerschaum.utils.dataframe.filter_unseen_df. - coerce_numeric (bool, default False):
If
True, convert float and int collisions to numeric. - coerce_timezone (bool, default True):
If
True, convert datetimes to UTC. - strip_timezone (bool, default False):
If
coerce_timezoneandstrip_timezoneareTrue, remove timezone information from datetimes. - as_polars (bool, default False):
If
True, return a Polars DataFrame. Supported Arrow-native dtypes are enforced with Polars; unsupported schemas use the established Pandas path before conversion. - debug (bool, default False): Verbosity toggle.
Returns
- The Pandas DataFrame with the types enforced.
def
get_datetime_bound_from_df( df: Union[pandas.DataFrame, Dict[str, List[Any]], List[Dict[str, Any]]], datetime_column: str, minimum: bool = True) -> Union[int, datetime.datetime, NoneType]:
1963def get_datetime_bound_from_df( 1964 df: Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]], 1965 datetime_column: str, 1966 minimum: bool = True, 1967) -> Union[int, datetime, None]: 1968 """ 1969 Return the minimum or maximum datetime (or integer) from a DataFrame. 1970 1971 Parameters 1972 ---------- 1973 df: Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]] 1974 The DataFrame, list, or dict which contains the range axis. 1975 1976 datetime_column: str 1977 The name of the datetime (or int) column. 1978 1979 minimum: bool 1980 Whether to return the minimum (default) or maximum value. 1981 1982 Returns 1983 ------- 1984 The minimum or maximum datetime value in the dataframe, or `None`. 1985 """ 1986 from meerschaum.utils.dtypes import to_datetime, value_is_null 1987 1988 if df is None: 1989 return None 1990 if not datetime_column: 1991 return None 1992 1993 def compare(a, b): 1994 if a is None: 1995 return b 1996 if b is None: 1997 return a 1998 if minimum: 1999 return a if a < b else b 2000 return a if a > b else b 2001 2002 if isinstance(df, list): 2003 if len(df) == 0: 2004 return None 2005 best_yet = df[0].get(datetime_column, None) 2006 for doc in df: 2007 val = doc.get(datetime_column, None) 2008 best_yet = compare(best_yet, val) 2009 return best_yet 2010 2011 if isinstance(df, dict): 2012 if datetime_column not in df: 2013 return None 2014 best_yet = df[datetime_column][0] 2015 for val in df[datetime_column]: 2016 best_yet = compare(best_yet, val) 2017 return best_yet 2018 2019 if 'DataFrame' in str(type(df)): 2020 from meerschaum.utils.dtypes import are_dtypes_equal 2021 pandas = mrsm.attempt_import('pandas') 2022 is_dask = 'dask' in df.__module__ 2023 2024 if datetime_column not in df.columns: 2025 return None 2026 2027 try: 2028 dt_val = ( 2029 df[datetime_column].min(skipna=True) 2030 if minimum 2031 else df[datetime_column].max(skipna=True) 2032 ) 2033 except Exception: 2034 dt_val = pandas.NA 2035 if is_dask and dt_val is not None and dt_val is not pandas.NA: 2036 dt_val = dt_val.compute() 2037 2038 return ( 2039 to_datetime(dt_val, as_pydatetime=True) 2040 if are_dtypes_equal(str(type(dt_val)), 'datetime') 2041 else (dt_val if not value_is_null(dt_val) else None) 2042 ) 2043 2044 return None
Return the minimum or maximum datetime (or integer) from a DataFrame.
Parameters
- df (Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]]): The DataFrame, list, or dict which contains the range axis.
- datetime_column (str): The name of the datetime (or int) column.
- minimum (bool): Whether to return the minimum (default) or maximum value.
Returns
- The minimum or maximum datetime value in the dataframe, or
None.
def
get_unique_index_values( df: Union[pandas.DataFrame, Dict[str, List[Any]], List[Dict[str, Any]]], indices: List[str]) -> Dict[str, List[Any]]:
2047def get_unique_index_values( 2048 df: Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]], 2049 indices: List[str], 2050) -> Dict[str, List[Any]]: 2051 """ 2052 Return a dictionary of the unique index values in a DataFrame. 2053 2054 Parameters 2055 ---------- 2056 df: Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]] 2057 The dataframe (or list or dict) which contains index values. 2058 2059 indices: List[str] 2060 The list of index columns. 2061 2062 Returns 2063 ------- 2064 A dictionary mapping indices to unique values. 2065 """ 2066 if df is None: 2067 return {} 2068 if 'dataframe' in str(type(df)).lower(): 2069 pandas = mrsm.attempt_import('pandas') 2070 return { 2071 col: list({ 2072 (val if val is not pandas.NA else None) 2073 for val in df[col].unique() 2074 }) 2075 for col in indices 2076 if col in df.columns 2077 } 2078 2079 unique_indices = defaultdict(lambda: set()) 2080 if isinstance(df, list): 2081 for doc in df: 2082 for index in indices: 2083 if index in doc: 2084 unique_indices[index].add(doc[index]) 2085 2086 elif isinstance(df, dict): 2087 for index in indices: 2088 if index in df: 2089 unique_indices[index] = unique_indices[index].union(set(df[index])) 2090 2091 return {key: list(val) for key, val in unique_indices.items()}
Return a dictionary of the unique index values in a DataFrame.
Parameters
- df (Union['pd.DataFrame', Dict[str, List[Any]], List[Dict[str, Any]]]): The dataframe (or list or dict) which contains index values.
- indices (List[str]): The list of index columns.
Returns
- A dictionary mapping indices to unique values.
def
df_is_chunk_generator(df: Any) -> bool:
2094def df_is_chunk_generator(df: Any) -> bool: 2095 """ 2096 Determine whether to treat `df` as a chunk generator. 2097 2098 Note this should only be used in a context where generators are expected, 2099 as it will return `True` for any iterable. 2100 2101 Parameters 2102 ---------- 2103 The DataFrame or chunk generator to evaluate. 2104 2105 Returns 2106 ------- 2107 A `bool` indicating whether to treat `df` as a generator. 2108 """ 2109 return ( 2110 not isinstance(df, (dict, list, str)) 2111 and 'DataFrame' not in str(type(df)) 2112 and isinstance(df, (Generator, Iterable, Iterator)) 2113 )
Determine whether to treat df as a chunk generator.
Note this should only be used in a context where generators are expected,
as it will return True for any iterable.
Parameters
- The DataFrame or chunk generator to evaluate.
Returns
- A
boolindicating whether to treatdfas a generator.
def
chunksize_to_npartitions(chunksize: Optional[int]) -> int:
2116def chunksize_to_npartitions(chunksize: Optional[int]) -> int: 2117 """ 2118 Return the Dask `npartitions` value for a given `chunksize`. 2119 """ 2120 if chunksize == -1: 2121 from meerschaum.config import get_config 2122 chunksize = get_config('system', 'connectors', 'sql', 'chunksize') 2123 if chunksize is None: 2124 return 1 2125 return -1 * chunksize
Return the Dask npartitions value for a given chunksize.
def
df_from_literal( pipe: Optional[meerschaum.Pipe] = None, literal: Optional[str] = None, debug: bool = False) -> pandas.DataFrame:
2128def df_from_literal( 2129 pipe: Optional[mrsm.Pipe] = None, 2130 literal: Optional[str] = None, 2131 debug: bool = False 2132) -> 'pd.DataFrame': 2133 """ 2134 Construct a dataframe from a literal value, using the pipe's datetime and value column names. 2135 2136 Parameters 2137 ---------- 2138 pipe: Optional['meerschaum.Pipe'], default None 2139 The pipe which will consume the literal value. 2140 2141 Returns 2142 ------- 2143 A 1-row pandas DataFrame from with the current UTC timestamp as the datetime columns 2144 and the literal as the value. 2145 """ 2146 from meerschaum.utils.packages import import_pandas 2147 from meerschaum.utils.warnings import error, warn 2148 from meerschaum.utils.debug import dprint 2149 from meerschaum.utils.dtypes import get_current_timestamp 2150 2151 if pipe is None or literal is None: 2152 error("Please provide a Pipe and a literal value") 2153 2154 dt_col = pipe.columns.get( 2155 'datetime', 2156 mrsm.get_config('pipes', 'autotime', 'column_name_if_datetime_missing') 2157 ) 2158 val_col = pipe.get_val_column(debug=debug) 2159 2160 val = literal 2161 if isinstance(literal, str): 2162 if debug: 2163 dprint(f"Received literal string: '{literal}'") 2164 import ast 2165 try: 2166 val = ast.literal_eval(literal) 2167 except Exception: 2168 warn( 2169 "Failed to parse value from string:\n" + f"{literal}" + 2170 "\n\nWill cast as a string instead."\ 2171 ) 2172 val = literal 2173 2174 now = get_current_timestamp(pipe.precision) 2175 pd = import_pandas() 2176 return pd.DataFrame({dt_col: [now], val_col: [val]})
Construct a dataframe from a literal value, using the pipe's datetime and value column names.
Parameters
- pipe (Optional['meerschaum.Pipe'], default None): The pipe which will consume the literal value.
Returns
- A 1-row pandas DataFrame from with the current UTC timestamp as the datetime columns
- and the literal as the value.
def
get_first_valid_dask_partition( ddf: dask.dataframe.dask_expr._collection.DataFrame) -> Optional[pandas.DataFrame]:
2179def get_first_valid_dask_partition(ddf: 'dask.dataframe.DataFrame') -> Union['pd.DataFrame', None]: 2180 """ 2181 Return the first valid Dask DataFrame partition (if possible). 2182 """ 2183 pdf = None 2184 for partition in ddf.partitions: 2185 try: 2186 pdf = partition.compute() 2187 except Exception: 2188 continue 2189 if len(pdf) > 0: 2190 return pdf 2191 _ = mrsm.attempt_import('partd', lazy=False) 2192 return ddf.compute()
Return the first valid Dask DataFrame partition (if possible).
def
query_df( df: pandas.DataFrame, params: Optional[Dict[str, Any]] = None, begin: Union[datetime.datetime, int, NoneType] = None, end: Union[datetime.datetime, int, NoneType] = None, datetime_column: Optional[str] = None, select_columns: Optional[List[str]] = None, omit_columns: Optional[List[str]] = None, inplace: bool = False, reset_index: bool = False, coerce_types: bool = False, debug: bool = False) -> pandas.DataFrame:
2195def query_df( 2196 df: 'pd.DataFrame', 2197 params: Optional[Dict[str, Any]] = None, 2198 begin: Union[datetime, int, None] = None, 2199 end: Union[datetime, int, None] = None, 2200 datetime_column: Optional[str] = None, 2201 select_columns: Optional[List[str]] = None, 2202 omit_columns: Optional[List[str]] = None, 2203 inplace: bool = False, 2204 reset_index: bool = False, 2205 coerce_types: bool = False, 2206 debug: bool = False, 2207) -> 'pd.DataFrame': 2208 """ 2209 Query the dataframe with the params dictionary. 2210 2211 Parameters 2212 ---------- 2213 df: pd.DataFrame 2214 The DataFrame to query against. 2215 2216 params: Optional[Dict[str, Any]], default None 2217 The parameters dictionary to use for the query. 2218 2219 begin: Union[datetime, int, None], default None 2220 If `begin` and `datetime_column` are provided, only return rows with a timestamp 2221 greater than or equal to this value. 2222 2223 end: Union[datetime, int, None], default None 2224 If `begin` and `datetime_column` are provided, only return rows with a timestamp 2225 less than this value. 2226 2227 datetime_column: Optional[str], default None 2228 A `datetime_column` must be provided to use `begin` and `end`. 2229 2230 select_columns: Optional[List[str]], default None 2231 If provided, only return these columns. 2232 2233 omit_columns: Optional[List[str]], default None 2234 If provided, do not include these columns in the result. 2235 2236 inplace: bool, default False 2237 If `True`, modify the DataFrame inplace rather than creating a new DataFrame. 2238 2239 reset_index: bool, default False 2240 If `True`, reset the index in the resulting DataFrame. 2241 2242 coerce_types: bool, default False 2243 If `True`, cast the dataframe and parameters as strings before querying. 2244 2245 Returns 2246 ------- 2247 A Pandas DataFrame query result. 2248 """ 2249 2250 def _process_select_columns(_df): 2251 if not select_columns: 2252 return 2253 for col in list(_df.columns): 2254 if col not in select_columns: 2255 del _df[col] 2256 2257 def _process_omit_columns(_df): 2258 if not omit_columns: 2259 return 2260 for col in list(_df.columns): 2261 if col in omit_columns: 2262 del _df[col] 2263 2264 if not params and not begin and not end: 2265 if not inplace: 2266 df = df.copy() 2267 _process_select_columns(df) 2268 _process_omit_columns(df) 2269 return df 2270 2271 from meerschaum.utils.debug import dprint 2272 from meerschaum.utils.misc import get_in_ex_params 2273 from meerschaum.utils.warnings import warn 2274 from meerschaum.utils.dtypes import are_dtypes_equal, value_is_null 2275 dateutil_parser = mrsm.attempt_import('dateutil.parser') 2276 pandas = mrsm.attempt_import('pandas') 2277 NA = pandas.NA 2278 2279 if params: 2280 proto_in_ex_params = get_in_ex_params(params) 2281 for key, (proto_in_vals, proto_ex_vals) in proto_in_ex_params.items(): 2282 if proto_ex_vals: 2283 coerce_types = True 2284 break 2285 params = params.copy() 2286 for key, val in {k: v for k, v in params.items()}.items(): 2287 if isinstance(val, (list, tuple, set)) or hasattr(val, 'astype'): 2288 if None in val: 2289 val = [item for item in val if item is not None] + [NA] 2290 params[key] = val 2291 if coerce_types: 2292 params[key] = [str(x) for x in val] 2293 else: 2294 if value_is_null(val): 2295 val = NA 2296 params[key] = NA 2297 if coerce_types: 2298 params[key] = str(val) 2299 2300 dtypes = {col: str(typ) for col, typ in df.dtypes.items()} 2301 2302 if inplace: 2303 df.fillna(NA, inplace=True) 2304 else: 2305 df = df.infer_objects().fillna(NA) 2306 2307 if isinstance(begin, str): 2308 begin = dateutil_parser.parse(begin) 2309 if isinstance(end, str): 2310 end = dateutil_parser.parse(end) 2311 2312 if begin is not None or end is not None: 2313 if not datetime_column or datetime_column not in df.columns: 2314 warn( 2315 f"The datetime column '{datetime_column}' is not present in the Dataframe, " 2316 + "ignoring begin and end...", 2317 ) 2318 begin, end = None, None 2319 2320 if debug: 2321 dprint(f"Querying dataframe:\n{params=} {begin=} {end=} {datetime_column=}") 2322 2323 if datetime_column and (begin is not None or end is not None): 2324 if debug: 2325 dprint("Checking for datetime column compatability.") 2326 2327 from meerschaum.utils.dtypes import coerce_timezone 2328 df_is_dt = are_dtypes_equal(str(df.dtypes[datetime_column]), 'datetime') 2329 begin_is_int = are_dtypes_equal(str(type(begin)), 'int') 2330 end_is_int = are_dtypes_equal(str(type(end)), 'int') 2331 2332 if df_is_dt: 2333 df_tz = ( 2334 getattr(df[datetime_column].dt, 'tz', None) 2335 if hasattr(df[datetime_column], 'dt') 2336 else None 2337 ) 2338 2339 if begin_is_int: 2340 begin = datetime.fromtimestamp(int(begin), timezone.utc).replace(tzinfo=None) 2341 if debug: 2342 dprint(f"`begin` will be cast to '{begin}'.") 2343 if end_is_int: 2344 end = datetime.fromtimestamp(int(end), timezone.utc).replace(tzinfo=None) 2345 if debug: 2346 dprint(f"`end` will be cast to '{end}'.") 2347 2348 begin = coerce_timezone(begin, strip_utc=(df_tz is None)) if begin is not None else None 2349 end = coerce_timezone(end, strip_utc=(df_tz is None)) if begin is not None else None 2350 2351 in_ex_params = get_in_ex_params(params) 2352 2353 masks = [ 2354 ( 2355 (df[datetime_column] >= begin) 2356 if begin is not None and datetime_column 2357 else True 2358 ) & ( 2359 (df[datetime_column] < end) 2360 if end is not None and datetime_column 2361 else True 2362 ) 2363 ] 2364 2365 masks.extend([ 2366 ( 2367 ( 2368 (df[col] if not coerce_types else df[col].astype(str)).isin(in_vals) 2369 if in_vals 2370 else True 2371 ) & ( 2372 ~(df[col] if not coerce_types else df[col].astype(str)).isin(ex_vals) 2373 if ex_vals 2374 else True 2375 ) 2376 ) 2377 for col, (in_vals, ex_vals) in in_ex_params.items() 2378 if col in df.columns 2379 ]) 2380 query_mask = masks[0] 2381 for mask in masks[1:]: 2382 query_mask = query_mask & mask 2383 2384 original_cols = df.columns 2385 2386 ### NOTE: We must cast bool columns to `boolean[pyarrow]` 2387 ### to allow for `<NA>` values. 2388 bool_cols = [ 2389 col 2390 for col, typ in df.dtypes.items() 2391 if are_dtypes_equal(str(typ), 'bool') 2392 ] 2393 for col in bool_cols: 2394 df[col] = df[col].astype('boolean[pyarrow]') 2395 2396 if not isinstance(query_mask, bool): 2397 df['__mrsm_mask'] = ( 2398 query_mask.astype('boolean[pyarrow]') 2399 if hasattr(query_mask, 'astype') 2400 else query_mask 2401 ) 2402 2403 if inplace: 2404 df.where(query_mask, other=NA, inplace=True) 2405 df.dropna(how='all', inplace=True) 2406 result_df = df 2407 else: 2408 result_df = df.where(query_mask, other=NA) 2409 result_df.dropna(how='all', inplace=True) 2410 2411 else: 2412 result_df = df 2413 2414 if '__mrsm_mask' in df.columns: 2415 del df['__mrsm_mask'] 2416 if '__mrsm_mask' in result_df.columns: 2417 del result_df['__mrsm_mask'] 2418 2419 if reset_index: 2420 result_df.reset_index(drop=True, inplace=True) 2421 2422 result_df = enforce_dtypes( 2423 result_df, 2424 dtypes, 2425 safe_copy=False, 2426 debug=debug, 2427 coerce_numeric=False, 2428 coerce_timezone=False, 2429 ) 2430 2431 if select_columns == ['*']: 2432 select_columns = None 2433 2434 if not select_columns and not omit_columns: 2435 return result_df[original_cols] 2436 2437 _process_select_columns(result_df) 2438 _process_omit_columns(result_df) 2439 2440 return result_df
Query the dataframe with the params dictionary.
Parameters
- df (pd.DataFrame): The DataFrame to query against.
- params (Optional[Dict[str, Any]], default None): The parameters dictionary to use for the query.
- begin (Union[datetime, int, None], default None):
If
beginanddatetime_columnare provided, only return rows with a timestamp greater than or equal to this value. - end (Union[datetime, int, None], default None):
If
beginanddatetime_columnare provided, only return rows with a timestamp less than this value. - datetime_column (Optional[str], default None):
A
datetime_columnmust be provided to usebeginandend. - select_columns (Optional[List[str]], default None): If provided, only return these columns.
- omit_columns (Optional[List[str]], default None): If provided, do not include these columns in the result.
- inplace (bool, default False):
If
True, modify the DataFrame inplace rather than creating a new DataFrame. - reset_index (bool, default False):
If
True, reset the index in the resulting DataFrame. - coerce_types (bool, default False):
If
True, cast the dataframe and parameters as strings before querying.
Returns
- A Pandas DataFrame query result.
def
to_json( df: pandas.DataFrame, safe_copy: bool = True, orient: str = 'records', date_format: str = 'iso', date_unit: str = 'us', double_precision: int = 15, geometry_format: str = 'geojson', **kwargs: Any) -> str:
2443def to_json( 2444 df: 'pd.DataFrame', 2445 safe_copy: bool = True, 2446 orient: str = 'records', 2447 date_format: str = 'iso', 2448 date_unit: str = 'us', 2449 double_precision: int = 15, 2450 geometry_format: str = 'geojson', 2451 **kwargs: Any 2452) -> str: 2453 """ 2454 Serialize the given dataframe as a JSON string. 2455 2456 Parameters 2457 ---------- 2458 df: pd.DataFrame 2459 The DataFrame to be serialized. 2460 2461 safe_copy: bool, default True 2462 If `False`, modify the DataFrame inplace. 2463 2464 date_format: str, default 'iso' 2465 The default format for timestamps. 2466 2467 date_unit: str, default 'us' 2468 The precision of the timestamps. 2469 2470 double_precision: int, default 15 2471 The number of decimal places to use when encoding floating point values (maximum 15). 2472 2473 geometry_format: str, default 'geojson' 2474 The serialization format for geometry data. 2475 Accepted values are `geojson`, `wkb_hex`, and `wkt`. 2476 2477 Returns 2478 ------- 2479 A JSON string. 2480 """ 2481 import warnings 2482 import functools 2483 from meerschaum.utils.packages import import_pandas 2484 from meerschaum.utils.dtypes import ( 2485 serialize_bytes, 2486 serialize_decimal, 2487 serialize_geometry, 2488 serialize_date, 2489 ) 2490 pd = import_pandas() 2491 uuid_cols = get_uuid_cols(df) 2492 bytes_cols = get_bytes_cols(df) 2493 numeric_cols = get_numeric_cols(df) 2494 date_cols = get_date_cols(df) 2495 geometry_cols = get_geometry_cols(df) 2496 geometry_cols_srids = { 2497 col: int((getattr(df[col].crs, 'srs', '') or '').split(':', maxsplit=1)[-1] or '0') 2498 for col in geometry_cols 2499 } if 'geodataframe' in str(type(df)).lower() else {} 2500 if safe_copy and bool(uuid_cols or bytes_cols or geometry_cols or numeric_cols): 2501 df = df.copy() 2502 if 'geodataframe' in str(type(df)).lower(): 2503 geometry_data = { 2504 col: df[col] 2505 for col in geometry_cols 2506 } 2507 df = pd.DataFrame({ 2508 col: df[col] 2509 for col in df.columns 2510 if col not in geometry_cols 2511 }) 2512 for col in geometry_cols: 2513 df[col] = pd.Series(ob for ob in geometry_data[col]) 2514 for col in uuid_cols: 2515 df[col] = df[col].astype(str) 2516 for col in bytes_cols: 2517 df[col] = df[col].apply(serialize_bytes) 2518 for col in numeric_cols: 2519 df[col] = df[col].apply(serialize_decimal) 2520 for col in date_cols: 2521 df[col] = df[col].apply(serialize_date) 2522 with warnings.catch_warnings(): 2523 warnings.simplefilter("ignore") 2524 for col in geometry_cols: 2525 srid = geometry_cols_srids.get(col, None) or None 2526 df[col] = pd.Series( 2527 serialize_geometry(val, geometry_format=geometry_format, srid=srid) 2528 for val in df[col] 2529 ) 2530 return df.infer_objects().fillna(pd.NA).to_json( 2531 date_format=date_format, 2532 date_unit=date_unit, 2533 double_precision=double_precision, 2534 orient=orient, 2535 **kwargs 2536 )
Serialize the given dataframe as a JSON string.
Parameters
- df (pd.DataFrame): The DataFrame to be serialized.
- safe_copy (bool, default True):
If
False, modify the DataFrame inplace. - date_format (str, default 'iso'): The default format for timestamps.
- date_unit (str, default 'us'): The precision of the timestamps.
- double_precision (int, default 15): The number of decimal places to use when encoding floating point values (maximum 15).
- geometry_format (str, default 'geojson'):
The serialization format for geometry data.
Accepted values are
geojson,wkb_hex, andwkt.
Returns
- A JSON string.
def
to_simple_lines(df: pandas.DataFrame) -> str:
2539def to_simple_lines(df: 'pd.DataFrame') -> str: 2540 """ 2541 Serialize a Pandas Dataframe as lines of simple dictionaries. 2542 2543 Parameters 2544 ---------- 2545 df: pd.DataFrame 2546 The dataframe to serialize into simple lines text. 2547 2548 Returns 2549 ------- 2550 A string of simple line dictionaries joined by newlines. 2551 """ 2552 from meerschaum.utils.misc import to_simple_dict 2553 if df is None or len(df) == 0: 2554 return '' 2555 2556 docs = df.to_dict(orient='records') 2557 return '\n'.join(to_simple_dict(doc) for doc in docs)
Serialize a Pandas Dataframe as lines of simple dictionaries.
Parameters
- df (pd.DataFrame): The dataframe to serialize into simple lines text.
Returns
- A string of simple line dictionaries joined by newlines.
def
parse_simple_lines(data: str) -> pandas.DataFrame:
2560def parse_simple_lines(data: str) -> 'pd.DataFrame': 2561 """ 2562 Parse simple lines text into a DataFrame. 2563 2564 Parameters 2565 ---------- 2566 data: str 2567 The simple lines text to parse into a DataFrame. 2568 2569 Returns 2570 ------- 2571 A dataframe containing the rows serialized in `data`. 2572 """ 2573 from meerschaum.utils.misc import string_to_dict 2574 from meerschaum.utils.packages import import_pandas 2575 pd = import_pandas() 2576 lines = data.splitlines() 2577 try: 2578 docs = [string_to_dict(line) for line in lines] 2579 df = pd.DataFrame(docs) 2580 except Exception: 2581 df = None 2582 2583 if df is None: 2584 raise ValueError("Cannot parse simple lines into a dataframe.") 2585 2586 return df
Parse simple lines text into a DataFrame.
Parameters
- data (str): The simple lines text to parse into a DataFrame.
Returns
- A dataframe containing the rows serialized in
data.