ENH: Read from postgis in chunks (#1123)

Co-authored-by: Oliver Schillinger <oliver.schillinger@innogy.com>
Co-authored-by: James McBride <jdmcbr@gmail.com>
Co-authored-by: Martin Fleischmann <martin@martinfleischmann.net>
Co-authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>

geopandas/geodataframe.py

@@ -399,6 +399,7 @@ class GeoDataFrame(GeoPandasBase, DataFrame):
         coerce_float=True,
         parse_dates=None,
         params=None,
+        chunksize=None,
     ):
         """
         Alternate constructor to create a ``GeoDataFrame`` from a sql query
@@ -428,6 +429,9 @@ class GeoDataFrame(GeoPandasBase, DataFrame):
               without native Datetime support, such as SQLite.
         params : list, tuple or dict, optional, default None
             List of parameters to pass to execute method.
+        chunksize : int, default None
+            If specified, return an iterator where chunksize is the number
+            of rows to include in each chunk.
 
         Examples
         --------
@@ -446,6 +450,7 @@ class GeoDataFrame(GeoPandasBase, DataFrame):
             coerce_float=coerce_float,
             parse_dates=parse_dates,
             params=params,
+            chunksize=chunksize,
         )
 
         return df

geopandas/io/sql.py

@@ -9,27 +9,15 @@ from geopandas import GeoDataFrame
 from .. import _compat as compat
 
 
-def _read_postgis(
-    sql,
-    con,
-    geom_col="geom",
-    crs=None,
-    index_col=None,
-    coerce_float=True,
-    parse_dates=None,
-    params=None,
-):
+def _df_to_geodf(df, geom_col="geom", crs=None):
     """
-    Returns a GeoDataFrame corresponding to the result of the query
-    string, which must contain a geometry column in WKB representation.
-
+    Transforms a pandas DataFrame into a GeoDataFrame.
+    The column 'geom_col' must be a geometry column in WKB representation.
+    To be used to convert df based on pd.read_sql to gdf.
     Parameters
     ----------
-    sql : string
-        SQL query to execute in selecting entries from database, or name
-        of the table to read from the database.
-    con : DB connection object or SQLAlchemy engine
-        Active connection to the database to query.
+    df : DataFrame
+        pandas DataFrame with geometry column in WKB representation.
     geom_col : string, default 'geom'
         column name to convert to shapely geometries
     crs : pyproj.CRS, optional
@@ -38,33 +26,11 @@ def _read_postgis(
         such as an authority string (eg "EPSG:4326") or a WKT string.
         If not set, tries to determine CRS from the SRID associated with the
         first geometry in the database, and assigns that to all geometries.
-
-    See the documentation for pandas.read_sql for further explanation
-    of the following parameters:
-    index_col, coerce_float, parse_dates, params
-
     Returns
     -------
     GeoDataFrame
-
-    Example
-    -------
-    PostGIS
-    >>> sql = "SELECT geom, kind FROM polygons"
-    SpatiaLite
-    >>> sql = "SELECT ST_AsBinary(geom) AS geom, kind FROM polygons"
-    >>> df = geopandas.read_postgis(sql, con)
     """
 
-    df = pd.read_sql(
-        sql,
-        con,
-        index_col=index_col,
-        coerce_float=coerce_float,
-        parse_dates=parse_dates,
-        params=params,
-    )
-
     if geom_col not in df:
         raise ValueError("Query missing geometry column '{}'".format(geom_col))
 
@@ -97,6 +63,82 @@ def _read_postgis(
     return GeoDataFrame(df, crs=crs, geometry=geom_col)
 
 
+def _read_postgis(
+    sql,
+    con,
+    geom_col="geom",
+    crs=None,
+    index_col=None,
+    coerce_float=True,
+    parse_dates=None,
+    params=None,
+    chunksize=None,
+):
+    """
+    Returns a GeoDataFrame corresponding to the result of the query
+    string, which must contain a geometry column in WKB representation.
+
+    Parameters
+    ----------
+    sql : string
+        SQL query to execute in selecting entries from database, or name
+        of the table to read from the database.
+    con : DB connection object or SQLAlchemy engine
+        Active connection to the database to query.
+    geom_col : string, default 'geom'
+        column name to convert to shapely geometries
+    crs : dict or str, optional
+        CRS to use for the returned GeoDataFrame; if not set, tries to
+        determine CRS from the SRID associated with the first geometry in
+        the database, and assigns that to all geometries.
+    chunksize : int, default None
+        If specified, return an iterator where chunksize is the number of rows to
+        include in each chunk.
+
+    See the documentation for pandas.read_sql for further explanation
+    of the following parameters:
+    index_col, coerce_float, parse_dates, params, chunksize
+
+    Returns
+    -------
+    GeoDataFrame
+
+    Example
+    -------
+    PostGIS
+    >>> sql = "SELECT geom, kind FROM polygons"
+    SpatiaLite
+    >>> sql = "SELECT ST_AsBinary(geom) AS geom, kind FROM polygons"
+    >>> df = geopandas.read_postgis(sql, con)
+    """
+
+    if chunksize is None:
+        # read all in one chunk and return a single GeoDataFrame
+        df = pd.read_sql(
+            sql,
+            con,
+            index_col=index_col,
+            coerce_float=coerce_float,
+            parse_dates=parse_dates,
+            params=params,
+            chunksize=chunksize,
+        )
+        return _df_to_geodf(df, geom_col=geom_col, crs=crs)
+
+    else:
+        # read data in chunks and return a generator
+        df_generator = pd.read_sql(
+            sql,
+            con,
+            index_col=index_col,
+            coerce_float=coerce_float,
+            parse_dates=parse_dates,
+            params=params,
+            chunksize=chunksize,
+        )
+        return (_df_to_geodf(df, geom_col=geom_col, crs=crs) for df in df_generator)
+
+
 def read_postgis(*args, **kwargs):
     import warnings
 

geopandas/io/tests/test_sql.py

@@ -6,6 +6,8 @@ see geopandas.tests.util for more information.
 """
 import os
 
+import pandas as pd
+
 import geopandas
 from geopandas import GeoDataFrame, read_file, read_postgis
 
@@ -292,6 +294,20 @@ class TestIO:
         df = read_postgis(sql, con, geom_col=geom_col)
         validate_boro_df(df)
 
+    def test_read_postgis_chunksize(self, connection_postgis, df_nybb):
+        """Test chunksize argument"""
+        chunksize = 2
+        con = connection_postgis
+        create_postgis(con, df_nybb)
+
+        sql = "SELECT * FROM nybb;"
+        df = pd.concat(read_postgis(sql, con, chunksize=chunksize))
+
+        validate_boro_df(df)
+        # no crs defined on the created geodatabase, and none specified
+        # by user; should not be set to 0, as from get_srid failure
+        assert df.crs is None
+
     def test_read_postgis_privacy(self, connection_postgis, df_nybb):
         con = connection_postgis
         create_postgis(con, df_nybb)