Repository: geofileops/geofileops (python)

Issue:
ENH: add a vectorized version of apply
`apply_vectorized`: will apply a geospatial operation on the geometry column in one go for functions that support vectorized operations on an array of shapely geometries...

Pull request description:
ENH: add a vectorized version of apply
Resolves #588

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 92db4e0..8bad9cb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### Improvements
 
+- Add function `apply_vectorized` to apply a vectorized function on a geometry column
+  (#588)
 - Add support for renaming column with only difference in casing (#549)
 - Add support to pass a single string for all `column` type parameters if a single
   column should be retained (#523)
diff --git a/docs/reference.rst b/docs/reference.rst
index 438837f..8f31405 100644
--- a/docs/reference.rst
+++ b/docs/reference.rst
@@ -16,6 +16,7 @@ Geometry tools
    :toctree: api/
 
    apply
+   apply_vectorized
    buffer
    clip_by_geometry
    convexhull
diff --git a/geofileops/geoops.py b/geofileops/geoops.py
index cdd46d5..9ce832d 100644
--- a/geofileops/geoops.py
+++ b/geofileops/geoops.py
@@ -395,7 +395,7 @@ def apply(
     batchsize: int = -1,
     force: bool = False,
 ):
-    """Apply a python lambda function on the geometry column of the input file.
+    """Apply a python function on the geometry column of the input file.
 
     The result is written to the output file specified.
 
@@ -496,6 +496,108 @@ def apply(
     )
 
 
+def apply_vectorized(
+    input_path: Union[str, "os.PathLike[Any]"],
+    output_path: Union[str, "os.PathLike[Any]"],
+    func: Callable[[Any], Any],
+    input_layer: Optional[str] = None,
+    output_layer: Optional[str] = None,
+    columns: Optional[list[str]] = None,
+    explodecollections: bool = False,
+    force_output_geometrytype: Union[GeometryType, str, None] = None,
+    gridsize: float = 0.0,
+    keep_empty_geoms: bool = False,
+    where_post: Optional[str] = None,
+    nb_parallel: int = -1,
+    batchsize: int = -1,
+    force: bool = False,
+):
+    """Apply a vectorized python function on the geometry column of the input file.
+
+    The result is written to the output file specified.
+
+    It is not possible to use the contents of other columns in the input file in the
+    python function. If you need this, use ``gfo.apply`` instead.
+
+    If ``explodecollections`` is False and the input and output file type is GeoPackage,
+    the fid will be preserved. In other cases this will typically not be the case.
+
+    Args:
+        input_path (PathLike): the input file
+        output_path (PathLike): the file to write the result to
+        func (Callable): vectorized lambda function to apply to the geometry column.
+            Vectorized means here that the function should accept a shapely geometry
+            array as input and will return a shapely geometry for each item in the input
+            array.
+        input_layer (str, optional): input layer name. If None, ``input_path`` should
+            contain only one layer. Defaults to None.
+        output_layer (str, optional): output layer name. If None, the ``output_path``
+            stem is used. Defaults to None.
+        columns (List[str], optional): list of columns to retain. If None, all standard
+            columns are retained. In addition to standard columns, it is also possible
+            to specify "fid", a unique index available in all input files. Note that the
+            "fid" will be aliased eg. to "fid_1". Defaults to None.
+        explodecollections (bool, optional): True to output only simple geometries.
+            Defaults to False.
+        force_output_geometrytype (GeometryType, optional): The output geometry type to
+            force. If None, a best-effort guess is made and will always result in a
+            multi-type. Defaults to None.
+        gridsize (float, optional): the size of the grid the coordinates of the ouput
+            will be rounded to. Eg. 0.001 to keep 3 decimals. Value 0.0 doesn't change
+            the precision. Defaults to 0.0.
+        keep_empty_geoms (bool, optional): True to keep rows with empty/null geometries
+            in the output. Defaults to False.
+        where_post (str, optional): SQL filter to apply after all other processing,
+            including e.g. ``explodecollections``. It should be in sqlite syntax and
+            |spatialite_reference_link| functions can be used. Defaults to None.
+        nb_parallel (int, optional): the number of parallel processes to use.
+            Defaults to -1: use all available CPUs.
+        batchsize (int, optional): indicative number of rows to process per
+            batch. A smaller batch size, possibly in combination with a
+            smaller ``nb_parallel``, will reduce the memory usage.
+            Defaults to -1: (try to) determine optimal size automatically.
+        force (bool, optional): overwrite existing output file(s).
+            Defaults to False.
+
+    Examples:
+        This example shows the usage of ``gfo.apply_vectorized``:
+
+        .. code-block:: python
+
+            import geofileops as gfo
+
+            gfo.apply_vectorized(
+                input_path=...,
+                output_path=...,
+                func=lambda geom: pygeoops.centerline(geom, densify_distance=0),
+            )
+
+    .. |spatialite_reference_link| raw:: html
+
+        <a href="https://www.gaia-gis.it/gaia-sins/spatialite-sql-latest.html" target="_blank">spatialite reference</a>
+
+    """  # noqa: E501
+    logger = logging.getLogger("geofileops.apply_vectorized")
+    logger.info(f"Start on {input_path}")
+
+    return _geoops_gpd.apply_vectorized(
+        input_path=Path(input_path),
+        output_path=Path(output_path),
+        func=func,
+        input_layer=input_layer,
+        output_layer=output_layer,
+        columns=columns,
+        explodecollections=explodecollections,
+        force_output_geometrytype=force_output_geometrytype,
+        gridsize=gridsize,
+        keep_empty_geoms=keep_empty_geoms,
+        where_post=where_post,
+        nb_parallel=nb_parallel,
+        batchsize=batchsize,
+        force=force,
+    )
+
+
 def buffer(
     input_path: Union[str, "os.PathLike[Any]"],
     output_path: Union[str, "os.PathLike[Any]"],
diff --git a/geofileops/util/_geoops_gpd.py b/geofileops/util/_geoops_gpd.py
index fb4d248..f9a7eab 100644
--- a/geofileops/util/_geoops_gpd.py
+++ b/geofileops/util/_geoops_gpd.py
@@ -359,6 +359,7 @@ class GeoOperation(enum.Enum):
     BUFFER = "buffer"
     CONVEXHULL = "convexhull"
     APPLY = "apply"
+    APPLY_VECTORIZED = "apply_vectorized"
 
 
 def apply(
@@ -409,6 +410,50 @@ def apply(
     )
 
 
+def apply_vectorized(
+    input_path: Path,
+    output_path: Path,
+    func: Callable[[Any], Any],
+    operation_name: Optional[str] = None,
+    input_layer: Optional[str] = None,
+    output_layer: Optional[str] = None,
+    columns: Optional[list[str]] = None,
+    explodecollections: bool = False,
+    force_output_geometrytype: Union[GeometryType, str, None] = None,
+    gridsize: float = 0.0,
+    keep_empty_geoms: bool = False,
+    where_post: Optional[str] = None,
+    nb_parallel: int = -1,
+    batchsize: int = -1,
+    force: bool = False,
+    parallelization_config: Optional[ParallelizationConfig] = None,
+):
+    # Init
+    operation_params = {"pickled_func": cloudpickle.dumps(func)}
+    if operation_name is not None:
+        operation_params["operation_name"] = operation_name
+
+    # Go!
+    return _apply_geooperation_to_layer(
+        input_path=input_path,
+        output_path=output_path,
+        operation=GeoOperation.APPLY_VECTORIZED,
+        operation_params=operation_params,
+        input_layer=input_layer,
+        output_layer=output_layer,
+        columns=columns,
+        explodecollections=explodecollections,
+        force_output_geometrytype=force_output_geometrytype,
+        gridsize=gridsize,
+        keep_empty_geoms=keep_empty_geoms,
+        where_post=where_post,
+        nb_parallel=nb_parallel,
+        batchsize=batchsize,
+        force=force,
+        parallelization_config=parallelization_config,
+    )
+
+
 def buffer(
     input_path: Path,
     output_path: Path,
@@ -945,6 +990,9 @@ def _apply_geooperation(
                 data_gdf.geometry = data_gdf.geometry.apply(func)
             else:
                 data_gdf.geometry = data_gdf.apply(func, axis=1)
+        elif operation is GeoOperation.APPLY_VECTORIZED:
+            func = pickle.loads(operation_params["pickled_func"])
+            data_gdf.geometry = func(data_gdf.geometry)
         else:
             raise ValueError(f"operation not supported: {operation}")
 


Test patch:
diff --git a/tests/test_geofileops_singlelayer_gpd.py b/tests/test_geofileops_singlelayer_gpd.py
index a683e15..d2a6a5f 100644
--- a/tests/test_geofileops_singlelayer_gpd.py
+++ b/tests/test_geofileops_singlelayer_gpd.py
@@ -211,6 +211,73 @@ def test_apply_geooperation_invalid_operation(tmp_path):
         )
 
 
+@pytest.mark.parametrize("suffix", SUFFIXES_GEOOPS)
+def test_apply_vectorized(tmp_path, suffix):
+    # Prepare test data
+    test_gdf = gpd.GeoDataFrame(
+        data=[
+            {"id": 1, "geometry": test_helper.TestData.polygon_small_island},
+            {"id": 2, "geometry": test_helper.TestData.polygon_with_island},
+            {"id": 3, "geometry": None},
+        ],
+        crs=31370,
+    )
+    input_path = tmp_path / f"polygons_small_holes_{suffix}"
+    gfo.to_file(test_gdf, input_path)
+    output_path = tmp_path / f"{input_path.stem}-output{suffix}"
+    input_layerinfo = gfo.get_layerinfo(input_path)
+    batchsize = math.ceil(input_layerinfo.featurecount / 2)
+    min_area = 2
+
+    def remove_inner_rings_vectorized(geometry, min_area_to_keep, crs=None):
+        return [
+            pygeoops.remove_inner_rings(
+                geom, min_area_to_keep=min_area_to_keep, crs=crs
+            )
+            for geom in geometry
+        ]
+
+    # Run test
+    output_gdf = gfo.apply_vectorized(
+        input_path=input_path,
+        output_path=output_path,
+        func=lambda geom: remove_inner_rings_vectorized(
+            geom, min_area_to_keep=min_area
+        ),
+        batchsize=batchsize,
+    )
+
+    # Now check if the output file is correctly created
+    assert output_path.exists()
+    exp_spatial_index = GeofileInfo(output_path).default_spatial_index
+    assert gfo.has_spatial_index(output_path) is exp_spatial_index
+
+    # Read result for some more detailed checks
+    output_gdf = gfo.read_file(output_path).sort_values("id").reset_index(drop=True)
+    output_layerinfo = gfo.get_layerinfo(output_path)
+
+    assert len(output_layerinfo.columns) == len(input_layerinfo.columns)
+
+    for index in range(2):
+        output_geometry = output_gdf["geometry"][index]
+        if index == 2:
+            assert output_geometry is None
+            continue
+        else:
+            assert output_geometry is not None
+        if isinstance(output_geometry, sh_geom.MultiPolygon):
+            assert len(output_geometry.geoms) == 1
+            output_geometry = output_geometry.geoms[0]
+        assert isinstance(output_geometry, sh_geom.Polygon)
+
+        if index == 0:
+            # In the 1st polygon the island must be removed
+            assert len(output_geometry.interiors) == 0
+        elif index == 1:
+            # In the 2nd polygon the island is larger, so should be there
+            assert len(output_geometry.interiors) == 1
+
+
 @pytest.mark.parametrize(
     "suffix, epsg", [(".gpkg", 31370), (".gpkg", 4326), (".shp", 31370)]
 )