Add option to read_dataframe to use sql statement (#70)

Author: theroggy

CHANGES.md

@@ -11,6 +11,7 @@
 -   generalize check for VSI files from `/vsizip` to `/vsi` (#29)
 -   add dtype for each field to `read_info` (#30)
 -   support writing empty GeoDataFrames (#38)
+-   support use of a sql statement in read_dataframe (#70)
 
 ### Breaking changes
 

docs/source/introduction.md

@@ -41,10 +41,6 @@ configuration options.
 You can certainly try to read or write using unsupported drivers that are
 available in your installation, but you may encounter errors.
 
-Note: different drivers have different tolerance for mixed geometry types, e.g.,
-MultiPolygon and Polygon in the same dataset. You will get exceptions if you
-attempt to write mixed geometries to a driver that does not support them.
-
 ## List available layers
 
 To list layers available in a data source:
@@ -181,6 +177,65 @@ with the bbox.
 
 Note: the `bbox` values must be in the same CRS as the dataset.
 
+## Execute a sql query
+
+You can use the `sql` parameter to execute a sql query on a dataset. 
+
+Depending on the dataset, you can use different sql dialects. By default, if 
+the dataset natively supports sql, the sql statement will be passed through 
+as such. Hence, the sql query should be written in the relevant native sql
+dialect (e.g. [GeoPackage](https://gdal.org/drivers/vector/gpkg.html)/
+[Sqlite](https://gdal.org/drivers/vector/sqlite.html), 
+[PostgreSQL](https://gdal.org/drivers/vector/pg.html)). If the datasource 
+doesn't natively support sql (e.g. 
+[ESRI Shapefile](https://gdal.org/drivers/vector/shapefile.html), 
+[FlatGeobuf](https://gdal.org/drivers/vector/flatgeobuf.html)), you can choose 
+between '[OGRSQL](https://gdal.org/user/ogr_sql_dialect.html#ogr-sql-dialect)' 
+(the default) and  
+'[SQLITE](https://gdal.org/user/sql_sqlite_dialect.html#sql-sqlite-dialect)'. 
+For SELECT statements the 'SQLITE' dialect tends to provide more spatial 
+features as all 
+[spatialite](https://www.gaia-gis.it/gaia-sins/spatialite-sql-latest.html) 
+functions can be used. If gdal is not built with spatialite support in SQLite,
+you can use ``sql_dialect="INDIRECT_SQLITE"`` to be able to use spatialite 
+functions on native SQLite files like Geopackage. 
+
+You can combine a sql query with other parameters that will filter the 
+dataset. When using ``columns``, ``skip_features``, ``max_features``, and/or 
+``where`` it is important to note that they will be applied AFTER the sql 
+statement, so these are some things you need to be aware of:
+- if you specify an alias for a column in the sql statement, you need to 
+  specify this alias when using the ``columns`` keyword.
+- ``skip_features`` and ``max_features`` will be applied on the rows returned 
+  by the sql query, not on the original dataset.
+
+For the ``bbox`` parameter, depending on the combination of the dialect of the 
+sql query and the dataset, a spatial index will be used or not, e.g.:
+- ESRI Shapefile: spatial index is used with 'OGRSQL', not with 'SQLITE'.
+- Geopackage: spatial index is always used.
+
+The following sql query returns the 5 Western European countries with the most 
+neighbours:
+
+```python
+>>> sql = """
+        SELECT geometry, name,
+               (SELECT count(*)  
+                  FROM ne_10m_admin_0_countries layer_sub
+                 WHERE ST_Intersects(layer.geometry, layer_sub.geometry)) AS nb_neighbours
+          FROM ne_10m_admin_0_countries layer
+         WHERE subregion = 'Western Europe'
+         ORDER BY nb_neighbours DESC
+         LIMIT 5"""
+>>> read_dataframe('ne_10m_admin_0_countries.shp', sql=sql, sql_dialect='SQLITE')
+          NAME  nb_neighbours                            geometry
+0       France             11  MULTIPOLYGON (((-54.11153 2.114...
+1      Germany             10  MULTIPOLYGON (((13.81572 48.766...
+2      Austria              9  POLYGON ((16.94504 48.60417, 16...
+3  Switzerland              6  POLYGON ((10.45381 46.86443, 10...
+4      Belgium              5  POLYGON ((2.52180 51.08754, 2.5...
+```
+
 ## Force geometries to be read as 2D geometries
 
 You can force a 3D dataset to 2D using `force_2d`:

pyogrio/_io.pyx

@@ -171,6 +171,40 @@ cdef OGRLayerH get_ogr_layer(GDALDatasetH ogr_dataset, layer) except NULL:
     return ogr_layer
 
 
+cdef OGRLayerH execute_sql(GDALDatasetH ogr_dataset, str sql, str sql_dialect=None) except NULL:
+    """Execute an SQL statement on a dataset.
+
+    Parameters
+    ----------
+    ogr_dataset : pointer to open OGR dataset
+    sql : str
+        The sql statement to execute
+    sql_dialect : str, optional (default: None)
+        The sql dialect the sql statement is written in
+
+    Returns
+    -------
+    pointer to OGR layer
+    """
+
+    try:
+        sql_b = sql.encode('utf-8')
+        sql_c = sql_b
+        if sql_dialect is None:
+            return exc_wrap_pointer(GDALDatasetExecuteSQL(ogr_dataset, sql_c, NULL, NULL))      
+
+        sql_dialect_b = sql_dialect.encode('utf-8')
+        sql_dialect_c = sql_dialect_b
+        return exc_wrap_pointer(GDALDatasetExecuteSQL(ogr_dataset, sql_c, NULL, sql_dialect_c))    
+
+    # GDAL does not always raise exception messages in this case
+    except NullPointerError:
+        raise DataLayerError(f"Error executing sql '{sql}'") from None
+
+    except CPLE_BaseError as exc:
+        raise DataLayerError(str(exc))
+
+
 cdef str get_crs(OGRLayerH ogr_layer):
     """Read CRS from layer as EPSG:<code> if available or WKT.
 
@@ -586,7 +620,6 @@ cdef get_features(
         except CPLE_BaseError as exc:
             raise FeatureError(str(exc))
 
-
         if return_fids:
             fid_view[i] = OGR_F_GetFID(ogr_feature)
 
@@ -738,6 +771,8 @@ def ogr_read(
     object where=None,
     tuple bbox=None,
     object fids=None,
+    str sql=None,
+    str sql_dialect=None,
     int return_fids=False,
     **kwargs):
 
@@ -752,90 +787,100 @@ def ogr_read(
     path_b = path.encode('utf-8')
     path_c = path_b
 
-    # layer defaults to index 0
-    if layer is None:
-        layer = 0
-
     if fids is not None:
-        if where is not None or bbox is not None or skip_features or max_features:
+        if where is not None or bbox is not None or sql is not None or skip_features or max_features:
             raise ValueError(
-                "cannot set both 'fids' and any of 'where', 'bbox', "
+                "cannot set both 'fids' and any of 'where', 'bbox', 'sql', "
                 "'skip_features' or 'max_features'"
             )
         fids = np.asarray(fids, dtype=np.intc)
 
+    if sql is not None and layer is not None:
+        raise ValueError("'sql' paramater cannot be combined with 'layer'")
+
     ogr_dataset = ogr_open(path_c, 0, kwargs)
-    ogr_layer = get_ogr_layer(ogr_dataset, layer)
+    try:
+        if sql is None:
+            # layer defaults to index 0
+            if layer is None:
+                layer = 0    
+            ogr_layer = get_ogr_layer(ogr_dataset, layer)
+        else:
+            ogr_layer = execute_sql(ogr_dataset, sql, sql_dialect)
 
-    crs = get_crs(ogr_layer)
+        crs = get_crs(ogr_layer)
 
-    # Encoding is derived from the dataset, from the user, or from the system locale
-    encoding = (
-        detect_encoding(ogr_dataset, ogr_layer)
-        or encoding
-        or locale.getpreferredencoding()
-    )
+        # Encoding is derived from the dataset, from the user, or from the system locale
+        encoding = (
+            detect_encoding(ogr_dataset, ogr_layer)
+            or encoding
+            or locale.getpreferredencoding()
+        )
 
-    fields = get_fields(ogr_layer, encoding)
+        fields = get_fields(ogr_layer, encoding)
 
-    if columns is not None:
-        # Fields are matched exactly by name, duplicates are dropped.
-        # Find index of each field into fields
-        idx = np.intersect1d(fields[:,2], columns, return_indices=True)[1]
-        fields = fields[idx, :]
+        if columns is not None:
+            # Fields are matched exactly by name, duplicates are dropped.
+            # Find index of each field into fields
+            idx = np.intersect1d(fields[:,2], columns, return_indices=True)[1]
+            fields = fields[idx, :]
 
-    geometry_type = get_geometry_type(ogr_layer)
+        geometry_type = get_geometry_type(ogr_layer)
 
-    if fids is not None:
-        geometries, field_data = get_features_by_fid(
-            ogr_layer,
-            fids,
-            fields,
-            encoding,
-            read_geometry=read_geometry and geometry_type is not None,
-            force_2d=force_2d,
-        )
+        if fids is not None:
+            geometries, field_data = get_features_by_fid(
+                ogr_layer,
+                fids,
+                fields,
+                encoding,
+                read_geometry=read_geometry and geometry_type is not None,
+                force_2d=force_2d,
+            )
 
-        # bypass reading fids since these should match fids used for read
-        if return_fids:
-            fid_data = fids.astype(np.int64)
+            # bypass reading fids since these should match fids used for read
+            if return_fids:
+                fid_data = fids.astype(np.int64)
+            else:
+                fid_data = None
         else:
-            fid_data = None
-    else:
-        # Apply the attribute filter
-        if where is not None and where != "":
-            apply_where_filter(ogr_layer, where)
+            # Apply the attribute filter
+            if where is not None and where != "":
+                apply_where_filter(ogr_layer, where)
 
-        # Apply the spatial filter
-        if bbox is not None:
-            apply_spatial_filter(ogr_layer, bbox)
+            # Apply the spatial filter
+            if bbox is not None:
+                apply_spatial_filter(ogr_layer, bbox)
 
-        # Limit feature range to available range
-        skip_features, max_features = validate_feature_range(
-            ogr_layer, skip_features, max_features
-        )
+            # Limit feature range to available range
+            skip_features, max_features = validate_feature_range(
+                ogr_layer, skip_features, max_features
+            )
 
-        fid_data, geometries, field_data = get_features(
-            ogr_layer,
-            fields,
-            encoding,
-            read_geometry=read_geometry and geometry_type is not None,
-            force_2d=force_2d,
-            skip_features=skip_features,
-            max_features=max_features,
-            return_fids=return_fids
-        )
+            fid_data, geometries, field_data = get_features(
+                ogr_layer,
+                fields,
+                encoding,
+                read_geometry=read_geometry and geometry_type is not None,
+                force_2d=force_2d,
+                skip_features=skip_features,
+                max_features=max_features,
+                return_fids=return_fids
+            )
 
-    meta = {
-        'crs': crs,
-        'encoding': encoding,
-        'fields': fields[:,2], # return only names
-        'geometry_type': geometry_type
-    }
+        meta = {
+            'crs': crs,
+            'encoding': encoding,
+            'fields': fields[:,2], # return only names
+            'geometry_type': geometry_type
+        }
 
-    if ogr_dataset != NULL:
-        GDALClose(ogr_dataset)
-    ogr_dataset = NULL
+    finally:
+        if ogr_dataset != NULL:
+            if sql is not None:
+                GDALDatasetReleaseResultSet(ogr_dataset, ogr_layer)
+
+            GDALClose(ogr_dataset)
+            ogr_dataset = NULL
 
     return (
         meta,
@@ -1015,7 +1060,6 @@ cdef void * create_crs(str crs) except NULL:
     return ogr_crs
 
 
-
 cdef infer_field_types(list dtypes):
     cdef int field_type = 0
     cdef int field_subtype = 0

pyogrio/_ogr.pxd

@@ -324,6 +324,12 @@ cdef extern from "gdal.h":
     int             GDALDatasetGetLayerCount(GDALDatasetH ds)
     OGRLayerH       GDALDatasetGetLayer(GDALDatasetH ds, int iLayer)
     OGRLayerH       GDALDatasetGetLayerByName(GDALDatasetH ds, char * pszName)
+    OGRLayerH       GDALDatasetExecuteSQL(
+                            GDALDatasetH ds, 
+                            const char* pszStatement, 
+                            OGRGeometryH hSpatialFilter, 
+                            const char* pszDialect)
+    void            GDALDatasetReleaseResultSet(GDALDatasetH, OGRLayerH)
     OGRErr          GDALDatasetStartTransaction(GDALDatasetH ds, int bForce)
     OGRErr          GDALDatasetCommitTransaction(GDALDatasetH ds)
     OGRErr          GDALDatasetRollbackTransaction(GDALDatasetH ds)