rq_geo_toolkit 2025.4.0__tar.gz

rq_geo_toolkit-2025.4.0/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.1
+Name: rq_geo_toolkit
+Version: 2025.4.0
+Summary: Collection of geo related functions for reuse in other libraries
+Author: Kamil Raczycki
+License: MIT
+Requires-Python: >=3.9
+Requires-Dist: geopandas>=1.0
+Requires-Dist: shapely>=2.0.6
+Requires-Dist: pyarrow>=16.0.0
+Requires-Dist: rich>=12.0.0
+Requires-Dist: geoarrow-rust-core>=0.3.0
+Requires-Dist: pooch>=1.6.0
+Requires-Dist: geopy>=2.0.0
+Requires-Dist: numpy>=1.26.0
+Requires-Dist: duckdb>=1.1.2
+Requires-Dist: psutil>=5.6.2
+Requires-Dist: polars>=1.9
+Requires-Dist: packaging>=17.0
+Description-Content-Type: text/markdown
+
+# RQ Geo Toolkit

rq_geo_toolkit-2025.4.0/README.md

@@ -0,0 +1 @@
+# RQ Geo Toolkit

rq_geo_toolkit-2025.4.0/pyproject.toml

@@ -0,0 +1,78 @@
+[project]
+name = "rq_geo_toolkit"
+version = "2025.4.0"
+description = "Collection of geo related functions for reuse in other libraries"
+authors = [
+    { name = "Kamil Raczycki", email = "" },
+]
+dependencies = [
+    "geopandas>=1.0",
+    "shapely>=2.0.6",
+    "pyarrow>=16.0.0",
+    "rich>=12.0.0",
+    "geoarrow-rust-core>=0.3.0",
+    "pooch>=1.6.0",
+    "geopy>=2.0.0",
+    "numpy>=1.26.0",
+    "duckdb>=1.1.2",
+    "psutil>=5.6.2",
+    "polars>=1.9",
+    "packaging>=17.0",
+]
+requires-python = ">=3.9"
+readme = "README.md"
+
+[project.license]
+text = "MIT"
+
+[build-system]
+requires = [
+    "pdm-backend",
+]
+build-backend = "pdm.backend"
+
+[tool.pdm]
+distribution = true
+
+[tool.pdm.dev-dependencies]
+dev = [
+    "bumpver",
+    "types-requests",
+    "setuptools>=45.0.0",
+]
+lint = [
+    "pre-commit>=4",
+    "mypy>=1",
+    "docformatter[tomli]",
+    "ruff>=0.1.0",
+]
+test = [
+    "pytest>=7.0.0",
+    "tox-pdm>=0.7.2",
+    "pytest-mock>=3.3.0",
+    "requests-mock>=1.12.1",
+    "pytest-check>=2.3.1",
+    "pytest-parametrization>=2022.2.1",
+    "pytest-doctestplus>=1.2.1",
+    "osmnx>=1.3.0",
+]
+
+[tool.pdm.scripts]
+post_install = "pre-commit install"
+
+[tool.bumpver]
+current_version = "2025.4.0"
+version_pattern = "YYYY.MM.INC0"
+commit_message = "chore(CI/CD): bump version {old_version} -> {new_version}"
+commit = true
+tag = false
+push = false
+
+[tool.bumpver.file_patterns]
+"pyproject.toml" = [
+    "^current_version = \"{version}\"$",
+    "^version = \"{version}\"$",
+]
+"rq_geo_toolkit/__init__.py" = [
+    "^__version__ = \"{version}\"$",
+]

rq_geo_toolkit-2025.4.0/rq_geo_toolkit/__init__.py

@@ -0,0 +1 @@
+__version__ = "2025.4.0"

rq_geo_toolkit-2025.4.0/rq_geo_toolkit/_exceptions.py

@@ -0,0 +1 @@
+class QueryNotGeocodedError(ValueError): ...

rq_geo_toolkit-2025.4.0/rq_geo_toolkit/_geopandas_api_version.py

@@ -0,0 +1,4 @@
+import geopandas as gpd
+from packaging import version
+
+GEOPANDAS_NEW_API = version.parse(gpd.__version__) >= version.parse("1.0.0")

rq_geo_toolkit-2025.4.0/rq_geo_toolkit/constants.py

@@ -0,0 +1,5 @@
+GEOMETRY_COLUMN = "geometry"
+
+PARQUET_ROW_GROUP_SIZE = 100_000
+PARQUET_COMPRESSION = "zstd"
+PARQUET_COMPRESSION_LEVEL = 3

rq_geo_toolkit-2025.4.0/rq_geo_toolkit/duckdb.py

@@ -0,0 +1,29 @@
+"""Helper functions for DuckDB."""
+
+from pathlib import Path
+from typing import Union
+
+import duckdb
+
+
+def sql_escape(value: str) -> str:
+    """Escape value for SQL query."""
+    return value.replace("'", "''")
+
+
+def set_up_duckdb_connection(
+    tmp_dir_path: Union[str, Path], preserve_insertion_order: bool = False
+) -> "duckdb.DuckDBPyConnection":
+    """Create DuckDB connection in a given directory."""
+    local_db_file = "db.duckdb"
+    connection = duckdb.connect(
+        database=str(Path(tmp_dir_path) / local_db_file),
+        config=dict(preserve_insertion_order=preserve_insertion_order),
+    )
+    connection.sql("SET enable_progress_bar = false;")
+    connection.sql("SET enable_progress_bar_print = false;")
+
+    connection.install_extension("spatial")
+    connection.load_extension("spatial")
+
+    return connection

rq_geo_toolkit-2025.4.0/rq_geo_toolkit/geocode.py

@@ -0,0 +1,70 @@
+"""Geocoding module for getting a geometry from query using Nominatim."""
+
+import hashlib
+import json
+from pathlib import Path
+from typing import Any, Optional, Union, cast, overload
+
+from geopy.geocoders.nominatim import Nominatim
+from geopy.location import Location
+from shapely.geometry import shape
+from shapely.geometry.base import BaseGeometry
+from shapely.ops import unary_union
+
+from rq_geo_toolkit._exceptions import QueryNotGeocodedError
+
+USER_AGENT = "RQ Geo Toolkit Python package (https://github.com/kraina-ai/rq_geo_toolkit)"
+
+@overload
+def geocode_to_geometry(query: str) -> BaseGeometry: ...
+
+
+@overload
+def geocode_to_geometry(query: list[str]) -> BaseGeometry: ...
+
+
+def geocode_to_geometry(query: Union[str, list[str]]) -> BaseGeometry:
+    """Geocode a query to a (Multi)Polygon geometry using Nominatim."""
+    if not isinstance(query, str):
+        return unary_union([geocode_to_geometry(sub_query) for sub_query in query])
+
+    h = hashlib.new("sha256")
+    h.update(query.encode())
+    query_hash = h.hexdigest()
+    query_file_path = Path("cache").resolve() / f"{query_hash}.json"
+
+    if not query_file_path.exists():
+        query_results = Nominatim(
+            user_agent=USER_AGENT
+        ).geocode(query, geometry="geojson", exactly_one=False)
+
+        if not query_results:
+            raise QueryNotGeocodedError(f"Zero results from Nominatim for query '{query}'.")
+
+        polygon_result = _get_first_polygon(query_results)
+
+        if not polygon_result:
+            raise QueryNotGeocodedError(f"No polygon found for query '{query}'.")
+
+        query_file_path.parent.mkdir(parents=True, exist_ok=True)
+        query_file_path.write_text(json.dumps(polygon_result))
+    else:
+        polygon_result = json.loads(query_file_path.read_text())
+
+    return unary_union(shape(polygon_result))
+
+
+def _get_first_polygon(results: list[Location]) -> Optional[dict[str, Any]]:
+    """Choose first result of geometry type (Multi)Polygon from list of
+    results.
+
+    Inspired by OSMnx implementation.
+    """
+    polygon_types = {"Polygon", "MultiPolygon"}
+
+    for result in results:
+        geojson_dict = cast(dict[str, Any], result.raw["geojson"])
+        if geojson_dict["type"] in polygon_types:
+            return geojson_dict
+
+    return None

rq_geo_toolkit-2025.4.0/rq_geo_toolkit/geoparquet_compression.py

@@ -0,0 +1,187 @@
+"""Module for sorting GeoParquet files."""
+
+import multiprocessing
+import tempfile
+from collections.abc import Callable
+from functools import partial
+from math import ceil
+from pathlib import Path
+from time import sleep
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import duckdb
+import psutil
+import pyarrow.parquet as pq
+from rich import print as rprint
+
+from rq_geo_toolkit.constants import PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL, PARQUET_ROW_GROUP_SIZE
+from rq_geo_toolkit.duckdb import set_up_duckdb_connection
+
+if TYPE_CHECKING:  # pragma: no cover
+    from rq_geo_toolkit.rich_utils import VERBOSITY_MODE
+
+MEMORY_1GB = 1024**3
+
+
+def compress_parquet_with_duckdb(
+    input_file_path: Path,
+    output_file_path: Path,
+    working_directory: Union[str, Path] = "files",
+    parquet_metadata: Optional[pq.FileMetaData] = None,
+    verbosity_mode: "VERBOSITY_MODE" = "transient",
+) -> Path:
+    """Compresses a GeoParquet file while keeping its metadata.
+
+    Args:
+        input_file_path (Path): Input GeoParquet file path.
+        output_file_path (Path): Output GeoParquet file path.
+        working_directory (Union[str, Path], optional): Directory where to save
+            the downloaded `*.parquet` files. Defaults to "files".
+        parquet_metadata (Optional[pq.FileMetaData], optional): GeoParquet file metadata used to
+            copy. If not provided, will load the metadata from the input file. Defaults to None.
+        verbosity_mode (Literal["silent", "transient", "verbose"], optional): Set progress
+            verbosity mode. Can be one of: silent, transient and verbose. Silent disables
+            output completely. Transient tracks progress, but removes output after finished.
+            Verbose leaves all progress outputs in the stdout. Defaults to "transient".
+    """
+    assert input_file_path.resolve().as_posix() != output_file_path.resolve().as_posix()
+
+    Path(working_directory).mkdir(parents=True, exist_ok=True)
+
+    if pq.read_metadata(input_file_path).num_rows == 0:
+        return input_file_path.rename(output_file_path)
+
+    with tempfile.TemporaryDirectory(dir=Path(working_directory).resolve()) as tmp_dir_name:
+        tmp_dir_path = Path(tmp_dir_name)
+
+        original_metadata_string = _parquet_schema_metadata_to_duckdb_kv_metadata(
+            parquet_metadata or pq.read_metadata(input_file_path)
+        )
+
+        _run_query_with_memory_limit(
+            tmp_dir_path=tmp_dir_path,
+            verbosity_mode=verbosity_mode,
+            current_memory_gb_limit=None,
+            current_threads_limit=None,
+            function=_compress_with_memory_limit,
+            args=(input_file_path, output_file_path, original_metadata_string),
+        )
+
+    return output_file_path
+
+
+def _compress_with_memory_limit(
+    input_file_path: Union[list[Path], Path],
+    output_file_path: Path,
+    original_metadata_string: str,
+    current_memory_gb_limit: float,
+    current_threads_limit: int,
+    tmp_dir_path: Path,
+) -> None:
+    connection = set_up_duckdb_connection(tmp_dir_path, preserve_insertion_order=True)
+
+    connection.execute("SET enable_geoparquet_conversion = false;")
+    connection.execute(f"SET memory_limit = '{current_memory_gb_limit}GB';")
+    connection.execute(f"SET threads = {current_threads_limit};")
+
+    if isinstance(input_file_path, Path):
+        sql_input_str = f"'{input_file_path}'"
+    else:
+        mapped_paths = ", ".join(f"'{path}'" for path in input_file_path)
+        sql_input_str = f"[{mapped_paths}]"
+
+    connection.execute(
+        f"""
+        COPY (
+            SELECT original_data.*
+            FROM read_parquet({sql_input_str}, hive_partitioning=false) original_data
+        ) TO '{output_file_path}' (
+            FORMAT parquet,
+            COMPRESSION {PARQUET_COMPRESSION},
+            COMPRESSION_LEVEL {PARQUET_COMPRESSION_LEVEL},
+            ROW_GROUP_SIZE {PARQUET_ROW_GROUP_SIZE},
+            KV_METADATA {original_metadata_string}
+        );
+        """
+    )
+
+    connection.close()
+
+
+def _run_query_with_memory_limit(
+    tmp_dir_path: Path,
+    verbosity_mode: "VERBOSITY_MODE",
+    current_memory_gb_limit: Optional[float],
+    current_threads_limit: Optional[int],
+    function: Callable[..., None],
+    args: Any,
+) -> tuple[float, int]:
+    current_memory_gb_limit = current_memory_gb_limit or ceil(
+        psutil.virtual_memory().total / MEMORY_1GB
+    )
+    current_threads_limit = current_threads_limit or multiprocessing.cpu_count()
+
+    while current_memory_gb_limit > 0:
+        try:
+            with (
+                tempfile.TemporaryDirectory(dir=Path(tmp_dir_path).resolve()) as tmp_dir_name,
+                multiprocessing.get_context("spawn").Pool() as pool,
+            ):
+                nested_tmp_dir_path = Path(tmp_dir_name)
+                r = pool.apply_async(
+                    func=partial(
+                        function,
+                        current_memory_gb_limit=current_memory_gb_limit,
+                        current_threads_limit=current_threads_limit,
+                        tmp_dir_path=nested_tmp_dir_path,
+                    ),
+                    args=args,
+                )
+                actual_memory = psutil.virtual_memory()
+                percentage_threshold = 95
+                if (actual_memory.total * 0.05) > MEMORY_1GB:
+                    percentage_threshold = (
+                        100 * (actual_memory.total - MEMORY_1GB) / actual_memory.total
+                    )
+                while not r.ready():
+                    actual_memory = psutil.virtual_memory()
+                    if actual_memory.percent > percentage_threshold:
+                        raise MemoryError()
+
+                    sleep(0.5)
+                r.get()
+            return current_memory_gb_limit, current_threads_limit
+        except (duckdb.OutOfMemoryException, MemoryError) as ex:
+            if current_memory_gb_limit < 1:
+                raise RuntimeError(
+                    "Not enough memory to run the ordering query. Please rerun without sorting."
+                ) from ex
+
+            if current_memory_gb_limit == 1:
+                current_memory_gb_limit /= 2
+            else:
+                current_memory_gb_limit = ceil(current_memory_gb_limit / 2)
+
+            current_threads_limit = ceil(current_threads_limit / 2)
+
+            if not verbosity_mode == "silent":
+                rprint(
+                    f"Encountered {ex.__class__.__name__} during operation."
+                    " Retrying with lower number of resources"
+                    f" ({current_memory_gb_limit:.2f}GB, {current_threads_limit} threads)."
+                )
+
+    raise RuntimeError("Not enough memory to run the query. Please rerun without sorting.")
+
+
+def _parquet_schema_metadata_to_duckdb_kv_metadata(parquet_file_metadata: pq.FileMetaData) -> str:
+    def escape_single_quotes(s: str) -> str:
+        return s.replace("'", "''")
+
+    kv_pairs = []
+    for key, value in parquet_file_metadata.metadata.items():
+        escaped_key = escape_single_quotes(key.decode())
+        escaped_value = escape_single_quotes(value.decode())
+        kv_pairs.append(f"'{escaped_key}': '{escaped_value}'")
+
+    return "{ " + ", ".join(kv_pairs) + " }"