Separate the logic between Pandas and GeoPandas

python/sedona/maps/SedonaMapUtils.py

@@ -17,8 +17,6 @@
 
 import json
 
-import geopandas
-
 from sedona.sql.types import GeometryType
 from sedona.utils.geoarrow import dataframe_to_arrow
 
@@ -41,18 +39,17 @@ def __convert_to_gdf_or_pdf__(cls, df, rename=True, geometry_col=None):
         # Convert the dataframe to arrow format, then to geopandas dataframe
         # This is faster than converting directly to geopandas dataframe via toPandas
         data_pyarrow = dataframe_to_arrow(df)
-        pandas_df = geopandas.GeoDataFrame.from_arrow(data_pyarrow)
 
         if (
             geometry_col is None
         ):  # No geometry column found even after searching schema, return Pandas Dataframe
-            return pandas_df
+            return data_pyarrow.to_pandas()
         try:
             import geopandas as gpd
         except ImportError:
             msg = "GeoPandas is missing. You can install it manually or via apache-sedona[kepler-map] or apache-sedona[pydeck-map]."
             raise ImportError(msg) from None
-        geo_df = gpd.GeoDataFrame(pandas_df, geometry=geometry_col)
+        geo_df = gpd.GeoDataFrame.from_arrow(data_pyarrow)
         if geometry_col != "geometry" and rename is True:
             geo_df.rename_geometry("geometry", inplace=True)
         return geo_df

python/sedona/raster_utils/SedonaUtils.py

@@ -17,6 +17,7 @@
 
 import geopandas
 
+from sedona.maps.SedonaMapUtils import SedonaMapUtils
 from sedona.utils.geoarrow import dataframe_to_arrow
 
 
@@ -27,6 +28,10 @@ def display_image(cls, df):
 
         # Convert the dataframe to arrow format, then to geopandas dataframe
         # This is faster than converting directly to geopandas dataframe via toPandas
+        geometry_col = SedonaMapUtils.__get_geometry_col__(df)
         data_pyarrow = dataframe_to_arrow(df)
-        pdf = geopandas.GeoDataFrame.from_arrow(data_pyarrow)
+        if geometry_col is None:
+            pdf = data_pyarrow.to_pandas()
+        else:
+            pdf = geopandas.GeoDataFrame.from_arrow(data_pyarrow)
         display(HTML(pdf.to_html(escape=False)))