dopecarpet/slice.py

import argparse
import duckdb
import geohash
import geopandas as gpd
import pandas as pd
from pathlib import Path
from shapely import wkb
from shapely.geometry import box
from tqdm import tqdm  # For progress bar
from typing import List, Dict, Optional

# Configure logging
import logging
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)
logger = logging.getLogger(__name__)

def get_geohashes_from_bbox(min_x, min_y, max_x, max_y) -> List[str]:
    """Generates all 4-character Geohashes that intersect a bounding box."""
    geohashes = set()
    step_size = 0.25  # 4-character Geohash grid size

    lat = min_y
    while lat <= max_y:
        lon = min_x
        while lon <= max_x:
            gh = geohash.encode(lat, lon, precision=4)
            geohashes.add(gh)
            lon += step_size
        lat += step_size

    return sorted(geohashes)

def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame, add_water: bool):
    """Appends new data to an existing Parquet file while ensuring correct merging and avoiding duplicates."""
    if file_path.exists():
        try:
            existing_data = gpd.read_parquet(file_path)
            
            # Remove sea polygons to avoid duplicates (only if we're adding water polygons)
            if add_water:
                existing_data = existing_data[~existing_data["tags"].apply(lambda tags: isinstance(tags, dict) and tags.get("natural") == "water")]

            # Ensure no duplicates based on feature_id (for land geometries)
            if "feature_id" in existing_data.columns and "feature_id" in new_data.columns:
                new_unique_data = new_data[~new_data["feature_id"].isin(existing_data["feature_id"])]
                combined_data = pd.concat([existing_data, new_unique_data], ignore_index=True)
            else:
                combined_data = pd.concat([existing_data, new_data], ignore_index=True)

            combined_data.to_parquet(file_path, index=False)
        except Exception as e:
            logger.warning(f"Error while merging {file_path}: {e}")
    else:
        new_data.to_parquet(file_path, index=False)

def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
    """Saves a GeoDataFrame as a GeoJSON file."""
    geojson_file = output_path / f"{geohash_code}.geojson"
    gdf.to_file(geojson_file, driver="GeoJSON")
    logger.info(f"✅ Saved GeoJSON: {geojson_file}")

def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
    """Saves a GeoDataFrame as a CBOR file (compact binary format)."""
    import cbor2
    cbor_file = output_path / f"{geohash_code}.cbor"
    geojson_dict = json.loads(gdf.to_json())
    with open(cbor_file, "wb") as f:
        cbor2.dump(geojson_dict, f)
    logger.info(f"✅ Saved CBOR: {cbor_file}")

def geohash_bbox(geohash_code):
    """Returns a bounding box polygon for a geohash."""
    bbox = geohash.bbox(geohash_code)
    return box(bbox["w"], bbox["s"], bbox["e"], bbox["n"])

def ensure_crs_consistency(gdf):
    """Ensures the GeoDataFrame is in EPSG:4326 to prevent CRS mismatches."""
    if gdf.crs is None or gdf.crs.to_string() != "EPSG:4326":
        gdf = gdf.to_crs("EPSG:4326")
    return gdf

def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.GeoDataFrame):
    """Adds water polygons to a geohash tile."""
    try:
        existing_gdf = gpd.read_parquet(geohash_file)
    except Exception as e:
        logger.warning(f"Failed to read {geohash_file}. Skipping. Error: {e}")
        return False

    existing_gdf = ensure_crs_consistency(existing_gdf)

    # Get the geohash bounding box
    bbox = geohash_bbox(geohash_code)

    # Extract water polygons that overlap with this geohash
    water_in_tile = water_gdf[water_gdf.intersects(bbox)].copy()

    if water_in_tile.empty:
        return False  # No water polygons added

    # Clip water polygons to the geohash tile boundary
    water_in_tile["geometry"] = water_in_tile.intersection(bbox)

    # Assign "water" metadata and a unique identifier for water polygons
    water_in_tile["feature_id"] = [f"water_{geohash_code}_{i}" for i in range(len(water_in_tile))]
    water_in_tile["tags"] = [{"natural": "water", "water": "sea"}] * len(water_in_tile)

    # Ensure columns match before merging
    for col in ["feature_id", "tags"]:
        if col not in existing_gdf.columns:
            existing_gdf[col] = None
        if col not in water_in_tile.columns:
            water_in_tile[col] = None

    # Merge updated water polygons into geohash file
    updated_gdf = gpd.GeoDataFrame(pd.concat([existing_gdf, water_in_tile], ignore_index=True), crs="EPSG:4326")

    # Save back to parquet
    updated_gdf.to_parquet(geohash_file, index=False)
    return True  # Water polygons added

def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: Optional[str], export_parquet: bool, export_geojson: bool, export_cbor: bool):
    """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles and optionally adds water polygons."""
    input_path, output_path = Path(input_file), Path(output_dir)
    if not input_path.exists():
        raise FileNotFoundError(f"File '{input_file}' not found.")
    output_path.mkdir(parents=True, exist_ok=True)

    # Load water dataset if provided
    water_gdf = None
    if water_file:
        water_path = Path(water_file)
        if not water_path.exists():
            raise FileNotFoundError(f"File '{water_file}' not found.")
        water_gdf = ensure_crs_consistency(gpd.read_parquet(water_path))

    # Connect to DuckDB and load spatial extension
    con = duckdb.connect()
    con.execute("INSTALL spatial; LOAD spatial;")

    # Load the GeoParquet file into DuckDB
    con.execute(f"CREATE TEMP TABLE geoparquet AS SELECT * FROM read_parquet('{input_path}') WHERE ST_GeometryType(geometry) IS NOT NULL AND ST_GeometryType(geometry) != 'POINT'")

    # Precompute bounding boxes for all features
    bbox_results = con.execute("""
        SELECT feature_id, rowid, ST_XMin(ST_Envelope(geometry)), ST_YMin(ST_Envelope(geometry)),
                     ST_XMax(ST_Envelope(geometry)), ST_YMax(ST_Envelope(geometry))
        FROM geoparquet;
    """).fetchall()

    # Map geohash codes to rowids
    geohash_mapping: Dict[str, List[int]] = {}
    for feature_id, rowid, min_x, min_y, max_x, max_y in bbox_results:
        intersecting_geohashes = get_geohashes_from_bbox(min_x, min_y, max_x, max_y)
        for geohash_code in intersecting_geohashes:
            if geohash_code not in geohash_mapping:
                geohash_mapping[geohash_code] = []
            geohash_mapping[geohash_code].append(rowid)

    # Process each geohash tile
    total_tiles = len(geohash_mapping)
    tiles_with_water = 0

    with tqdm(total=total_tiles, desc="Processing geohash tiles") as pbar:
        for geohash_code, rowids in geohash_mapping.items():
            geohash_bbox = geohash.bbox(geohash_code)

            # Fetch and clip geometries for this geohash tile
            filtered_data = con.execute(f"""
                SELECT feature_id, tags, 
                       ST_AsWKB(ST_Intersection(geometry, 
                                                ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, 
                                                                {geohash_bbox["e"]}, {geohash_bbox["n"]}))) 
                       AS clipped_geom 
                FROM geoparquet 
                WHERE rowid IN ({','.join(map(str, rowids))})
                AND ST_Intersects(geometry, 
                                  ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, 
                                                  {geohash_bbox["e"]}, {geohash_bbox["n"]}));
            """).fetchdf()

            if not filtered_data.empty:
                # Convert WKB geometries to Shapely geometries
                filtered_data["geometry"] = filtered_data["clipped_geom"].apply(
                    lambda x: wkb.loads(bytes(x)) if isinstance(x, (bytes, bytearray)) else None
                )
                filtered_data.drop(columns=["clipped_geom"], inplace=True)

                # Drop rows with invalid geometries
                filtered_data = filtered_data.dropna(subset=["geometry"])

                if not filtered_data.empty:
                    gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326")

                    # Export to GeoParquet
                    if export_parquet:
                        geohash_file = output_path / f"{geohash_code}.parquet"
                        append_to_parquet(geohash_file, gdf, add_water=(water_gdf is not None))
                        if water_gdf is not None and add_water_to_geohash(geohash_code, geohash_file, water_gdf):
                            tiles_with_water += 1

                    # Export to GeoJSON
                    if export_geojson:
                        save_as_geojson(geohash_code, gdf, output_path)

                    # Export to CBOR
                    if export_cbor:
                        save_as_cbor(geohash_code, gdf, output_path)

            pbar.update(1)

    con.close()
    print(f"🎉 Processing complete! Generated {total_tiles} tiles, {tiles_with_water} of which had water polygons added.")

def main():
    parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, optionally add water polygons, and export in desired formats.")
    parser.add_argument("-i", "--input", required=True, help="Path to the input GeoParquet file.")
    parser.add_argument("-o", "--output", required=True, help="Directory to save output files.")
    parser.add_argument("-w", "--water", required=False, help="Path to the water.parquet file.")
    parser.add_argument("--parquet", action="store_true", help="Export as GeoParquet")
    parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON")
    parser.add_argument("--cbor", action="store_true", help="Export as CBOR")
    
    args = parser.parse_args()

    if not (args.parquet or args.geojson or args.cbor):
        args.parquet = True  # Default to GeoParquet if no options are given

    try:
        slice_and_split_geoparquet(args.input, args.output, args.water, args.parquet, args.geojson, args.cbor)
    except Exception as e:
        logger.error(f"❌ Error: {e}")

if __name__ == "__main__":
    main()