point splitter

CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 
 ## [0.54.0] - 2026-10-28
+- Add public `group_geometries` helper in `openeo.extra.job_management` for geometry-grouping based job splitting, including optional tile-grid co-grouping.
 
 ### Changed
 

docs/cookbook/job_manager.rst

@@ -202,6 +202,7 @@ to resume where you left off.
     Parquet support requires the ``pyarrow`` package
     (see :ref:`optional dependencies <installation-optional-dependencies>`).
 
+
 STAC API (experimental)
 -----------------------
 
@@ -247,7 +248,7 @@ Example: splitting a bounding box into 10x10 km tiles in EPSG:3857:
 
     # Split into 10x10 km tiles in EPSG:3857
     gdf = split_area(
-        aoi=extent,
+        spatial_extent=extent,
         tile_size=10_000,  # tile size in meters (units of the projection)
         projection="EPSG:3857",  # projection for tiling
     )
@@ -300,7 +301,80 @@ or any other custom partitioning, you can pass a pre-computed
     tile_grid = gpd.read_parquet("my_tile_grid.parquet")
 
     # Only tiles intersecting the AOI are returned
-    gdf = split_area(aoi=extent , tile_grid=tile_grid)
+    gdf = split_area(spatial_extent=extent, tile_grid=tile_grid)
+
+Grouping geometries (e.g. point-based extractions)
+--------------------------------------------------
+
+When the work to split is a set of geometries (e.g. sample locations)
+rather than a single large area, e.g. for point-based extraction jobs,
+:py:func:`~openeo.extra.job_management._job_splitting.group_geometries` groups
+the rows of a GeoDataFrame into batches, optionally of at most ``max_geometries`` rows.
+Each row is treated as a single feature to group, based on its representative point,
+regardless of geometry type (e.g. ``Point``, ``MultiPoint``, ``Polygon``).
+Null or empty geometries are not allowed.
+Geometries are not split themselves: use ``split_area`` for a large area of interest.
+Like with ``split_area``, geometries without
+CRS information are assumed to be in the CRS of the tile grid (if any).
+Groups are formed by recursive bisection at the median along the longer axis,
+so that batches stay geographically clustered and balanced.
+The returned GeoDataFrames keep the original index, to link back to the input rows:
+
+.. code-block:: python
+
+    import geopandas as gpd
+    from openeo.extra.job_management import group_geometries
+
+    points = gpd.read_file("my_sample_locations.gpkg")
+    batches = group_geometries(points, max_geometries=500)
+    # batches is a list of GeoDataFrames, each with at most 500 rows
+
+Optionally, a ``tile_grid`` (e.g. a satellite tiling grid) can be passed so
+that each feature is first assigned to the tile containing its representative point
+and batches are formed per tile.
+Features outside every tile are batched separately:
+
+.. code-block:: python
+
+    tile_grid = gpd.read_parquet("my_tile_grid.parquet")
+    batches = group_geometries(points, max_geometries=500, tile_grid=tile_grid)
+
+Alternatively, a regular grid can be specified through ``tile_size`` and ``projection``.
+This grid is anchored at the origin of the projection,
+so the grouping of existing features does not change when other features are added:
+
+.. code-block:: python
+
+    # group points per 50x50 km tile in EPSG:3857
+    batches = group_geometries(points, max_geometries=500, tile_size=50_000, projection="EPSG:3857")
+
+Without ``max_geometries`` (default ``None``), there is no size limit and
+geometries are simply grouped per tile:
+
+.. code-block:: python
+
+    # one batch per tile
+    batches = group_geometries(points, tile_grid=tile_grid)
+
+Current limitations:
+
+- Geometries are never split: a polygon crossing tile edges, or larger than a tile,
+  ends up in a single batch and extends beyond its tile.
+  Use ``split_area`` for large areas of interest.
+- Multi-part geometries (e.g. ``MultiPolygon``) are kept together as one feature,
+  placed by a single representative point, even if their parts are far apart.
+  To group parts individually, explode them into separate rows first
+  (the original index is kept, to link back to the input rows;
+  note that a single feature may then be spread over multiple batches):
+
+  .. code-block:: python
+
+      batches = group_geometries(points.explode(index_parts=False), max_geometries=500)
+
+- Batches are limited by number of geometries only, not by area or processing cost.
+  Use a tile grid to limit the spatial extent of batches.
+- With overlapping tiles (e.g. Sentinel-2 MGRS), a feature is assigned to the first tile
+  (in grid order) containing its representative point, not to the tile with the largest overlap.
 
 Customizing Job Handling
 ========================
@@ -543,6 +617,7 @@ Job Database
 
 .. autoclass:: openeo.extra.job_management.ParquetJobDatabase
 
+
 .. autofunction:: openeo.extra.job_management.create_job_db
 
 .. autofunction:: openeo.extra.job_management.get_job_db

openeo/extra/job_management/__init__.py

@@ -6,7 +6,7 @@
     create_job_db,
     get_job_db,
 )
-from openeo.extra.job_management._job_splitting import split_area
+from openeo.extra.job_management._job_splitting import group_geometries, split_area
 from openeo.extra.job_management._manager import MultiBackendJobManager
 from openeo.extra.job_management.process_based import ProcessBasedJobCreator
 
@@ -19,5 +19,6 @@
     "create_job_db",
     "get_job_db",
     "split_area",
+    "group_geometries",
     "MultiBackendJobManager",
 ]

openeo/extra/job_management/_job_splitting.py

@@ -4,6 +4,8 @@
 from typing import Dict, List, Optional, Tuple, Union
 
 import geopandas as gpd
+import numpy as np
+import pandas as pd
 import shapely
 import shapely.geometry.base
 from shapely.geometry import MultiPolygon, Polygon
@@ -12,6 +14,7 @@
 
 _log = logging.getLogger(__name__)
 
+
 class JobSplittingFailure(Exception):
     pass
 
@@ -72,25 +75,6 @@ def _parse_input_geometry(
                     "Invalid bounding box: west must not equal east (zero-width bounding box)."
                 )
 
-            # if bbox["west"] > bbox["east"]:
-            #     # In a geographic CRS this indicates an antimeridian crossing;
-            #     # split into two polygons on each side of the ±180° meridian.
-            #     if source_epsg is not None and _is_geographic_crs(source_epsg):
-            #         _log.info(
-            #             "Bounding box crosses the antimeridian (west=%s > east=%s); "
-            #             "splitting into two polygons.",
-            #             bbox["west"],
-            #             bbox["east"],
-            #         )
-            #         west_part = shapely.geometry.box(bbox["west"], bbox["south"], 180.0, bbox["north"])
-            #         east_part = shapely.geometry.box(-180.0, bbox["south"], bbox["east"], bbox["north"])
-            #         return MultiPolygon([west_part, east_part]), source_epsg
-            #     raise JobSplittingFailure(
-            #         "Invalid bounding box: west must be less than east. "
-            #         "Antimeridian-crossing bounding boxes are only supported "
-            #         "for geographic CRSs (e.g. EPSG:4326) with an explicit 'crs' field."
-            #     )
-
             return bbox.as_geometry(), source_epsg
         elif isinstance(geometry, (Polygon, MultiPolygon)):
             return geometry, None
@@ -374,3 +358,178 @@ def split_area(
 
     grid = _SizeBasedTileGrid(epsg=normalize_crs(projection), size=tile_size)
     return grid.get_tiles(aoi)
+
+
+def _predefined_tile_keys(tile_grid: "_PredefinedTileGrid", points: np.ndarray) -> np.ndarray:
+    """For each point (in the grid CRS), the position of the first tile containing it (-1: none)."""
+    # TODO: polygons are assigned by their representative point only, so they may extend beyond their tile.
+    #   Assigning them to e.g. the tile with the largest overlap could be a future addition.
+    n_tiles = len(tile_grid._gdf)
+    point_idx, tile_idx = tile_grid._gdf.sindex.query(points, predicate="intersects")
+    # Lowest tile position wins for points on shared tile edges.
+    keys = np.full(len(points), n_tiles)
+    np.minimum.at(keys, point_idx, tile_idx)
+    keys[keys == n_tiles] = -1
+    return keys
+
+
+def _bisect_group(
+    positions: np.ndarray,
+    *,
+    xs: np.ndarray,
+    ys: np.ndarray,
+    max_geometries: Optional[int],
+) -> List[np.ndarray]:
+    """
+    Recursively bisect a group of rows (given as positions) at the median of their locations,
+    along the longer axis of the locations' spread, until groups have at most *max_geometries* rows.
+    """
+    result = []
+    stack = [positions]
+    while stack:
+        pos = stack.pop()
+        if max_geometries is None or len(pos) <= max_geometries:
+            result.append(pos)
+            continue
+        px, py = xs[pos], ys[pos]
+        axis = px if np.ptp(px) >= np.ptp(py) else py
+        pos = pos[np.argsort(axis, kind="stable")]
+        cut = len(pos) // 2
+        stack.append(pos[cut:])
+        stack.append(pos[:cut])
+    return result
+
+
+def group_geometries(
+    geometries: gpd.GeoDataFrame,
+    *,
+    max_geometries: Optional[int] = None,
+    projection: Optional[str] = None,
+    tile_size: Optional[float] = None,
+    tile_grid: Optional[Union[_PredefinedTileGrid, _SizeBasedTileGrid, gpd.GeoDataFrame]] = None,
+) -> List[gpd.GeoDataFrame]:
+    """
+    Group geometries (e.g. sample locations for point-based extraction)
+    into spatially clustered groups of at most *max_geometries* each.
+
+    Each row is treated as a single "feature" to group, regardless of its
+    geometry type (e.g. ``Point``, ``MultiPoint``, ``Polygon``):
+    *max_geometries* limits the number of rows (features) per group,
+    not the number of individual coordinates. To limit individual
+    coordinates instead, explode multi-part geometries into separate rows
+    first, e.g. with ``geometries.explode(index_parts=False)``.
+
+    Features are located by their ``representative_point()`` (computed in the input CRS),
+    they are not split themselves. For a single large area of interest that has to be split,
+    use :func:`split_area` instead.
+
+    Grouping is done by recursive bisection (k-d tree style): a group is cut at the
+    median of the feature locations along the longer axis of their spread,
+    until it has at most *max_geometries* features.
+    This keeps groups spatially clustered and balanced.
+
+    If *tile_grid* is given (e.g. a Sentinel-2 MGRS/UTM tile grid), each feature is first
+    assigned to the tile containing its representative point, and every tile is bisected separately.
+    Features whose representative point is not in any tile are grouped (and bisected) separately.
+    Alternatively, *tile_size* and *projection* can be given to group
+    features in a regular grid of square tiles anchored at the origin of the projection
+    (so grouping does not change when other features are added).
+
+    Current limitations:
+
+    - Features are never split: a polygon crossing tile edges, or larger than a tile,
+      is assigned to a single tile and extends beyond it. Use :func:`split_area` for large areas.
+    - Groups are limited by feature count only, not by area or processing cost.
+      Use a tile grid to limit the spatial extent of groups.
+    - Multi-part geometries are placed by a single representative point,
+      so their parts can be far away from the rest of their group.
+    - With overlapping tiles, a feature is assigned to the first tile (in grid order)
+      containing its representative point, not to the tile with the largest overlap.
+
+    :param geometries: GeoDataFrame with a ``geometry`` column.
+        Null or empty geometries are not allowed.
+    :param max_geometries: maximum number of geometries per group.
+        ``None`` (default) means no limit.
+    :param projection: EPSG string (e.g. ``"EPSG:3857"``) of the regular grid
+        to group by. Must be combined with *tile_size*.
+    :param tile_size: tile edge length in the unit of *projection*.
+        Must be combined with *projection*.
+    :param tile_grid: optional :class:`_PredefinedTileGrid` or
+        :class:`_SizeBasedTileGrid` instance, or GeoDataFrame of tile polygons
+        (with CRS set) to co-join geometries by before splitting, e.g. a
+        satellite tiling grid.
+        Mutually exclusive with *projection* / *tile_size*.
+    :return: list of GeoDataFrames, with the original index of *geometries* preserved.
+    :raises JobSplittingFailure: on invalid input (e.g. non-positive
+        *max_geometries*, null/empty geometries, contradictory grid arguments).
+
+    Like with :func:`split_area`, geometries without CRS information
+    are assumed to be in the tile grid's CRS.
+    Without tile grid, no CRS is needed and the returned groups have no CRS either.
+    """
+    if not isinstance(geometries, gpd.GeoDataFrame):
+        raise JobSplittingFailure(f"Expected a GeoDataFrame, got {type(geometries).__name__}.")
+    if "geometry" not in geometries.columns:
+        raise JobSplittingFailure("The GeoDataFrame must contain a 'geometry' column.")
+    if max_geometries is not None and max_geometries <= 0:
+        raise JobSplittingFailure(f"'max_geometries' must be positive, got {max_geometries!r}.")
+
+    if tile_grid is not None:
+        if projection is not None or tile_size is not None:
+            raise JobSplittingFailure(
+                "Cannot combine 'tile_grid' with 'projection' or 'tile_size'. "
+                "Either pass a tile_grid, or pass projection + tile_size."
+            )
+        if isinstance(tile_grid, gpd.GeoDataFrame):
+            tile_grid = _PredefinedTileGrid(tiles=tile_grid)
+        elif not isinstance(tile_grid, (_PredefinedTileGrid, _SizeBasedTileGrid)):
+            raise JobSplittingFailure(
+                "Unsupported tile_grid type: "
+                f"{type(tile_grid).__name__}. Expected _PredefinedTileGrid, _SizeBasedTileGrid, or GeoDataFrame."
+            )
+    elif projection is not None or tile_size is not None:
+        if projection is None or tile_size is None:
+            raise JobSplittingFailure("'projection' and 'tile_size' must be provided together.")
+        tile_grid = _SizeBasedTileGrid.from_size_projection(size=tile_size, projection=projection)
+
+    n_invalid = int((geometries.geometry.isna() | geometries.geometry.is_empty).sum())
+    if n_invalid:
+        raise JobSplittingFailure(f"Found {n_invalid} null or empty geometries, which cannot be grouped.")
+
+    if geometries.crs is None and tile_grid is not None:
+        _log.warning(
+            "Geometries have no CRS information; assuming they are already in the tile grid's CRS (EPSG:%d).",
+            tile_grid.crs,
+        )
+        geometries = geometries.set_crs(f"EPSG:{tile_grid.crs}")
+    if geometries.crs is not None and geometries.crs.to_epsg() is not None:
+        _TileGridInterface._check_antimeridian_crossing(geometries.total_bounds, geometries.crs.to_epsg())
+
+    if geometries.empty:
+        return []
+
+    # TODO: limiting groups by area or processing cost (e.g. a weight per feature) is not supported yet.
+    # Work with positions (not index labels), so duplicated index labels are no problem.
+    locations = geometries.geometry.representative_point()
+    if tile_grid is not None:
+        locations = locations.to_crs(epsg=tile_grid.crs)
+    xs, ys = locations.x.to_numpy(), locations.y.to_numpy()
+
+    all_positions = np.arange(len(geometries))
+    if tile_grid is None:
+        position_groups = [all_positions]
+    else:
+        if isinstance(tile_grid, _SizeBasedTileGrid):
+            # Regular grid anchored at the origin of the CRS.
+            keys = pd.DataFrame({"x": np.floor(xs / tile_grid.size), "y": np.floor(ys / tile_grid.size)})
+        elif isinstance(tile_grid, _PredefinedTileGrid):
+            keys = pd.DataFrame({"tile": _predefined_tile_keys(tile_grid, locations.to_numpy())})
+        else:
+            raise JobSplittingFailure(f"Unsupported tile grid type: {type(tile_grid).__name__}.")
+        position_groups = list(keys.groupby(list(keys.columns), sort=True).indices.values())
+
+    result = []
+    for positions in position_groups:
+        for group in _bisect_group(positions, xs=xs, ys=ys, max_geometries=max_geometries):
+            result.append(geometries.iloc[np.sort(group)])
+    return result