From 0beb61ca3d791fbb2b877d07f43a6f878874b43d Mon Sep 17 00:00:00 2001 From: Oliver Schillinger Date: Sat, 23 May 2020 09:33:24 +0200 Subject: [PATCH] ENH: Read from postgis in chunks (#1123) Co-authored-by: Oliver Schillinger Co-authored-by: James McBride Co-authored-by: Martin Fleischmann Co-authored-by: Joris Van den Bossche --- geopandas/geodataframe.py | 5 ++ geopandas/io/sql.py | 122 ++++++++++++++++++++++----------- geopandas/io/tests/test_sql.py | 16 +++++ 3 files changed, 103 insertions(+), 40 deletions(-) diff --git a/geopandas/geodataframe.py b/geopandas/geodataframe.py index 35231fc..238c41d 100644 --- a/geopandas/geodataframe.py +++ b/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 diff --git a/geopandas/io/sql.py b/geopandas/io/sql.py index 61ff0f3..56512e0 100644 --- a/geopandas/io/sql.py +++ b/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 diff --git a/geopandas/io/tests/test_sql.py b/geopandas/io/tests/test_sql.py index 88c9394..bf3beff 100644 --- a/geopandas/io/tests/test_sql.py +++ b/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)