feat: Geospatial - parse Athena geospatial types via geopandas (#2346)

* feat: Add geopandas

* Deserialize shapely geospatial types

* Add test case

* Convert to GeoDataFrame if any of the columns contains geometry

* Expand test case

* Make geometry optional

* Rename geopandas extra

* Fix _fix_csv_types_generator

* Resolve poetry conflicts

Signed-off-by: Anton Kukushkin <[email protected]>

* Add geopandas extra

Signed-off-by: Anton Kukushkin <[email protected]>

* Regenerate poetry.lock

Signed-off-by: Anton Kukushkin <[email protected]>

* Fix static checks

Signed-off-by: Anton Kukushkin <[email protected]>

* Fix static checks

Signed-off-by: Anton Kukushkin <[email protected]>

* Revert lock update

Signed-off-by: Anton Kukushkin <[email protected]>

---------

Signed-off-by: Anton Kukushkin <[email protected]>

Author: kukushking

awswrangler/_data_types.py

@@ -371,6 +371,8 @@ def athena2pandas(dtype: str, dtype_backend: Optional[str] = None) -> str:  # py
         return "bytes" if dtype_backend != "pyarrow" else "binary[pyarrow]"
     if dtype in ("array", "row", "map"):
         return "object"
+    if dtype == "geometry":
+        return "string"
     raise exceptions.UnsupportedType(f"Unsupported Athena type: {dtype}")
 
 

awswrangler/athena/_read.py

@@ -32,9 +32,25 @@
 
 from ._cache import _cache_manager, _CacheInfo, _check_for_cached_results
 
+shapely_wkt = _utils.import_optional_dependency("shapely.wkt")
+geopandas = _utils.import_optional_dependency("geopandas")
+
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
+@_utils.check_optional_dependency(shapely_wkt, "shapely")
+@_utils.check_optional_dependency(geopandas, "geopandas")
+def _cast_geometry(df: pd.DataFrame, parse_geometry: List[str] = None):
+    def load_geom_wkt(x):
+        """Load geometry from well-known text."""
+        return shapely_wkt.loads(x)
+
+    for col in parse_geometry:
+        df[col] = geopandas.GeoSeries(df[col].apply(load_geom_wkt))
+
+    return geopandas.GeoDataFrame(df)
+
+
 def _extract_ctas_manifest_paths(path: str, boto3_session: Optional[boto3.Session] = None) -> List[str]:
     """Get the list of paths of the generated files."""
     bucket_name, key_path = _utils.parse_path(path)
@@ -46,11 +62,11 @@ def _extract_ctas_manifest_paths(path: str, boto3_session: Optional[boto3.Sessio
 
 
 def _fix_csv_types_generator(
-    dfs: Iterator[pd.DataFrame], parse_dates: List[str], binaries: List[str]
+    dfs: Iterator[pd.DataFrame], parse_dates: List[str], binaries: List[str], parse_geometry: List[str]
 ) -> Iterator[pd.DataFrame]:
     """Apply data types cast to a Pandas DataFrames Generator."""
     for df in dfs:
-        yield _fix_csv_types(df=df, parse_dates=parse_dates, binaries=binaries)
+        yield _fix_csv_types(df=df, parse_dates=parse_dates, binaries=binaries, parse_geometry=parse_geometry)
 
 
 def _add_query_metadata_generator(
@@ -62,7 +78,9 @@ def _add_query_metadata_generator(
         yield df
 
 
-def _fix_csv_types(df: pd.DataFrame, parse_dates: List[str], binaries: List[str]) -> pd.DataFrame:
+def _fix_csv_types(
+    df: pd.DataFrame, parse_dates: List[str], binaries: List[str], parse_geometry: List[str]
+) -> pd.DataFrame:
     """Apply data types cast to a Pandas DataFrames."""
     if len(df.index) > 0:
         for col in parse_dates:
@@ -74,6 +92,10 @@ def _fix_csv_types(df: pd.DataFrame, parse_dates: List[str], binaries: List[str]
                 )
         for col in binaries:
             df[col] = df[col].str.encode(encoding="utf-8")
+
+    if geopandas and parse_geometry:
+        df = _cast_geometry(df, parse_geometry=parse_geometry)
+
     return df
 
 
@@ -197,7 +219,12 @@ def _fetch_csv_result(
     )
     _logger.debug("Start type casting...")
     if _chunksize is None:
-        df = _fix_csv_types(df=ret, parse_dates=query_metadata.parse_dates, binaries=query_metadata.binaries)
+        df = _fix_csv_types(
+            df=ret,
+            parse_dates=query_metadata.parse_dates,
+            binaries=query_metadata.binaries,
+            parse_geometry=query_metadata.parse_geometry,
+        )
         df = _apply_query_metadata(df=df, query_metadata=query_metadata)
         if keep_files is False:
             s3.delete_objects(
@@ -207,7 +234,12 @@ def _fetch_csv_result(
                 s3_additional_kwargs=s3_additional_kwargs,
             )
         return df
-    dfs = _fix_csv_types_generator(dfs=ret, parse_dates=query_metadata.parse_dates, binaries=query_metadata.binaries)
+    dfs = _fix_csv_types_generator(
+        dfs=ret,
+        parse_dates=query_metadata.parse_dates,
+        binaries=query_metadata.binaries,
+        parse_geometry=query_metadata.parse_geometry,
+    )
     dfs = _add_query_metadata_generator(dfs=dfs, query_metadata=query_metadata)
     if keep_files is False:
         return _delete_after_iterate(

awswrangler/athena/_utils.py

@@ -49,6 +49,7 @@ class _QueryMetadata(NamedTuple):
     dtype: Dict[str, str]
     parse_timestamps: List[str]
     parse_dates: List[str]
+    parse_geometry: List[str]
     converters: Dict[str, Any]
     binaries: List[str]
     output_location: Optional[str]
@@ -235,6 +236,7 @@ def _get_query_metadata(  # pylint: disable=too-many-statements
     dtype: Dict[str, str] = {}
     parse_timestamps: List[str] = []
     parse_dates: List[str] = []
+    parse_geometry: List[str] = []
     converters: Dict[str, Any] = {}
     binaries: List[str] = []
     col_name: str
@@ -252,6 +254,8 @@ def _get_query_metadata(  # pylint: disable=too-many-statements
             binaries.append(col_name)
         elif pandas_type == "decimal":
             converters[col_name] = lambda x: Decimal(str(x)) if str(x) not in ("", "none", " ", "<NA>") else None
+        elif col_type == "geometry" and pandas_type == "string":
+            parse_geometry.append(col_name)
         else:
             dtype[col_name] = pandas_type
 
@@ -269,6 +273,7 @@ def _get_query_metadata(  # pylint: disable=too-many-statements
         dtype=dtype,
         parse_timestamps=parse_timestamps,
         parse_dates=parse_dates,
+        parse_geometry=parse_geometry,
         converters=converters,
         binaries=binaries,
         output_location=output_location,