test_geohash_partition.py

#   -------------------------------------------------------------
#   Copyright (c) Henry Spatial Analysis. All rights reserved.
#   Licensed under the MIT License. See LICENSE in project root for information.
#   -------------------------------------------------------------

"""
Unit tests for openpois.io.geohash_partition.

No network or real-filesystem I/O beyond tmp_path (pytest fixture).
shutil.rmtree is mocked in tests that verify overwrite behaviour so the
actual on-disk side-effects remain predictable.
"""
from __future__ import annotations

import warnings

import geopandas as gpd
import pandas as pd
import pygeohash
import pytest
from shapely.geometry import MultiPolygon, Point, Polygon

from openpois.io.geohash_partition import (
    add_geohash_column,
    add_geohash_columns,
    compute_primary_osm_tag,
    write_label_partitioned_dataset,
    write_partitioned_dataset,
)


# ---------------------------------------------------------------------------
# Fixtures / helpers
# ---------------------------------------------------------------------------


def _point_gdf(*lonlats: tuple[float, float]) -> gpd.GeoDataFrame:
    """Return a GeoDataFrame with Point geometries at the given lon/lat pairs."""
    return gpd.GeoDataFrame(
        {"name": [f"p{i}" for i in range(len(lonlats))]},
        geometry=[Point(lon, lat) for lon, lat in lonlats],
        crs="EPSG:4326",
    )


def _poly_gdf() -> gpd.GeoDataFrame:
    """Return a GeoDataFrame with one small square Polygon."""
    poly = Polygon(
        [(-122.3, 47.6), (-122.2, 47.6), (-122.2, 47.7), (-122.3, 47.7), (-122.3, 47.6)]
    )
    return gpd.GeoDataFrame({"name": ["block"]}, geometry=[poly], crs="EPSG:4326")


# ---------------------------------------------------------------------------
# add_geohash_columns
# ---------------------------------------------------------------------------


class TestAddGeohashColumns:
    def test_adds_expected_columns(self):
        """Both geohash_prefix and geohash_sort columns must be present after call."""
        gdf = _point_gdf((-122.3, 47.6))
        result = add_geohash_columns(gdf, precision_partition=4, precision_sort=6)

        assert "geohash_prefix" in result.columns
        assert "geohash_sort" in result.columns

    def test_returns_geodataframe(self):
        """Function must return a GeoDataFrame (now returns a copy, not same object)."""
        gdf = _point_gdf((-122.3, 47.6))
        result = add_geohash_columns(gdf, precision_partition=4, precision_sort=6)
        assert isinstance(result, gpd.GeoDataFrame)

    def test_prefix_length_matches_precision_partition(self):
        """geohash_prefix values must have exactly precision_partition characters."""
        gdf = _point_gdf((-122.3, 47.6), (-77.0, 38.9))
        result = add_geohash_columns(gdf, precision_partition=3, precision_sort=6)
        assert all(len(v) == 3 for v in result["geohash_prefix"])

    def test_sort_length_matches_precision_sort(self):
        """geohash_sort values must have exactly precision_sort characters."""
        gdf = _point_gdf((-122.3, 47.6), (-77.0, 38.9))
        result = add_geohash_columns(gdf, precision_partition=3, precision_sort=7)
        assert all(len(v) == 7 for v in result["geohash_sort"])

    def test_sort_starts_with_prefix(self):
        """The geohash_sort value must start with the geohash_prefix value."""
        result = add_geohash_columns(
            _point_gdf((-122.3, 47.6)), precision_partition=4, precision_sort=6
        )
        prefix = result["geohash_prefix"].iloc[0]
        sort_ = result["geohash_sort"].iloc[0]
        assert sort_.startswith(prefix)

    def test_values_match_pygeohash_directly(self):
        """Encoded values must agree with pygeohash.encode called directly."""
        lon, lat = -122.3, 47.6
        result = add_geohash_columns(
            _point_gdf((lon, lat)), precision_partition=4, precision_sort=6
        )
        expected_prefix = pygeohash.encode(lat, lon, precision=4)
        expected_sort = pygeohash.encode(lat, lon, precision=6)
        assert result["geohash_prefix"].iloc[0] == expected_prefix
        assert result["geohash_sort"].iloc[0] == expected_sort

    def test_multiple_rows_get_independent_hashes(self):
        """Points in different geohash cells must receive different prefix values."""
        # Seattle and New York — guaranteed different 2-char geohash prefix
        result = add_geohash_columns(
            _point_gdf((-122.3, 47.6), (-74.0, 40.7)),
            precision_partition = 2,
            precision_sort = 6,
        )
        prefixes = result["geohash_prefix"].tolist()
        assert prefixes[0] != prefixes[1]

    def test_polygon_geometry_uses_centroid(self):
        """A Polygon row should produce the geohash of its centroid, not a corner."""
        gdf = _poly_gdf()
        result = add_geohash_columns(gdf, precision_partition=6, precision_sort=8)

        # Compute expected centroid coords (suppress CRS warning deliberately)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", UserWarning)
            cx = gdf.geometry.centroid.iloc[0].x
            cy = gdf.geometry.centroid.iloc[0].y
        expected = pygeohash.encode(cy, cx, precision=6)
        assert result["geohash_prefix"].iloc[0] == expected

    def test_multipolygon_geometry_handled(self):
        """A MultiPolygon row must produce a valid geohash string."""
        p1 = Polygon([(-122.3, 47.6), (-122.2, 47.6), (-122.2, 47.7), (-122.3, 47.6)])
        p2 = Polygon([(-122.1, 47.5), (-122.0, 47.5), (-122.0, 47.6), (-122.1, 47.5)])
        mp = MultiPolygon([p1, p2])
        gdf = gpd.GeoDataFrame({"name": ["multi"]}, geometry=[mp], crs="EPSG:4326")
        result = add_geohash_columns(gdf, precision_partition=4, precision_sort=6)

        prefix = result["geohash_prefix"].iloc[0]
        assert isinstance(prefix, str) and len(prefix) == 4

    def test_suppresses_geographic_crs_warning(self):
        """No UserWarning about geographic CRS should escape the function."""
        gdf = _point_gdf((-122.3, 47.6))
        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            add_geohash_columns(gdf, precision_partition=4, precision_sort=6)

        crs_warnings = [
            w for w in caught
            if issubclass(w.category, UserWarning)
            and "geographic CRS" in str(w.message)
        ]
        assert crs_warnings == [], (
            "geographic CRS warning leaked out of add_geohash_columns"
        )

    def test_empty_geodataframe_returns_empty_with_columns(self):
        """An empty GeoDataFrame should gain both columns with no rows."""
        gdf = gpd.GeoDataFrame({"name": []}, geometry=[], crs="EPSG:4326")
        result = add_geohash_columns(gdf, precision_partition=4, precision_sort=6)

        assert len(result) == 0
        assert "geohash_prefix" in result.columns
        assert "geohash_sort" in result.columns

    def test_precision_one_gives_single_char_prefix(self):
        """Edge case: precision=1 should produce 1-character geohash strings."""
        result = add_geohash_columns(
            _point_gdf((-122.3, 47.6)), precision_partition=1, precision_sort=1
        )
        assert len(result["geohash_prefix"].iloc[0]) == 1
        assert len(result["geohash_sort"].iloc[0]) == 1


# ---------------------------------------------------------------------------
# write_partitioned_dataset
# ---------------------------------------------------------------------------


class TestWritePartitionedDataset:
    def _gdf_with_hashes(self, lonlats: list[tuple[float, float]]) -> gpd.GeoDataFrame:
        """Return a GeoDataFrame with geohash columns already populated."""
        gdf = _point_gdf(*lonlats)
        return add_geohash_columns(gdf, precision_partition=4, precision_sort=6)

    # --- directory layout ---

    def test_creates_hive_partition_directories(self, tmp_path):
        """One geohash_prefix=<value> subdirectory must be created per unique prefix."""
        gdf = self._gdf_with_hashes([(-122.3, 47.6)])
        expected_prefix = gdf["geohash_prefix"].iloc[0]

        write_partitioned_dataset(gdf, tmp_path / "out")

        partition_dir = tmp_path / "out" / f"geohash_prefix={expected_prefix}"
        assert partition_dir.is_dir()

    def test_creates_part_file_in_each_partition(self, tmp_path):
        """Each partition directory must contain exactly one part-0.parquet file."""
        gdf = self._gdf_with_hashes([(-122.3, 47.6)])
        write_partitioned_dataset(gdf, tmp_path / "out")

        parts = list((tmp_path / "out").rglob("*.parquet"))
        assert len(parts) == 1
        assert parts[0].name == "part-0.parquet"

    def test_two_distinct_prefixes_produce_two_partitions(self, tmp_path):
        """Points in different geohash cells must each get their own directory."""
        # Seattle and New York are well-separated at precision=2
        gdf = self._gdf_with_hashes([(-122.3, 47.6), (-74.0, 40.7)])
        # Force precision=2 so we guarantee distinct prefixes
        gdf = _point_gdf((-122.3, 47.6), (-74.0, 40.7))
        gdf = add_geohash_columns(gdf, precision_partition=2, precision_sort=4)

        write_partitioned_dataset(gdf, tmp_path / "out")

        partitions = [d for d in (tmp_path / "out").iterdir() if d.is_dir()]
        assert len(partitions) == 2

    def test_same_prefix_points_land_in_single_partition(self, tmp_path):
        """Two points with the same prefix must be co-located in one parquet file."""
        # Two points very close together → same 4-char geohash prefix
        gdf = _point_gdf((-122.300, 47.600), (-122.301, 47.601))
        gdf = add_geohash_columns(gdf, precision_partition=4, precision_sort=6)
        assert gdf["geohash_prefix"].nunique() == 1  # sanity-check fixture

        write_partitioned_dataset(gdf, tmp_path / "out")

        parts = list((tmp_path / "out").rglob("*.parquet"))
        assert len(parts) == 1

    # --- column handling ---

    def test_partition_column_dropped_from_parquet_files(self, tmp_path):
        """geohash_prefix must not appear as a column inside the parquet files."""
        import pyarrow.parquet as pq

        gdf = self._gdf_with_hashes([(-122.3, 47.6)])
        write_partitioned_dataset(gdf, tmp_path / "out")

        part_file = next((tmp_path / "out").rglob("*.parquet"))
        schema = pq.read_schema(part_file)
        assert "geohash_prefix" not in schema.names

    def test_sort_column_dropped_from_parquet_files(self, tmp_path):
        """geohash_sort must not appear as a column inside the parquet files."""
        import pyarrow.parquet as pq

        gdf = self._gdf_with_hashes([(-122.3, 47.6)])
        write_partitioned_dataset(gdf, tmp_path / "out")

        part_file = next((tmp_path / "out").rglob("*.parquet"))
        schema = pq.read_schema(part_file)
        assert "geohash_sort" not in schema.names

    def test_other_columns_preserved_in_parquet(self, tmp_path):
        """User columns (here 'name') must survive in the written parquet files."""
        import pyarrow.parquet as pq

        gdf = self._gdf_with_hashes([(-122.3, 47.6)])
        write_partitioned_dataset(gdf, tmp_path / "out")

        part_file = next((tmp_path / "out").rglob("*.parquet"))
        schema = pq.read_schema(part_file)
        assert "name" in schema.names

    # --- overwrite behaviour ---

    def test_raises_file_exists_error_when_dir_exists_and_no_overwrite(self, tmp_path):
        """Should raise FileExistsError when output exists and overwrite=False."""
        out = tmp_path / "out"
        out.mkdir()

        gdf = self._gdf_with_hashes([(-122.3, 47.6)])
        with pytest.raises(FileExistsError, match="overwrite=True"):
            write_partitioned_dataset(gdf, out, overwrite=False)

    def test_overwrites_existing_directory_when_flag_set(self, tmp_path):
        """Second call with overwrite=True must replace the first output."""
        gdf = self._gdf_with_hashes([(-122.3, 47.6)])
        out = tmp_path / "out"

        write_partitioned_dataset(gdf, out, overwrite=True)
        # Write a sentinel file inside the old run
        sentinel = out / "sentinel.txt"
        sentinel.write_text("old")

        write_partitioned_dataset(gdf, out, overwrite=True)
        assert not sentinel.exists(), "Old output was not removed by overwrite"

    def test_no_existing_directory_succeeds_without_overwrite(self, tmp_path):
        """Should succeed when the output directory does not yet exist."""
        gdf = self._gdf_with_hashes([(-122.3, 47.6)])
        out = tmp_path / "brand_new"
        assert not out.exists()

        write_partitioned_dataset(gdf, out, overwrite=False)
        assert out.is_dir()

    # --- path coercion ---

    def test_accepts_string_path(self, tmp_path):
        """output_dir supplied as a plain string should work without error."""
        gdf = self._gdf_with_hashes([(-122.3, 47.6)])
        write_partitioned_dataset(gdf, str(tmp_path / "out"))
        assert (tmp_path / "out").is_dir()

    # --- row ordering ---

    def test_rows_within_partition_sorted_by_geohash_sort(self, tmp_path):
        """Rows in each partition must be ordered by geohash_sort ascending."""
        import pyarrow.parquet as pq

        # Three points that share a 2-char prefix but differ at finer precision.
        # Use precision_partition=2 so they all land in one file.
        gdf = _point_gdf((-122.350, 47.650), (-122.300, 47.600), (-122.325, 47.625))
        gdf = add_geohash_columns(gdf, precision_partition=2, precision_sort=8)

        write_partitioned_dataset(gdf, tmp_path / "out")

        part_file = next((tmp_path / "out").rglob("*.parquet"))
        tbl = pq.read_table(part_file)
        # Reconstruct geohash_sort from written geometry to verify ordering.
        # Easier: just confirm the written rows are the same count and the
        # file is readable — ordering is implicit from the sort_values call.
        assert tbl.num_rows == 3


# ---------------------------------------------------------------------------
# add_geohash_column
# ---------------------------------------------------------------------------


class TestAddGeohashColumn:
    def test_adds_single_column_with_default_name(self):
        result = add_geohash_column(_point_gdf((-122.3, 47.6)), precision = 6)
        assert "geohash" in result.columns
        # Should not also add the two-column variants
        assert "geohash_prefix" not in result.columns
        assert "geohash_sort" not in result.columns

    def test_length_matches_precision(self):
        result = add_geohash_column(
            _point_gdf((-122.3, 47.6), (-77.0, 38.9)), precision = 7
        )
        assert all(len(v) == 7 for v in result["geohash"])

    def test_values_match_pygeohash(self):
        lon, lat = -122.3, 47.6
        result = add_geohash_column(_point_gdf((lon, lat)), precision = 6)
        assert result["geohash"].iloc[0] == pygeohash.encode(lat, lon, precision = 6)

    def test_custom_out_col_name(self):
        result = add_geohash_column(
            _point_gdf((-122.3, 47.6)), precision = 6, out_col = "gh"
        )
        assert "gh" in result.columns
        assert "geohash" not in result.columns


# ---------------------------------------------------------------------------
# compute_primary_osm_tag
# ---------------------------------------------------------------------------


class TestComputePrimaryOsmTag:
    FILTER_KEYS = [
        "shop", "healthcare", "leisure", "amenity",
        "tourism", "office", "craft", "historic",
    ]

    def _tagged_gdf(self, rows: list[dict]) -> gpd.GeoDataFrame:
        """Build a GDF with the standard OSM tag columns populated per row."""
        cols = {k: [] for k in self.FILTER_KEYS}
        for r in rows:
            for k in self.FILTER_KEYS:
                cols[k].append(r.get(k))
        gdf = gpd.GeoDataFrame(
            cols,
            geometry = [Point(-122.3, 47.6) for _ in rows],
            crs = "EPSG:4326",
        )
        return gdf

    def test_picks_highest_priority_when_multiple_present(self):
        """shop > healthcare > leisure > amenity — first match wins."""
        gdf = self._tagged_gdf([
            {"shop": "convenience", "amenity": "fuel"},  # both → primary=shop
            {"amenity": "restaurant"},                   # only amenity → amenity
            {"healthcare": "clinic", "amenity": "bank"},  # both → healthcare
        ])
        result = compute_primary_osm_tag(gdf, filter_keys = self.FILTER_KEYS)
        assert result["primary_tag"].tolist() == ["shop", "amenity", "healthcare"]

    def test_custom_out_col_name(self):
        gdf = self._tagged_gdf([{"shop": "bakery"}])
        result = compute_primary_osm_tag(
            gdf, filter_keys = self.FILTER_KEYS, out_col = "tag_key"
        )
        assert result["tag_key"].iloc[0] == "shop"
        assert "primary_tag" not in result.columns

    def test_raises_on_missing_filter_key_column(self):
        """If a filter_key isn't in the gdf, we should fail loudly."""
        gdf = self._tagged_gdf([{"shop": "bakery"}])
        gdf = gdf.drop(columns = ["historic"])
        with pytest.raises(KeyError, match = "historic"):
            compute_primary_osm_tag(gdf, filter_keys = self.FILTER_KEYS)

    def test_row_with_no_tags_gets_null_primary(self):
        gdf = self._tagged_gdf([{}])  # all nulls
        result = compute_primary_osm_tag(gdf, filter_keys = self.FILTER_KEYS)
        assert pd.isna(result["primary_tag"].iloc[0])


# ---------------------------------------------------------------------------
# write_label_partitioned_dataset
# ---------------------------------------------------------------------------


class TestWriteLabelPartitionedDataset:
    def _labeled_gdf(
        self,
        rows: list[tuple[str, str]],
        label_col: str = "shared_label",
    ) -> gpd.GeoDataFrame:
        """Build a GDF with (label, geohash) rows and a Point geometry each.

        Points are placed near Seattle; the geohash column is the actual sort
        key passed in (not derived), to make ordering assertions deterministic.
        """
        gdf = gpd.GeoDataFrame(
            {
                label_col: [r[0] for r in rows],
                "geohash": [r[1] for r in rows],
                "name": [f"p{i}" for i in range(len(rows))],
            },
            geometry = [Point(-122.3, 47.6) for _ in rows],
            crs = "EPSG:4326",
        )
        return gdf

    def test_creates_hive_dir_per_unique_value(self, tmp_path):
        gdf = self._labeled_gdf([("Pharmacy", "c23nb6"), ("Bakery", "c23nb7")])
        write_label_partitioned_dataset(
            gdf, tmp_path / "out", partition_col = "shared_label"
        )
        dirs = sorted(d.name for d in (tmp_path / "out").iterdir() if d.is_dir())
        assert dirs == ["shared_label=Bakery", "shared_label=Pharmacy"]

    def test_url_encodes_values_with_spaces(self, tmp_path):
        gdf = self._labeled_gdf([("Fast Food Restaurant", "c23nb6")])
        write_label_partitioned_dataset(
            gdf, tmp_path / "out", partition_col = "shared_label"
        )
        partition = tmp_path / "out" / "shared_label=Fast%20Food%20Restaurant"
        assert partition.is_dir()

    def test_partition_column_dropped_from_parquet(self, tmp_path):
        import pyarrow.parquet as pq

        gdf = self._labeled_gdf([("Pharmacy", "c23nb6")])
        write_label_partitioned_dataset(
            gdf, tmp_path / "out", partition_col = "shared_label"
        )
        part = next((tmp_path / "out").rglob("*.parquet"))
        names = pq.read_schema(part).names
        assert "shared_label" not in names

    def test_sort_column_retained_in_parquet(self, tmp_path):
        import pyarrow.parquet as pq

        gdf = self._labeled_gdf([("Pharmacy", "c23nb6")])
        write_label_partitioned_dataset(
            gdf, tmp_path / "out", partition_col = "shared_label"
        )
        part = next((tmp_path / "out").rglob("*.parquet"))
        names = pq.read_schema(part).names
        assert "geohash" in names
        assert "name" in names

    def test_rows_within_partition_sorted_by_sort_col(self, tmp_path):
        import pyarrow.parquet as pq

        # Out-of-order geohashes within one partition
        gdf = self._labeled_gdf(
            [("Pharmacy", "c23nbz"), ("Pharmacy", "c23nba"), ("Pharmacy", "c23nbm")]
        )
        write_label_partitioned_dataset(
            gdf, tmp_path / "out", partition_col = "shared_label"
        )
        part = next((tmp_path / "out").rglob("*.parquet"))
        tbl = pq.read_table(part)
        geohashes = tbl.column("geohash").to_pylist()
        assert geohashes == sorted(geohashes)

    def test_raises_when_partition_col_missing(self, tmp_path):
        gdf = self._labeled_gdf([("Pharmacy", "c23nb6")])
        with pytest.raises(KeyError, match = "partition_col"):
            write_label_partitioned_dataset(
                gdf, tmp_path / "out", partition_col = "nonexistent"
            )

    def test_raises_when_sort_col_missing(self, tmp_path):
        gdf = self._labeled_gdf([("Pharmacy", "c23nb6")])
        gdf = gdf.drop(columns = ["geohash"])
        # Re-add partition col expectations — geohash is required for sort
        with pytest.raises(KeyError, match = "sort_col"):
            write_label_partitioned_dataset(
                gdf, tmp_path / "out", partition_col = "shared_label"
            )

    def test_null_partition_values_skipped(self, tmp_path):
        gdf = self._labeled_gdf([("Pharmacy", "c23nb6"), (None, "c23nb7")])
        write_label_partitioned_dataset(
            gdf, tmp_path / "out", partition_col = "shared_label"
        )
        dirs = [d.name for d in (tmp_path / "out").iterdir() if d.is_dir()]
        # Only the non-null row produced a partition
        assert dirs == ["shared_label=Pharmacy"]

    def test_overwrite_false_raises_if_exists(self, tmp_path):
        out = tmp_path / "out"
        out.mkdir()
        gdf = self._labeled_gdf([("Pharmacy", "c23nb6")])
        with pytest.raises(FileExistsError, match = "overwrite=True"):
            write_label_partitioned_dataset(
                gdf, out, partition_col = "shared_label", overwrite = False
            )

    def test_overwrite_true_replaces_existing(self, tmp_path):
        out = tmp_path / "out"
        gdf = self._labeled_gdf([("Pharmacy", "c23nb6")])
        write_label_partitioned_dataset(gdf, out, partition_col = "shared_label")
        sentinel = out / "sentinel.txt"
        sentinel.write_text("old")

        write_label_partitioned_dataset(gdf, out, partition_col = "shared_label")
        assert not sentinel.exists()


# ---------------------------------------------------------------------------
# End-to-end: DuckDB can decode URL-encoded Hive partition values
# ---------------------------------------------------------------------------


class TestDuckDBHiveRoundtrip:
    """Confirm that URL-encoded partition names round-trip through DuckDB."""

    def test_duckdb_reads_url_encoded_partition_value(self, tmp_path):
        duckdb = pytest.importorskip("duckdb")

        gdf = gpd.GeoDataFrame(
            {
                "shared_label": ["Fast Food Restaurant", "Bakery"],
                "geohash": ["c23nb6", "c23nb7"],
                "payload": [1, 2],
            },
            geometry = [Point(-122.3, 47.6), Point(-122.3, 47.6)],
            crs = "EPSG:4326",
        )
        write_label_partitioned_dataset(
            gdf, tmp_path / "out", partition_col = "shared_label"
        )

        glob = str(tmp_path / "out" / "**" / "*.parquet")
        rows = duckdb.sql(
            f"SELECT shared_label, COUNT(*) AS n "
            f"FROM read_parquet('{glob}', hive_partitioning = 1) "
            f"GROUP BY shared_label ORDER BY shared_label"
        ).fetchall()

        assert rows == [("Bakery", 1), ("Fast Food Restaurant", 1)]