ssb-sgis 1.1.17__py3-none-any.whl → 1.2.1__py3-none-any.whl

sgis/geopandas_tools/sfilter.py

@@ -4,11 +4,12 @@ import numpy as np
 import pandas as pd
 from geopandas import GeoDataFrame
 from geopandas import GeoSeries
-from geopandas import __version__ as geopandas_version
 from shapely import Geometry
-from shapely import STRtree
 
+from ..conf import _get_instance
+from ..conf import config
 from .conversion import to_gdf
+from .runners import RTreeQueryRunner
 
 gdf_type_error_message = "'gdf' should be of type GeoDataFrame or GeoSeries."
 
@@ -18,6 +19,8 @@ def sfilter(
     other: GeoDataFrame | GeoSeries | Geometry,
     predicate: str = "intersects",
     distance: int | float | None = None,
+    n_jobs: int | None = None,
+    rtree_runner: RTreeQueryRunner | None = None,
 ) -> GeoDataFrame:
     """Filter a GeoDataFrame or GeoSeries by spatial predicate.
 
@@ -33,6 +36,9 @@ def sfilter(
         other: The geometry object to filter 'gdf' by.
         predicate: Spatial predicate to use. Defaults to 'intersects'.
         distance: Max distance to allow if predicate=="dwithin".
+        n_jobs: Number of workers.
+        rtree_runner: Optionally debug/manipulate the spatial indexing operations.
+            See the 'runners' module for example implementations.
 
     Returns:
         A copy of 'gdf' with only the rows matching the
@@ -80,7 +86,9 @@ def sfilter(
 
     other = _sfilter_checks(other, crs=gdf.crs)
 
-    indices = _get_sfilter_indices(gdf, other, predicate, distance)
+    indices = _get_sfilter_indices(
+        gdf, other, predicate, distance, n_jobs, rtree_runner
+    )
 
     return gdf.iloc[indices]
 
@@ -90,6 +98,8 @@ def sfilter_split(
     other: GeoDataFrame | GeoSeries | Geometry,
     predicate: str = "intersects",
     distance: int | float | None = None,
+    n_jobs: int = 1,
+    rtree_runner: RTreeQueryRunner | None = None,
 ) -> tuple[GeoDataFrame, GeoDataFrame]:
     """Split a GeoDataFrame or GeoSeries by spatial predicate.
 
@@ -101,6 +111,9 @@ def sfilter_split(
         other: The geometry object to filter 'gdf' by.
         predicate: Spatial predicate to use. Defaults to 'intersects'.
         distance: Max distance to allow if predicate=="dwithin".
+        n_jobs: Number of workers.
+        rtree_runner: Optionally debug/manipulate the spatial indexing operations.
+            See the 'runners' module for example implementations.
 
     Returns:
         A tuple of GeoDataFrames, one with the rows that match the spatial predicate
@@ -151,7 +164,9 @@ def sfilter_split(
 
     other = _sfilter_checks(other, crs=gdf.crs)
 
-    indices = _get_sfilter_indices(gdf, other, predicate, distance)
+    indices = _get_sfilter_indices(
+        gdf, other, predicate, distance, n_jobs, rtree_runner
+    )
 
     return (
         gdf.iloc[indices],
@@ -164,6 +179,8 @@ def sfilter_inverse(
     other: GeoDataFrame | GeoSeries | Geometry,
     predicate: str = "intersects",
     distance: int | float | None = None,
+    n_jobs: int = 1,
+    rtree_runner: RTreeQueryRunner | None = None,
 ) -> GeoDataFrame | GeoSeries:
     """Filter a GeoDataFrame or GeoSeries by inverse spatial predicate.
 
@@ -174,6 +191,9 @@ def sfilter_inverse(
         other: The geometry object to filter 'gdf' by.
         predicate: Spatial predicate to use. Defaults to 'intersects'.
         distance: Max distance to allow if predicate=="dwithin".
+        n_jobs: Number of workers.
+        rtree_runner: Optionally debug/manipulate the spatial indexing operations.
+            See the 'runners' module for example implementations.
 
     Returns:
         A copy of 'gdf' with only the rows that do not match the
@@ -215,11 +235,10 @@ def sfilter_inverse(
     """
     if not isinstance(gdf, (GeoDataFrame | GeoSeries)):
         raise TypeError(gdf_type_error_message)
-
     other = _sfilter_checks(other, crs=gdf.crs)
-
-    indices = _get_sfilter_indices(gdf, other, predicate, distance)
-
+    indices = _get_sfilter_indices(
+        gdf, other, predicate, distance, n_jobs, rtree_runner
+    )
     return gdf.iloc[pd.Index(range(len(gdf))).difference(pd.Index(indices))]
 
 
@@ -252,6 +271,8 @@ def _get_sfilter_indices(
     right: GeoDataFrame | GeoSeries | Geometry,
     predicate: str,
     distance: int | float | None,
+    n_jobs: int,
+    rtree_runner: RTreeQueryRunner | None,
 ) -> np.ndarray:
     """Compute geometric comparisons and get matching indices.
 
@@ -264,6 +285,9 @@ def _get_sfilter_indices(
     right : GeoDataFrame
     predicate : string
         Binary predicate to query.
+    n_jobs: Number of workers.
+    rtree_runner: Optionally debug/manipulate the spatial indexing operations.
+        See the 'runners' module for example implementations.
 
     Returns:
     -------
@@ -273,6 +297,9 @@ def _get_sfilter_indices(
     """
     original_predicate = predicate
 
+    if rtree_runner is None:
+        rtree_runner = _get_instance(config, "rtree_runner", n_jobs=n_jobs)
+
     with warnings.catch_warnings():
         # We don't need to show our own warning here
         # TODO remove this once the deprecation has been enforced
@@ -285,25 +312,16 @@ def _get_sfilter_indices(
             # contains is a faster predicate
             # see discussion at https://github.com/geopandas/geopandas/pull/1421
             predicate = "contains"
-            sindex, kwargs = _get_spatial_tree(left)
-            input_geoms = right.geometry if isinstance(right, GeoDataFrame) else right
+            arr1 = right.geometry.values
+            arr2 = left.geometry.values
         else:
             # all other predicates are symmetric
             # keep them the same
-            sindex, kwargs = _get_spatial_tree(right)
-            input_geoms = left.geometry if isinstance(left, GeoDataFrame) else left
+            arr1 = left.geometry.values
+            arr2 = right.geometry.values
 
-    l_idx, r_idx = sindex.query(
-        input_geoms, predicate=predicate, distance=distance, **kwargs
-    )
+    left, right = rtree_runner.run(arr1, arr2, predicate=predicate, distance=distance)
 
     if original_predicate == "within":
-        return np.sort(np.unique(r_idx))
-
-    return np.sort(np.unique(l_idx))
-
-
-def _get_spatial_tree(df):
-    if int(geopandas_version[0]) >= 1:
-        return df.sindex, {"sort": False}
-    return STRtree(df.geometry.values), {}
+        return np.sort(np.unique(right))
+    return np.sort(np.unique(left))

sgis/geopandas_tools/utils.py

@@ -0,0 +1,37 @@
+import numpy as np
+import pandas as pd
+from geopandas import GeoSeries
+from shapely import make_valid
+from shapely import union_all
+
+from .geometry_types import to_single_geom_type
+
+
+def _unary_union_for_notna(geoms, **kwargs):
+    try:
+        return make_valid(union_all(geoms, **kwargs))
+    except TypeError:
+        return make_valid(union_all([geom for geom in geoms.dropna().values], **kwargs))
+
+
+def make_valid_and_keep_geom_type(geoms: np.ndarray, geom_type: str) -> GeoSeries:
+    """Make GeometryCollections into (Multi)Polygons, (Multi)LineStrings or (Multi)Points.
+
+    Because GeometryCollections might appear after dissolving (union_all).
+    And this makes shapely difference/intersection fail.
+
+    Args:
+        geoms: Array of geometries.
+        geom_type: geometry type to be kept.
+    """
+    geoms = GeoSeries(geoms)
+    geoms.index = range(len(geoms))
+    geoms.loc[:] = make_valid(geoms.to_numpy())
+    geoms_with_correct_type = geoms.explode(index_parts=False).pipe(
+        to_single_geom_type, geom_type
+    )
+    only_one = geoms_with_correct_type.groupby(level=0).transform("size") == 1
+    one_hit = geoms_with_correct_type[only_one]
+    many_hits = geoms_with_correct_type[~only_one].groupby(level=0).agg(union_all)
+    geoms_with_wrong_type = geoms.loc[~geoms.index.isin(geoms_with_correct_type.index)]
+    return pd.concat([one_hit, many_hits, geoms_with_wrong_type]).sort_index()

sgis/helpers.py

@@ -198,7 +198,7 @@ def get_all_files(root: str, recursive: bool = True) -> list[str]:
 
 
 def return_two_vals(
-    vals: tuple[str, str] | list[str] | str | int | float
+    vals: tuple[str, str] | list[str] | str | int | float,
 ) -> tuple[str | int | float, str | int | float]:
     """Return a two-length tuple from a str/int/float or list/tuple of length 1 or 2.
 

sgis/io/dapla_functions.py

@@ -128,11 +128,11 @@ def read_geopandas(
         return gpd.GeoDataFrame(
             _read_partitioned_parquet(
                 gcs_path,
-                read_func=pq.read_table,
                 file_system=file_system,
                 mask=mask,
                 filters=filters,
                 child_paths=child_paths,
+                use_threads=use_threads,
                 **kwargs,
             )
         )
@@ -145,7 +145,7 @@ def read_geopandas(
         read_func = gpd.read_file
 
     with file_system.open(gcs_path, mode="rb") as file:
-        return _read_geopandas(
+        return _read_geopandas_single_path(
             file,
             read_func=read_func,
             file_format=file_format,
@@ -163,18 +163,10 @@ def _read_geopandas_from_iterable(
         paths = list(paths.index)
     elif mask is None:
         paths = list(paths)
-    else:
-        if not isinstance(paths, GeoSeries):
-            bounds_series: GeoSeries = get_bounds_series(
-                paths,
-                file_system,
-                use_threads=use_threads,
-                pandas_fallback=pandas_fallback,
-            )
-        else:
-            bounds_series = paths
-        new_bounds_series = sfilter(bounds_series, mask)
-        if not len(new_bounds_series):
+    elif isinstance(paths, GeoSeries):
+        bounds_series = sfilter(paths, mask)
+        if not len(bounds_series):
+            # return GeoDataFrame with correct columns
             if isinstance(kwargs.get("columns"), Iterable):
                 cols = {col: [] for col in kwargs["columns"]}
             else:
@@ -186,29 +178,14 @@ def _read_geopandas_from_iterable(
                         if file_system.isfile(path):
                             raise ArrowInvalid(e, path) from e
             return GeoDataFrame(cols | {"geometry": []})
-        paths = list(new_bounds_series.index)
+        paths = list(bounds_series.index)
 
-    # recursive read with threads
-    threads = (
-        min(len(paths), int(multiprocessing.cpu_count())) or 1 if use_threads else 1
+    results: list[pyarrow.Table] = _read_pyarrow_with_treads(
+        paths, file_system=file_system, mask=mask, use_threads=use_threads, **kwargs
     )
-    with joblib.Parallel(n_jobs=threads, backend="threading") as parallel:
-        dfs: list[GeoDataFrame] = parallel(
-            joblib.delayed(read_geopandas)(
-                x,
-                file_system=file_system,
-                pandas_fallback=pandas_fallback,
-                mask=mask,
-                use_threads=use_threads,
-                **kwargs,
-            )
-            for x in paths
-        )
-
-    if dfs:
-        df = pd.concat(dfs, ignore_index=True)
+    if results:
         try:
-            df = GeoDataFrame(df)
+            return _concat_pyarrow_to_geopandas(results, paths, file_system)
         except Exception as e:
             if not pandas_fallback:
                 print(e)
@@ -219,6 +196,49 @@ def _read_geopandas_from_iterable(
     return df
 
 
+def _read_pyarrow_with_treads(
+    paths: list[str | Path | os.PathLike], file_system, use_threads, mask, **kwargs
+) -> list[pyarrow.Table]:
+    read_partial = functools.partial(
+        _read_pyarrow, mask=mask, file_system=file_system, **kwargs
+    )
+    if not use_threads:
+        return [x for x in map(read_partial, paths) if x is not None]
+    with ThreadPoolExecutor() as executor:
+        return [x for x in executor.map(read_partial, paths) if x is not None]
+
+
+def intersects(file, mask, file_system) -> bool:
+    bbox, _ = _get_bounds_parquet_from_open_file(file, file_system)
+    return shapely.box(*bbox).intersects(to_shapely(mask))
+
+
+def _read_pyarrow(path: str, file_system, mask=None, **kwargs) -> pyarrow.Table | None:
+    try:
+        with file_system.open(path, "rb") as file:
+            if mask is not None and not intersects(file, mask, file_system):
+                return
+
+            # 'get' instead of 'pop' because dict is mutable
+            schema = kwargs.get("schema", pq.read_schema(file))
+            new_kwargs = {
+                key: value for key, value in kwargs.items() if key != "schema"
+            }
+
+            return pq.read_table(file, schema=schema, **new_kwargs)
+    except ArrowInvalid as e:
+        glob_func = _get_glob_func(file_system)
+        if not len(
+            {
+                x
+                for x in glob_func(str(Path(path) / "**"))
+                if not paths_are_equal(path, x)
+            }
+        ):
+            raise e
+        # allow not being able to read empty directories that are hard to delete in gcs
+
+
 def _get_bounds_parquet(
     path: str | Path, file_system: GCSFileSystem, pandas_fallback: bool = False
 ) -> tuple[list[float], dict] | tuple[None, None]:
@@ -662,10 +682,10 @@ def expression_match_path(expression: ds.Expression, path: str) -> bool:
     return bool(len(table))
 
 
-def _read_geopandas(
+def _read_geopandas_single_path(
     file,
-    read_func: Callable = gpd.read_parquet,
-    file_format: str = "parquet",
+    read_func: Callable,
+    file_format: str,
     **kwargs,
 ):
     try:
@@ -681,32 +701,29 @@ def _read_geopandas(
         raise e.__class__(f"{e.__class__.__name__}: {e} for {file}.") from e
 
 
-def _read_pandas(gcs_path: str, **kwargs):
+def _read_pandas(gcs_path: str, use_threads: bool = True, **kwargs):
     file_system = _get_file_system(None, kwargs)
 
     if not isinstance(gcs_path, (str | Path | os.PathLike)):
-        # recursive read with threads
-        threads = (
-            min(len(gcs_path), int(multiprocessing.cpu_count())) or 1
-            if kwargs.get("use_threads")
-            else 1
+        results: list[pyarrow.Table] = _read_pyarrow_with_treads(
+            gcs_path,
+            file_system=file_system,
+            mask=None,
+            use_threads=use_threads,
+            **kwargs,
         )
-        with joblib.Parallel(n_jobs=threads, backend="threading") as parallel:
-            return pd.concat(
-                parallel(
-                    joblib.delayed(_read_pandas)(x, file_system=file_system, **kwargs)
-                    for x in gcs_path
-                )
-            )
+        results = pyarrow.concat_tables(results, promote_options="permissive")
+        return results.to_pandas()
 
     child_paths = get_child_paths(gcs_path, file_system)
     if child_paths:
         return _read_partitioned_parquet(
             gcs_path,
-            read_func=pd.read_parquet,
             file_system=file_system,
             mask=None,
             child_paths=child_paths,
+            use_threads=use_threads,
+            to_geopandas=False,
             **kwargs,
         )
 
@@ -716,11 +733,12 @@ def _read_pandas(gcs_path: str, **kwargs):
 
 def _read_partitioned_parquet(
     path: str,
-    read_func: Callable,
     filters=None,
     file_system=None,
     mask=None,
     child_paths: list[str] | None = None,
+    use_threads: bool = True,
+    to_geopandas: bool = True,
     **kwargs,
 ):
     file_system = _get_file_system(file_system, kwargs)
@@ -731,62 +749,22 @@ def _read_partitioned_parquet(
 
     filters = _filters_to_expression(filters)
 
-    def intersects(file, mask) -> bool:
-        bbox, _ = _get_bounds_parquet_from_open_file(file, file_system)
-        return shapely.box(*bbox).intersects(to_shapely(mask))
-
-    def read(child_path: str) -> pyarrow.Table | None:
-        try:
-            with file_system.open(child_path, "rb") as file:
-                if mask is not None and not intersects(file, mask):
-                    return
-
-                # 'get' instead of 'pop' because dict is mutable
-                schema = kwargs.get("schema", pq.read_schema(file))
-                new_kwargs = {
-                    key: value for key, value in kwargs.items() if key != "schema"
-                }
-
-                return read_func(file, schema=schema, filters=filters, **new_kwargs)
-        except ArrowInvalid as e:
-            if not len(
-                {
-                    x
-                    for x in glob_func(str(Path(child_path) / "**"))
-                    if not paths_are_equal(child_path, x)
-                }
-            ):
-                raise e
-            # allow not being able to read hard-to-delete empty directories
+    results: list[pyarrow.Table] = _read_pyarrow_with_treads(
+        (
+            path
+            for path in child_paths
+            if filters is None or expression_match_path(filters, path)
+        ),
+        file_system=file_system,
+        mask=mask,
+        use_threads=use_threads,
+        **kwargs,
+    )
 
-    with ThreadPoolExecutor() as executor:
-        results = [
-            df
-            for df in (
-                executor.map(
-                    read,
-                    (
-                        path
-                        for path in child_paths
-                        if filters is None or expression_match_path(filters, path)
-                    ),
-                )
-            )
-            if df is not None
-        ]
-
-    if results:
-        if all(isinstance(x, DataFrame) for x in results):
-            return pd.concat(results)
-        else:
-            geo_metadata = _get_geo_metadata(next(iter(child_paths)), file_system)
-            return _arrow_to_geopandas(
-                pyarrow.concat_tables(
-                    results,
-                    promote_options="permissive",
-                ),
-                geo_metadata,
-            )
+    if results and to_geopandas:
+        return _concat_pyarrow_to_geopandas(results, child_paths, file_system)
+    elif results:
+        return pyarrow.concat_tables(results, promote_options="permissive").to_pandas()
 
     # add columns to empty DataFrame
     first_path = next(iter(child_paths + [path]))
@@ -796,6 +774,17 @@ def _read_partitioned_parquet(
     return df
 
 
+def _concat_pyarrow_to_geopandas(
+    results: list[pyarrow.Table], paths: list[str], file_system: Any
+):
+    results = pyarrow.concat_tables(
+        results,
+        promote_options="permissive",
+    )
+    geo_metadata = _get_geo_metadata(next(iter(paths)), file_system)
+    return _arrow_to_geopandas(results, geo_metadata)
+
+
 def paths_are_equal(path1: Path | str, path2: Path | str) -> bool:
     return Path(path1).parts == Path(path2).parts
 

sgis/maps/map.py

@@ -307,7 +307,9 @@ class Map:
         notna = array[array.notna()]
         isna = array[array.isna()]
 
-        unique_multiplied = (notna * self._multiplier).astype(np.int64)
+        unique_multiplied = (notna.astype(np.float64) * self._multiplier).astype(
+            np.int64
+        )
 
         return pd.concat([unique_multiplied, isna]).sort_index()
 

sgis/parallel/parallel.py

@@ -75,13 +75,15 @@ def parallel_overlay(
     Returns:
         A GeoDataFrame containing the result of the overlay operation.
     """
+    if how != "intersection":
+        raise ValueError("parallel_overlay only supports how='intersection'.")
     return pd.concat(
         chunkwise(
             _clean_overlay_with_print,
             df1,
             kwargs={
                 "df2": df2,
-                # "to_print": to_print,
+                "to_print": to_print,
                 "how": how,
             }
             | kwargs,
@@ -672,7 +674,7 @@ class Parallel:
     def chunkwise(
         self,
         func: Callable,
-        iterable: Collection[Iterable[Any]],
+        *iterables: Collection[Iterable[Any]],
         args: tuple | None = None,
         kwargs: dict | None = None,
         max_rows_per_chunk: int | None = None,
@@ -682,8 +684,8 @@ class Parallel:
         Args:
             func: Function to run chunkwise. It should take
                 (a chunk of) the iterable as first argument.
-            iterable: Iterable to split in chunks and passed
-                as first argument to 'func'.
+            iterables: Iterable(s) to split in chunks and passed
+                as first argument(s) to 'func'. Iterables must have same length.
             args: Positional arguments in 'func' after the DataFrame.
             kwargs: Additional keyword arguments in 'func'.
             max_rows_per_chunk: Alternatively decide number of chunks
@@ -691,7 +693,7 @@ class Parallel:
         """
         return chunkwise(
             func,
-            iterable,
+            *iterables,
             args=args,
             kwargs=kwargs,
             processes=self.processes,
@@ -1067,7 +1069,7 @@ def _fix_missing_muni_numbers(
 
 def chunkwise(
     func: Callable,
-    iterable: Collection[Iterable[Any]],
+    *iterables: Collection[Iterable[Any]],
     args: tuple | None = None,
     kwargs: dict | None = None,
     processes: int = 1,
@@ -1082,7 +1084,7 @@ def chunkwise(
     Args:
         func: The function to apply to each chunk. This function must accept a DataFrame as
             its first argument and return a DataFrame.
-        iterable: Iterable to be chunked and processed.
+        iterables: Iterable(s) to be chunked and processed. Must have same length.
         args: Additional positional arguments to pass to 'func'.
         kwargs: Keyword arguments to pass to 'func'.
         processes: The number of parallel jobs to run. Defaults to 1 (no parallel execution).
@@ -1096,30 +1098,36 @@ def chunkwise(
     args = args or ()
     kwargs = kwargs or {}
 
+    if len({len(x) for x in iterables}) not in [0, 1]:
+        raise ValueError(
+            f"iterables must have same length. Got {', '.join([len(x) for x in iterables])}"
+        )
+
     if max_rows_per_chunk is None:
         n_chunks: int = processes
     else:
-        n_chunks: int = len(iterable) // max_rows_per_chunk
-
+        n_chunks: int = len(next(iter(iterables))) // max_rows_per_chunk
     if n_chunks <= 1:
-        return [func(iterable, *args, **kwargs)]
+        return [func(*iterables, *args, **kwargs)]
 
-    chunks = np.array_split(np.arange(len(iterable)), n_chunks)
+    chunks = np.array_split(np.arange(len(next(iter(iterables)))), n_chunks)
 
-    if hasattr(iterable, "iloc"):
-        iterable_chunked: list[pd.DataFrame | pd.Series] = [
-            iterable.iloc[chunk] for chunk in chunks
-        ]
-    elif is_array_like(iterable):
-        iterable_chunked: list[np.ndarray] = [iterable[chunk] for chunk in chunks]
-    else:
-        to_type: type = iterable.__class__
-        iterable_chunked: list[Iterable] = [
-            to_type(chunk) for chunk in np.array_split(list(iterable), n_chunks)
-        ]
-    return Parallel(processes, backend=backend).map(
+    def get_chunk(iterable, chunk):
+        if hasattr(iterable, "iloc"):
+            return iterable.iloc[chunk]
+        elif is_array_like(iterable):
+            return iterable[chunk]
+        else:
+            to_type: type = iterable.__class__
+            return to_type([x for i, x in enumerate(iterable) if i in chunk])
+
+    iterables_chunked: list[list[Iterable[Any]]] = [
+        [get_chunk(iterable, chunk) for iterable in iterables] for chunk in chunks
+    ]
+
+    return Parallel(processes, backend=backend).starmap(
         func,
-        iterable_chunked,
+        iterables_chunked,
         args=args,
         kwargs=kwargs,
     )