Merge pull request #1338 from CartoDB/fix-quad-grid

Fix QuadGrid class

cartoframes/analysis/grid.py

@@ -1,27 +1,33 @@
 import mercantile
-from shapely.geometry import box
 import pandas as pd
-from cartoframes import CartoDataFrame
+
+from shapely.geometry import box
+from ..core.cartodataframe import CartoDataFrame
 
 
 class QuadGrid():
 
     def polyfill(self, input_gdf, zoom_level):
+        if not hasattr(input_gdf, 'geometry'):
+            raise ValueError('This dataframe has no valid geometry.')
 
-        dfs = []
+        geometry_name = input_gdf.geometry.name
 
+        dfs = []
         for _, row in input_gdf.iterrows():
-            bounds = row['geometry'].bounds
+            input_geometry = row[geometry_name]
+            bounds = input_geometry.bounds
             tiles = mercantile.tiles(bounds[0], bounds[1], bounds[2], bounds[3], zoom_level)
-            resp = []
-            for t in tiles:
-                r = row.copy()
-                geometry = box(*mercantile.bounds(t))
-                if geometry.intersects(row['geometry']):
-                    r['geometry'] = geometry
-                    r['quadkey'] = mercantile.quadkey(t)
-                    resp.append(r)
-
-            dfs.append(pd.DataFrame(resp))
-
-        return CartoDataFrame(pd.concat(dfs).reset_index(drop=True), crs='epsg:4326')
+            new_rows = []
+            for tile in tiles:
+                new_row = row.copy()
+                new_geometry = box(*mercantile.bounds(tile))
+                if new_geometry.intersects(input_geometry):
+                    new_row[geometry_name] = new_geometry
+                    new_row['quadkey'] = mercantile.quadkey(tile)
+                    new_rows.append(new_row)
+            dfs.append(pd.DataFrame(new_rows))
+
+        df = pd.concat(dfs).reset_index(drop=True)
+
+        return CartoDataFrame(df, geometry=geometry_name, crs='epsg:4326')

tests/unit/analysis/grid_quadkey_bbox.csv

@@ -1,4 +1,4 @@
-id,geometry,quadkey
+id,geom,quadkey
 1,"POLYGON ((1.0546875 1.933226826477117, 1.0546875 2.021065118766994, 0.966796875 2.021065118766994, 0.966796875 1.933226826477117, 1.0546875 1.933226826477117))",122222203013
 1,"POLYGON ((1.0546875 1.845383988573192, 1.0546875 1.933226826477117, 0.966796875 1.933226826477117, 0.966796875 1.845383988573192, 1.0546875 1.845383988573192))",122222203031
 1,"POLYGON ((1.0546875 1.757536811308316, 1.0546875 1.845383988573192, 0.966796875 1.845383988573192, 0.966796875 1.757536811308316, 1.0546875 1.757536811308316))",122222203033

tests/unit/analysis/grid_quadkey_pol.csv

@@ -1,4 +1,4 @@
-id,geometry,quadkey
+id,geom,quadkey
 1,"POLYGON ((-6.6796875 37.64903402157864, -6.6796875 37.71859032558815, -6.767578125 37.71859032558815, -6.767578125 37.64903402157864, -6.6796875 37.64903402157864))",033110330011
 1,"POLYGON ((-6.591796875 37.71859032558815, -6.591796875 37.78808138412046, -6.6796875 37.78808138412046, -6.6796875 37.71859032558815, -6.591796875 37.71859032558815))",033110312322
 1,"POLYGON ((-6.591796875 37.64903402157864, -6.591796875 37.71859032558815, -6.6796875 37.71859032558815, -6.6796875 37.64903402157864, -6.591796875 37.64903402157864))",033110330100

tests/unit/analysis/test_grid.py

@@ -5,7 +5,6 @@
 import os
 
 from shapely.geometry import box, shape
-from shapely import wkt
 from cartoframes.analysis.grid import QuadGrid
 from cartoframes import CartoDataFrame
 
@@ -15,7 +14,7 @@
 # DATA FRAME SRC BBOX
 pol_1 = box(1, 1, 2, 2)
 pol_2 = box(3, 3, 4, 4)
-GDF_BOX = gpd.GeoDataFrame({'id': [1, 2], 'geometry': [pol_1, pol_2]}, columns=['id', 'geometry'])
+GDF_BOX = gpd.GeoDataFrame({'id': [1, 2], 'geom': [pol_1, pol_2]}, columns=['id', 'geom'], geometry='geom')
 
 pol_geojson = {
     "type": "Polygon",
@@ -61,7 +60,7 @@
     ]
 }
 
-GDF_IRREGULAR = gpd.GeoDataFrame({'id': [1], 'geometry': [shape(pol_geojson)]}, columns=['id', 'geometry'])
+GDF_IRREGULAR = gpd.GeoDataFrame({'id': [1], 'geom': [shape(pol_geojson)]}, columns=['id', 'geom'], geometry='geom')
 
 BASE_FOLDER = os.path.dirname(os.path.abspath(__file__))
 
@@ -70,10 +69,8 @@ class TestGrid(object):
 
     def _load_test_gdf(self, fname):
         fname = os.path.join(BASE_FOLDER, fname)
-        df = pd.read_csv(fname, dtype={'id': np.int64, 'geometry': object, 'quadkey': object})
-        df['geometry'] = df['geometry'].apply(wkt.loads)
-        cdf_test = CartoDataFrame(df)
-        cdf_test.crs = 'epsg:4326'
+        df = pd.read_csv(fname, dtype={'id': np.int64, 'geom': object, 'quadkey': object})
+        cdf_test = CartoDataFrame(df, geometry='geom', crs='epsg:4326')
         return cdf_test
 
     def test_quadgrid_polyfill_box(self, mocker):