Return GeoDataFrame in the to_dataframe function

cartoframes/data/observatory/catalog/dataset.py

@@ -408,7 +408,7 @@ def to_csv(self, file_path, credentials=None, limit=None, order_by=None, sql_que
 
     @check_do_enabled
     def to_dataframe(self, credentials=None, limit=None, order_by=None, sql_query=None, add_geom=None):
-        """Download dataset data as a pandas.DataFrame. You need Data Observatory enabled in your CARTO
+        """Download dataset data as a geopandas.GeoDataFrame. You need Data Observatory enabled in your CARTO
         account, please contact us at support@carto.com for more information.
 
         For premium datasets (those with `is_public_data` set to False), you need a subscription to the dataset.
@@ -427,7 +427,7 @@ def to_dataframe(self, credentials=None, limit=None, order_by=None, sql_query=No
 
 
         Returns:
-            pandas.DataFrame
+            geopandas.GeoDataFrame
 
         Raises:
             DOError: if you have not a valid license for the dataset being downloaded,

cartoframes/data/observatory/catalog/entity.py

@@ -1,8 +1,10 @@
 import pandas as pd
 
 from abc import ABC
+from geopandas import GeoDataFrame
 
 from carto.do_dataset import DODataset
+from ....utils.geom_utils import set_geometry
 from ....utils.logger import log
 
 _DATASET_READ_MSG = '''To load it as a DataFrame you can do:
@@ -18,6 +20,8 @@
     gdf = GeoDataFrame(df, geometry=decode_geometry(df['geom']))
 '''
 
+GEOM_COL = 'geom'
+
 
 class CatalogEntity(ABC):
     """This is an internal class the rest of the classes related to the catalog discovery extend.
@@ -142,7 +146,12 @@ def _download(self, credentials, file_path=None, limit=None, order_by=None, sql_
                 log.info(_GEOGRAPHY_READ_MSG.format(file_path))
         else:
             dataframe = pd.read_csv(rows)
-            return dataframe
+            gdf = GeoDataFrame(dataframe)
+
+            if GEOM_COL in gdf:
+                set_geometry(gdf, GEOM_COL, inplace=True)
+
+            return gdf
 
     def _get_remote_full_table_name(self, user_project, user_dataset, public_project):
         project, dataset, table = self.id.split('.')

tests/e2e/data/observatory/catalog/test_download.py

@@ -2,7 +2,9 @@
 import json
 import pandas
 import pytest
+
 from pathlib import Path
+from geopandas import GeoDataFrame
 
 from cartoframes.auth import Credentials
 from cartoframes.data.observatory import Dataset, Geography
@@ -72,6 +74,9 @@ def test_dataset_to_csv_private(self):
 
     def test_dataset_to_dataframe_public(self):
         df = public_dataset.to_dataframe(self.credentials, limit=PUBLIC_LIMIT)
+        assert isinstance(df, GeoDataFrame)
+        assert df.geom.type == 'Polygon'
+
         df.to_csv(self.tmp_file, index=False)
 
         df = pandas.read_csv(self.tmp_file)
@@ -81,6 +86,9 @@ def test_dataset_to_dataframe_public(self):
 
     def test_dataset_to_dataframe_private(self):
         df = private_dataset.to_dataframe(self.credentials, limit=PRIVATE_LIMIT)
+        assert isinstance(df, GeoDataFrame)
+        assert df.geom.type == 'Polygon'
+
         df.to_csv(self.tmp_file, index=False)
 
         df = pandas.read_csv(self.tmp_file)
@@ -111,6 +119,9 @@ def test_geography_to_csv_private(self):
     @pytest.mark.skip()  # TODO implement equals check using a tolerance
     def test_geography_to_dataframe_public(self):
         df = public_geography.to_dataframe(self.credentials, limit=PUBLIC_LIMIT, order_by='geoid')
+        assert isinstance(df, GeoDataFrame)
+        assert df.geom.type == 'Polygon'
+
         df.to_csv(self.tmp_file, index=False)
 
         df = pandas.read_csv(self.tmp_file)
@@ -121,6 +132,9 @@ def test_geography_to_dataframe_public(self):
     @pytest.mark.skip()  # TODO implement equals check using a tolerance
     def test_geography_to_dataframe_private(self):
         df = private_geography.to_dataframe(self.credentials, limit=PRIVATE_LIMIT, order_by='geoid')
+        assert isinstance(df, GeoDataFrame)
+        assert df.geom.type == 'Polygon'
+
         df.to_csv(self.tmp_file, index=False)
 
         df = pandas.read_csv(self.tmp_file)
@@ -165,6 +179,9 @@ def test_geography_to_csv_public_with_sql_query(self):
     def test_dataset_to_dataframe_public_with_sql_query(self):
         sql_query = 'select * from {dataset} order by geoid limit 2'
         df = public_dataset.to_dataframe(self.credentials, sql_query=sql_query)
+        assert isinstance(df, GeoDataFrame)
+        assert df.geom.type == 'Polygon'
+
         df.to_csv(self.tmp_file, index=False)
 
         df = pandas.read_csv(self.tmp_file)
@@ -177,6 +194,9 @@ def test_dataset_to_dataframe_public_with_sql_query_and_add_geom(self):
         sql_query = 'select * from {dataset} order by geoid limit 2'
         add_geom = True
         df = public_dataset.to_dataframe(self.credentials, sql_query=sql_query, add_geom=add_geom)
+        assert isinstance(df, GeoDataFrame)
+        assert df.geom.type == 'Polygon'
+
         df.to_csv(self.tmp_file, index=False)
 
         df = pandas.read_csv(self.tmp_file)
@@ -188,6 +208,9 @@ def test_dataset_to_dataframe_public_with_sql_query_and_add_geom(self):
     def test_geography_to_dataframe_public_with_sql_query(self):
         sql_query = 'select * from {geography} order by geoid limit 2'
         df = public_geography.to_dataframe(self.credentials, sql_query=sql_query)
+        assert isinstance(df, GeoDataFrame)
+        assert df.geom.type == 'Polygon'
+
         df.to_csv(self.tmp_file, index=False)
 
         df = pandas.read_csv(self.tmp_file)