ENH: Calculate feature count ourselves if GDAL returns unknown count (-1), better support OSM data

CHANGES.md

@@ -6,11 +6,18 @@
 
 -   Support writing dataframes without geometry column (#267)
 
+-   Calculate feature count by iterating over features if GDAL returns an
+    unknown count for a data layer (e.g., OSM driver); this may have signficant
+    performance impacts for some data sources that would otherwise return an
+    unknown count (count is used in `read_info`, `read`, `read_dataframe`) (#271).
+
 ### Bug fixes
 
 -   Fix int32 overflow when reading int64 columns (#260)
 -   Fix `fid_as_index=True` doesn't set fid as index using `read_dataframe` with
     `use_arrow=True` (#265)
+-   Fix errors reading OSM data due to invalid feature count and incorrect
+    reading of OSM layers beyond the first layer (#271)
 
 ## 0.6.0 (2023-04-27)
 

docs/source/known_issues.md

@@ -65,3 +65,45 @@ convert the datetimes to string.
 
 Timezone information is ignored at the moment, both when reading and when writing
 datetime columns.
+
+## Support for OpenStreetMap (OSM) data
+
+OpenStreetMap data do not natively support calculating the feature count by data
+layer due to the internal data structures. To get around this, Pyogrio iterates
+over all features first to calculate the feature count that is used to allocate
+arrays that contain the geometries and attributes read from the data layer, and
+then iterates over all feature again to populate those arrays. Further, data
+within the file are not structured at the top level to support fast reading by
+layer, which means that reading data by layer may need to read all records
+within the data source, not just those belonging to a particular layer. This is
+inefficient and slow, and is exacerbated when attemping to read from
+remotely-hosted data sources rather than local files.
+
+You may also be instructed by GDAL to enable interleaved reading mode via an
+error message when you try to read a large file without it, which you can do in
+one of two ways:
+
+1. Set config option used for all operations
+
+```python
+from pyogrio import set_gdal_config_options
+
+set_gdal_config_options({"OGR_INTERLEAVED_READING": True})
+```
+
+2. Set dataset open option
+
+```python
+
+from pyogrio import read_dataframe
+
+df = read_dataframe(path, INTERLEAVED_READING=True)
+```
+
+We recommend the following to sidestep performance issues:
+
+-   download remote OSM data sources to local files before attempting
+    to read
+-   the `use_arrow=True` option may speed up reading from OSM files
+-   if possible, use a different tool such as `ogr2ogr` to translate the OSM
+    data source into a more performant format for reading by layer, such as GPKG

pyogrio/_io.pyx

@@ -195,6 +195,17 @@ cdef OGRLayerH get_ogr_layer(GDALDatasetH ogr_dataset, layer) except NULL:
     except CPLE_BaseError as exc:
         raise DataLayerError(str(exc))
 
+    # if the driver is OSM, we need to execute SQL to set the layer to read in
+    # order to read it properly
+    if get_driver(ogr_dataset) == "OSM":
+        # Note: this returns NULL and does not need to be freed via
+        # GDALDatasetReleaseResultSet()
+        layer_name = get_string(OGR_L_GetName(ogr_layer))
+        sql_b = f"SET interest_layers = {layer_name}".encode('utf-8')
+        sql_c = sql_b
+
+        GDALDatasetExecuteSQL(ogr_dataset, sql_c, NULL, NULL)
+
     return ogr_layer
 
 
@@ -309,6 +320,62 @@ cdef get_driver(OGRDataSourceH ogr_dataset):
     return driver
 
 
+cdef get_feature_count(OGRLayerH ogr_layer):
+    """Get the feature count of a layer.
+
+    If GDAL returns an unknown count (-1), this iterates over every feature
+    to calculate the count.
+
+    Parameters
+    ----------
+    ogr_layer : pointer to open OGR layer
+
+    Returns
+    -------
+    int
+        count of features
+    """
+
+    cdef OGRFeatureH ogr_feature = NULL
+    cdef int feature_count = OGR_L_GetFeatureCount(ogr_layer, 1)
+
+    # if GDAL refuses to give us the feature count, we have to loop over all
+    # features ourselves and get the count.  This can happen for some drivers
+    # (e.g., OSM) or if a where clause is invalid but not rejected as error
+    if feature_count == -1:
+        # make sure layer is read from beginning
+        OGR_L_ResetReading(ogr_layer)
+
+        feature_count = 0
+        while True:
+            try:
+                ogr_feature = exc_wrap_pointer(OGR_L_GetNextFeature(ogr_layer))
+                feature_count +=1
+
+            except NullPointerError:
+                # No more rows available, so stop reading
+                break
+
+            # driver may raise other errors, e.g., for OSM if node ids are not
+            # increasing, the default config option OSM_USE_CUSTOM_INDEXING=YES
+            # causes errors iterating over features
+            except CPLE_BaseError as exc:
+                # if an invalid where clause is used for a GPKG file, it is not
+                # caught as an error until attempting to iterate over features;
+                # catch it here
+                if "failed to prepare SQL" in str(exc):
+                    raise ValueError(f"Invalid SQL query: {str(exc)}") from None
+
+                raise DataLayerError(f"Could not iterate over features: {str(exc)}") from None
+
+            finally:
+                if ogr_feature != NULL:
+                    OGR_F_Destroy(ogr_feature)
+                    ogr_feature = NULL
+
+    return feature_count
+
+
 cdef set_metadata(GDALMajorObjectH obj, object metadata):
     """Set metadata on a dataset or layer
 
@@ -371,13 +438,19 @@ cdef detect_encoding(OGRDataSourceH ogr_dataset, OGRLayerH ogr_layer):
     -------
     str or None
     """
+
     if OGR_L_TestCapability(ogr_layer, OLCStringsAsUTF8):
         return 'UTF-8'
 
     driver = get_driver(ogr_dataset)
     if driver == 'ESRI Shapefile':
         return 'ISO-8859-1'
 
+    if driver == "OSM":
+        # always set OSM data to UTF-8
+        # per https://help.openstreetmap.org/questions/2172/what-encoding-does-openstreetmap-use
+        return "UTF-8"
+
     return None
 
 
@@ -486,11 +559,10 @@ cdef apply_where_filter(OGRLayerH ogr_layer, str where):
     if err != OGRERR_NONE:
         try:
             exc_check()
-            name = OGR_L_GetName(ogr_layer)
         except CPLE_BaseError as exc:
             raise ValueError(str(exc))
 
-        raise ValueError(f"Invalid SQL query for layer '{name}': '{where}'")
+        raise ValueError(f"Invalid SQL query for layer '{OGR_L_GetName(ogr_layer)}': '{where}'")
 
 
 cdef apply_spatial_filter(OGRLayerH ogr_layer, bbox):
@@ -526,11 +598,10 @@ cdef validate_feature_range(OGRLayerH ogr_layer, int skip_features=0, int max_fe
     skip_features : number of features to skip from beginning of available range
     max_features : maximum number of features to read from available range
     """
-    feature_count = OGR_L_GetFeatureCount(ogr_layer, 1)
+    feature_count = get_feature_count(ogr_layer)
     num_features = max_features
 
-    if feature_count <= 0:
-        # the count comes back as -1 if the where clause above is invalid but not rejected as error
+    if feature_count == 0:
         name = OGR_L_GetName(ogr_layer)
         warnings.warn(f"Layer '{name}' does not have any features to read")
         return 0, 0
@@ -736,9 +807,6 @@ cdef get_features(
                 break
 
             except CPLE_BaseError as exc:
-                if "failed to prepare SQL" in str(exc):
-                    raise ValueError(f"Invalid SQL query") from exc
-
                 raise FeatureError(str(exc))
 
             if i >= num_features:
@@ -886,10 +954,7 @@ cdef get_bounds(
                 break
 
             except CPLE_BaseError as exc:
-                if "failed to prepare SQL" in str(exc):
-                    raise ValueError(f"Invalid SQL query") from exc
-                else:
-                    raise FeatureError(str(exc))
+                raise FeatureError(str(exc))
 
             if i >= num_features:
                 raise FeatureError(
@@ -1327,7 +1392,7 @@ def ogr_read_info(
             'fields': fields[:,2], # return only names
             'dtypes': fields[:,3],
             'geometry_type': get_geometry_type(ogr_layer),
-            'features': OGR_L_GetFeatureCount(ogr_layer, 1),
+            'features': get_feature_count(ogr_layer),
             'driver': get_driver(ogr_dataset),
             "capabilities": {
                 "random_read": OGR_L_TestCapability(ogr_layer, OLCRandomRead),

pyogrio/tests/fixtures/README.md

@@ -83,3 +83,7 @@ This was extracted from https://prd-tnm.s3.amazonaws.com/StagedProducts/Hydrogra
 ```bash
 ogr2ogr test_mixed_surface.gpkg NHDPLUS_H_0308_HU4_GDB.gdb NHDWaterbody -where '"NHDPlusID" = 15000300070477' -select "NHDPlusID"
 ```
+
+### OSM PBF test
+
+This was downloaded from https://github.com/openstreetmap/OSM-binary/blob/master/resources/sample.pbf

pyogrio/tests/test_core.py

@@ -12,7 +12,7 @@
     get_gdal_config_option,
     get_gdal_data_path,
 )
-from pyogrio.errors import DataSourceError
+from pyogrio.errors import DataSourceError, DataLayerError
 
 from pyogrio._env import GDALEnv
 
@@ -266,6 +266,23 @@ def test_read_info_invalid_dataset_kwargs(naturalearth_lowres):
         read_info(naturalearth_lowres, INVALID="YES")
 
 
+def test_read_info_force_feature_count_exception(data_dir):
+    with pytest.raises(DataLayerError, match="Could not iterate over features"):
+        read_info(data_dir / "sample.osm.pbf", layer="lines")
+
+
+def test_read_info_force_feature_count(data_dir):
+    # the sample OSM file has non-increasing node IDs which causes the default
+    # custom indexing to raise an exception iterating over features
+    meta = read_info(data_dir / "sample.osm.pbf", USE_CUSTOM_INDEXING=False)
+    assert meta["features"] == 8
+
+    meta = read_info(
+        data_dir / "sample.osm.pbf", layer="lines", USE_CUSTOM_INDEXING=False
+    )
+    assert meta["features"] == 36
+
+
 @pytest.mark.parametrize(
     "name,value,expected",
     [