col named 'geometry' fixes

Author: ATL2001

src/spatial_polars/spatialframe.py

@@ -201,7 +201,7 @@ def write_geoparquet(
 
     def to_geoarrow(
         self,
-        geometry_name: List[str]|str = ["geometry"],
+        geometry_name: List[str] | str = ["geometry"],
     ):
         r"""
         Converts the dataframe to geoarrow table.
@@ -260,7 +260,7 @@ def to_geoarrow(
             )
 
         pa_table = self._df.drop(geometry_name).to_arrow()
-        
+
         for this_g_name in geometry_name:
             crs = pyproj.CRS(self._df[this_g_name].struct.field("crs")[0]).to_wkt(
                 version="WKT1_GDAL"
@@ -654,6 +654,7 @@ def viz(
 
         """
         from lonboard import viz
+
         geoarrow_table = self.to_geoarrow(geometry_name)
 
         return viz(
@@ -698,7 +699,7 @@ def to_scatterplotlayer(
         visible: bool = True,
         antialiasing: bool = True,
         billboard: bool = False,
-        **kwargs
+        **kwargs,
     ) -> ScatterplotLayer:
         """
         Makes a Lonboard [ScatterplotLayer][lonboard.ScatterplotLayer] from the SpatialFrame.
@@ -927,7 +928,7 @@ def to_scatterplotlayer(
             radius_units=radius_units,
             stroked=stroked,
             visible=visible,
-            **kwargs
+            **kwargs,
         )
         return layer
 
@@ -954,7 +955,7 @@ def to_pathlayer(
         width_max_pixels: Optional[float] = None,
         width_scale: float = 1,
         width_units: Literal["meters", "common", "pixels"] = "meters",
-        **kwargs
+        **kwargs,
     ) -> PathLayer:
         """
         Makes a Lonboard [PathLayer][lonboard.PathLayer] from the SpatialFrame.
@@ -1041,7 +1042,6 @@ def to_pathlayer(
         from lonboard import PathLayer
         from lonboard.colormap import apply_continuous_cmap, apply_categorical_cmap
 
-
         validate_cmap_input(
             self._df, cmap_col, cmap_type, cmap, alpha, normalize_cmap_col
         )
@@ -1094,7 +1094,7 @@ def to_pathlayer(
             width_max_pixels=width_max_pixels,
             width_scale=width_scale,
             width_units=width_units,
-            **kwargs
+            **kwargs,
         )
         return layer
 
@@ -1130,7 +1130,7 @@ def to_polygonlayer(
         pickable: bool = True,
         visible: bool = True,
         wireframe: bool = False,
-        **kwargs
+        **kwargs,
     ) -> PolygonLayer:
         """
         Makes a Lonboard [PolygonLayer][lonboard.PolygonLayer] from the SpatialFrame.
@@ -1354,7 +1354,7 @@ def to_polygonlayer(
             stroked=stroked,
             visible=visible,
             wireframe=wireframe,
-            **kwargs
+            **kwargs,
         )
         return layer
 
@@ -1502,6 +1502,14 @@ def from_WKB(df: pl.DataFrame, wkb_col: str, crs: Any = 4326):
         """
         crs_wkt = pyproj.CRS.from_user_input(crs).to_wkt()
 
+        if wkb_col == "geometry":
+            return df.with_columns(
+                pl.struct(
+                    c(wkb_col).alias("wkb_geometry"),
+                    pl.lit(crs_wkt, dtype=pl.Categorical).alias("crs"),
+                ).alias("geometry")
+            )
+
         return df.with_columns(
             pl.struct(
                 c(wkb_col).alias("wkb_geometry"),
@@ -1573,6 +1581,15 @@ def from_WKT(df, wkt_col: str, crs: Any = 4326):
         geoms = shapely.from_wkt(df.select(wkt_col).to_series().to_numpy().copy())
         wkb_array = shapely.to_wkb(geoms)
         crs_wkt = pyproj.CRS.from_user_input(crs).to_wkt()
+
+        if wkt_col == "geometry":
+            return df.with_columns(
+                pl.struct(
+                    pl.Series("wkb_geometry", wkb_array, dtype=pl.Binary),
+                    pl.lit(crs_wkt, dtype=pl.Categorical).alias("crs"),
+                ).alias("geometry")
+            )
+
         return df.with_columns(
             pl.struct(
                 pl.Series("wkb_geometry", wkb_array, dtype=pl.Binary),

tests/test_from_polars.py

@@ -31,6 +31,19 @@ def test_from_WKB(arch_mound_df: pl.DataFrame) -> None:
     assert_frame_equal(s_df, arch_mound_df)
 
 
+def test_from_WKB_named_geometry(arch_mound_df: pl.DataFrame) -> None:
+    arch_wkb = shapely.Point(-90.18497, 38.62456).wkb
+    monks_mound_wkb = shapely.Point(-90.06211, 38.66072).wkb
+    df = pl.DataFrame(
+        {
+            "Place": ["Gateway Arch", "Monks Mound"],
+            "geometry": [arch_wkb, monks_mound_wkb],
+        }
+    )
+    s_df = SpatialFrame.from_WKB(df, "geometry")
+    assert_frame_equal(s_df, arch_mound_df)
+
+
 def test_from_WKT(arch_mound_df: pl.DataFrame) -> None:
     arch_wkt = shapely.Point(-90.18497, 38.62456).wkt
     monks_mound_wkt = shapely.Point(-90.06211, 38.66072).wkt
@@ -42,3 +55,16 @@ def test_from_WKT(arch_mound_df: pl.DataFrame) -> None:
     )
     s_df = SpatialFrame.from_WKT(df, "wkt")
     assert_frame_equal(s_df, arch_mound_df)
+
+
+def test_from_WKT_named_geometry(arch_mound_df: pl.DataFrame) -> None:
+    arch_wkt = shapely.Point(-90.18497, 38.62456).wkt
+    monks_mound_wkt = shapely.Point(-90.06211, 38.66072).wkt
+    df = pl.DataFrame(
+        {
+            "Place": ["Gateway Arch", "Monks Mound"],
+            "geometry": [arch_wkt, monks_mound_wkt],
+        }
+    )
+    s_df = SpatialFrame.from_WKT(df, "geometry")
+    assert_frame_equal(s_df, arch_mound_df)