meerschaum 2.8.4__py3-none-any.whl → 2.9.0rc1__py3-none-any.whl

meerschaum/utils/dataframe.py

@@ -153,6 +153,7 @@ def filter_unseen_df(
         attempt_cast_to_numeric,
         attempt_cast_to_uuid,
         attempt_cast_to_bytes,
+        attempt_cast_to_geometry,
         coerce_timezone,
         serialize_decimal,
     )
@@ -350,6 +351,10 @@ def filter_unseen_df(
     new_bytes_cols = get_bytes_cols(new_df)
     bytes_cols = set(new_bytes_cols + old_bytes_cols)
 
+    old_geometry_cols = get_geometry_cols(old_df)
+    new_geometry_cols = get_geometry_cols(new_df)
+    geometry_cols = set(new_geometry_cols + old_geometry_cols)
+
     joined_df = merge(
         new_df.infer_objects(copy=False).fillna(NA),
         old_df.infer_objects(copy=False).fillna(NA),
@@ -400,6 +405,14 @@ def filter_unseen_df(
         except Exception:
             warn(f"Unable to parse bytes column '{bytes_col}':\n{traceback.format_exc()}")
 
+    for geometry_col in geometry_cols:
+        if geometry_col not in delta_df.columns:
+            continue
+        try:
+            delta_df[geometry_col] = delta_df[geometry_col].apply(attempt_cast_to_geometry)
+        except Exception:
+            warn(f"Unable to parse bytes column '{bytes_col}':\n{traceback.format_exc()}")
+
     return delta_df
 
 
@@ -858,6 +871,44 @@ def get_bytes_cols(df: 'pd.DataFrame') -> List[str]:
     ]
 
 
+def get_geometry_cols(df: 'pd.DataFrame') -> List[str]:
+    """
+    Get the columns which contain shapely objects from a Pandas DataFrame.
+
+    Parameters
+    ----------
+    df: pd.DataFrame
+        The DataFrame which may contain bytes strings.
+
+    Returns
+    -------
+    A list of columns to treat as `geometry`.
+    """
+    if df is None:
+        return []
+
+    is_dask = 'dask' in df.__module__
+    if is_dask:
+        df = get_first_valid_dask_partition(df)
+
+    if len(df) == 0:
+        return []
+
+    cols_indices = {
+        col: df[col].first_valid_index()
+        for col in df.columns
+    }
+    return [
+        col
+        for col, ix in cols_indices.items()
+        if (
+            ix is not None
+            and
+            'shapely' in str(type(df.loc[ix][col]))
+        )
+    ]
+
+
 def enforce_dtypes(
     df: 'pd.DataFrame',
     dtypes: Dict[str, str],
@@ -911,6 +962,7 @@ def enforce_dtypes(
         attempt_cast_to_numeric,
         attempt_cast_to_uuid,
         attempt_cast_to_bytes,
+        attempt_cast_to_geometry,
         coerce_timezone as _coerce_timezone,
     )
     from meerschaum.utils.dtypes.sql import get_numeric_precision_scale
@@ -937,6 +989,11 @@ def enforce_dtypes(
         for col, typ in dtypes.items()
         if typ.startswith('numeric')
     ]
+    geometry_cols = [
+        col
+        for col, typ in dtypes.items()
+        if typ.startswith('geometry') or typ.startswith('geography')
+    ]
     uuid_cols = [
         col
         for col, typ in dtypes.items()
@@ -1026,6 +1083,24 @@ def enforce_dtypes(
             if col in df.columns:
                 df[col] = _coerce_timezone(df[col], strip_utc=strip_timezone)
 
+    if geometry_cols:
+        geopandas = mrsm.attempt_import('geopandas')
+        if debug:
+            dprint(f"Checking for geometry: {geometry_cols}")
+        parsed_geom_cols = []
+        for col in geometry_cols:
+            try:
+                df[col] = df[col].apply(attempt_cast_to_geometry)
+                parsed_geom_cols.append(col)
+            except Exception as e:
+                if debug:
+                    dprint(f"Unable to parse column '{col}' as geometry:\n{e}")
+
+        if parsed_geom_cols:
+            if debug:
+                dprint(f"Converting to GeoDataFrame (geometry column: '{parsed_geom_cols[0]}')...")
+            df = geopandas.GeoDataFrame(df, geometry=parsed_geom_cols[0])
+
     df_dtypes = {c: str(t) for c, t in df.dtypes.items()}
     if are_dtypes_equal(df_dtypes, pipe_pandas_dtypes):
         if debug:
@@ -1602,13 +1677,19 @@ def to_json(
     -------
     A JSON string.
     """
+    import warnings
     from meerschaum.utils.packages import import_pandas
-    from meerschaum.utils.dtypes import serialize_bytes, serialize_decimal
+    from meerschaum.utils.dtypes import (
+        serialize_bytes,
+        serialize_decimal,
+        serialize_geometry,
+    )
     pd = import_pandas()
     uuid_cols = get_uuid_cols(df)
     bytes_cols = get_bytes_cols(df)
     numeric_cols = get_numeric_cols(df)
-    if safe_copy and bool(uuid_cols or bytes_cols):
+    geometry_cols = get_geometry_cols(df)
+    if safe_copy and bool(uuid_cols or bytes_cols or geometry_cols or numeric_cols):
         df = df.copy()
     for col in uuid_cols:
         df[col] = df[col].astype(str)
@@ -1616,6 +1697,10 @@ def to_json(
         df[col] = df[col].apply(serialize_bytes)
     for col in numeric_cols:
         df[col] = df[col].apply(serialize_decimal)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        for col in geometry_cols:
+            df[col] = df[col].apply(serialize_geometry)
     return df.infer_objects(copy=False).fillna(pd.NA).to_json(
         date_format=date_format,
         date_unit=date_unit,

meerschaum/utils/dtypes/__init__.py

@@ -12,7 +12,7 @@ from datetime import timezone, datetime
 from decimal import Decimal, Context, InvalidOperation, ROUND_HALF_UP
 
 import meerschaum as mrsm
-from meerschaum.utils.typing import Dict, Union, Any, Optional
+from meerschaum.utils.typing import Dict, Union, Any, Optional, Tuple
 from meerschaum.utils.warnings import warn
 
 MRSM_ALIAS_DTYPES: Dict[str, str] = {
@@ -27,10 +27,13 @@ MRSM_ALIAS_DTYPES: Dict[str, str] = {
     'bytea': 'bytes',
     'guid': 'uuid',
     'UUID': 'uuid',
+    'geom': 'geometry',
 }
 MRSM_PD_DTYPES: Dict[Union[str, None], str] = {
     'json': 'object',
     'numeric': 'object',
+    'geometry': 'object',
+    'geography': 'object',
     'uuid': 'object',
     'datetime': 'datetime64[ns, UTC]',
     'bool': 'bool[pyarrow]',
@@ -60,6 +63,12 @@ def to_pandas_dtype(dtype: str) -> str:
     if dtype.startswith('numeric'):
         return MRSM_PD_DTYPES['numeric']
 
+    if dtype.startswith('geometry'):
+        return MRSM_PD_DTYPES['geometry']
+
+    if dtype.startswith('geography'):
+        return MRSM_PD_DTYPES['geography']
+
     ### NOTE: Kind of a hack, but if the first word of the given dtype is in all caps,
     ### treat it as a SQL db type.
     if dtype.split(' ')[0].isupper():
@@ -147,6 +156,10 @@ def are_dtypes_equal(
     if ldtype in bytes_dtypes and rdtype in bytes_dtypes:
         return True
 
+    geometry_dtypes = ('geometry', 'object', 'geography')
+    if ldtype in geometry_dtypes and rdtype in geometry_dtypes:
+        return True
+
     if ldtype.lower() == rdtype.lower():
         return True
 
@@ -277,6 +290,56 @@ def attempt_cast_to_bytes(value: Any) -> Any:
         return value
 
 
+def attempt_cast_to_geometry(value: Any) -> Any:
+    """
+    Given a value, attempt to coerce it into a `shapely` (`geometry`) object.
+    """
+    shapely = mrsm.attempt_import('shapely', lazy=False)
+    if 'shapely' in str(type(value)):
+        return value
+
+    value_is_wkt = geometry_is_wkt(value)
+    if value_is_wkt is None:
+        return value
+
+    try:
+        return (
+            shapely.wkt.loads(value)
+            if value_is_wkt
+            else shapely.wkb.loads(value)
+        )
+    except Exception:
+        return value
+
+
+def geometry_is_wkt(value: Union[str, bytes]) -> Union[bool, None]:
+    """
+    Determine whether an input value should be treated as WKT or WKB geometry data.
+
+    Parameters
+    ----------
+    value: Union[str, bytes]
+        The input data to be parsed into geometry data.
+
+    Returns
+    -------
+    A `bool` (`True` if `value` is WKT and `False` if it should be treated as WKB).
+    Return `None` if `value` should be parsed as neither.
+    """
+    import re
+    if isinstance(value, bytes):
+        return False
+    
+    wkt_pattern = r'^\s*(POINT|LINESTRING|POLYGON|MULTIPOINT|MULTILINESTRING|MULTIPOLYGON|GEOMETRYCOLLECTION)\s*\(.*\)\s*$'
+    if re.match(wkt_pattern, value, re.IGNORECASE):
+        return True
+    
+    if all(c in '0123456789ABCDEFabcdef' for c in value) and len(value) % 2 == 0:
+        return False
+    
+    return None
+
+
 def value_is_null(value: Any) -> bool:
     """
     Determine if a value is a null-like string.
@@ -458,6 +521,37 @@ def serialize_bytes(data: bytes) -> str:
     return base64.b64encode(data).decode('utf-8')
 
 
+def serialize_geometry(geom: Any, as_wkt: bool = False) -> str:
+    """
+    Serialize geometry data as a hex-encoded well-known-binary string. 
+
+    Parameters
+    ----------
+    geom: Any
+        The potential geometry data to be serialized.
+
+    as_wkt, bool, default False
+        If `True`, serialize geometry data as well-known text (WKT)
+        instead of well-known binary (WKB).
+
+    Returns
+    -------
+    A string containing the geometry data.
+    """
+    if hasattr(geom, 'wkb_hex'):
+        return geom.wkb_hex if not as_wkt else geom.wkt
+
+    return str(geom)
+
+
+def deserialize_geometry(geom_wkb: Union[str, bytes]):
+    """
+    Deserialize a WKB string into a shapely geometry object.
+    """
+    shapely = mrsm.attempt_import(lazy=False)
+    return shapely.wkb.loads(geom_wkb)
+
+
 def deserialize_bytes_string(data: str | None, force_hex: bool = False) -> bytes | None:
     """
     Given a serialized ASCII string of bytes data, return the original bytes.
@@ -559,7 +653,94 @@ def json_serialize_value(x: Any, default_to_str: bool = True) -> str:
     if isinstance(x, Decimal):
         return serialize_decimal(x)
 
+    if 'shapely' in str(type(x)):
+        return serialize_geometry(x)
+
     if value_is_null(x):
         return None
 
     return str(x) if default_to_str else x
+
+
+def get_geometry_type_srid(
+    dtype: str = 'geometry',
+    default_type: str = 'geometry',
+    default_srid: int = 4326,
+) -> Union[Tuple[str, int], Tuple[str, None]]:
+    """
+    Given the specified geometry `dtype`, return a tuple in the form (type, SRID).
+
+    Parameters
+    ----------
+    dtype: Optional[str], default None
+        Optionally provide a specific `geometry` syntax (e.g. `geometry[MultiLineString, 4326]`).
+        You may specify a supported `shapely` geometry type and an SRID in the dtype modifier:
+
+        - `Point`
+        - `LineString`
+        - `LinearRing`
+        - `Polygon`
+        - `MultiPoint`
+        - `MultiLineString`
+        - `MultiPolygon`
+        - `GeometryCollection`
+
+    Returns
+    -------
+    A tuple in the form (type, SRID).
+    Defaults to `(default_type, default_srid)`.
+
+    Examples
+    --------
+    >>> from meerschaum.utils.dtypes import get_geometry_type_srid
+    >>> get_geometry_type_srid()
+    ('geometry', 4326)
+    >>> get_geometry_type_srid('geometry[]')
+    ('geometry', 4326)
+    >>> get_geometry_type_srid('geometry[Point, 0]')
+    ('Point', 0)
+    >>> get_geometry_type_srid('geometry[0, Point]')
+    ('Point', 0)
+    >>> get_geometry_type_srid('geometry[0]')
+    ('geometry', 0)
+    >>> get_geometry_type_srid('geometry[MULTILINESTRING, 4326]')
+    ('MultiLineString', 4326)
+    >>> get_geometry_type_srid('geography')
+    ('geometry', 4326)
+    >>> get_geometry_type_srid('geography[POINT]')
+    ('Point', 4376)
+    """
+    from meerschaum.utils.misc import is_int
+    bare_dtype = dtype.split('[', maxsplit=1)[0]
+    modifier = dtype.split(bare_dtype, maxsplit=1)[-1].lstrip('[').rstrip(']')
+    if not modifier:
+        return default_type, default_srid
+
+    shapely_geometry_base = mrsm.attempt_import('shapely.geometry.base')
+    geometry_types = {
+        typ.lower(): typ
+        for typ in shapely_geometry_base.GEOMETRY_TYPES
+    }
+
+    parts = [part.lower().replace('srid=', '').replace('type=', '').strip() for part in modifier.split(',')]
+    parts_casted = [
+        (
+            int(part)
+            if is_int(part)
+            else part
+        ) for part in parts]
+
+    srid = default_srid
+    geometry_type = default_type
+
+    for part in parts_casted:
+        if isinstance(part, int):
+            srid = part
+            break
+
+    for part in parts:
+        if part.lower() in geometry_types:
+            geometry_type = geometry_types.get(part)
+            break
+
+    return geometry_type, srid