slice and water

slice.py

@@ -7,8 +7,14 @@ import json
 import cbor2
 from pathlib import Path
 from shapely import wkb
+from shapely.geometry import box
 from typing import List, Dict
 
+# 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()
@@ -37,12 +43,12 @@ def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame):
 
             combined_data = combined_data.reset_index(drop=True)
             combined_data.to_parquet(file_path, index=False)
-            print(f"✅ Appended new data to {file_path}")
+            logger.info(f"✅ Appended new data to {file_path}")
         except Exception as e:
-            print(f"❌ Error while merging {file_path}: {e}")
+            logger.warning(f"❌ Error while merging {file_path}: {e}")
     else:
         new_data.to_parquet(file_path, index=False)
-        print(f"✅ Created new file: {file_path}")
+        logger.info(f"✅ Created new file: {file_path}")
 
 def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
     """Saves a GeoDataFrame as a GeoJSON file."""
@@ -51,7 +57,7 @@ def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path)
 
     with open(geojson_file, "w", encoding="utf-8") as f:
         json.dump(geojson_dict, f)
-    print(f"✅ Saved GeoJSON: {geojson_file}")
+    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)."""
@@ -60,16 +66,76 @@ def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
 
     with open(cbor_file, "wb") as f:
         cbor2.dump(geojson_dict, f)
-    print(f"✅ Saved CBOR: {cbor_file}")
+    logger.info(f"✅ Saved CBOR: {cbor_file}")
 
-def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet: bool, export_geojson: bool, export_cbor: bool):
+def geohash_bbox(geohash_code):
-    """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles."""
+    """Returns a bounding box polygon for a geohash."""
-    input_path, output_path = Path(input_file), Path(output_dir)
+    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."""
+    logger.info(f"Processing {geohash_file}...")
+
+    # Load existing geohash file
+    try:
+        existing_gdf = gpd.read_parquet(geohash_file)
+    except Exception as e:
+        logger.warning(f"Failed to read {geohash_file}. Skipping. Error: {e}")
+        return
+
+    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:
+        logger.info(f"No water found for {geohash_code}. Skipping water addition.")
+        return
+
+    # Clip water polygons to the geohash tile boundary
+    water_in_tile["geometry"] = water_in_tile.intersection(bbox)
+
+    # Assign "water" metadata
+    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)
+    logger.info(f"Updated geohash file {geohash_file} with water polygons.")
+
+def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: str, export_parquet: bool, export_geojson: bool, export_cbor: bool):
+    """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles and adds water polygons."""
+    input_path, output_path, water_path = Path(input_file), Path(output_dir), Path(water_file)
     if not input_path.exists():
         raise FileNotFoundError(f"File '{input_file}' not found.")
+    if not water_path.exists():
+        raise FileNotFoundError(f"File '{water_file}' not found.")
     output_path.mkdir(parents=True, exist_ok=True)
 
-    print(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...")
+    logger.info(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...")
+
+    # Load water dataset
+    logger.info(f"Loading water dataset from {water_file}...")
+    water_gdf = ensure_crs_consistency(gpd.read_parquet(water_path))
 
     # Connect to DuckDB and load spatial extension
     con = duckdb.connect()
@@ -129,6 +195,7 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet:
                 if export_parquet:
                     geohash_file = output_path / f"{geohash_code}.parquet"
                     append_to_parquet(geohash_file, gdf)
+                    add_water_to_geohash(geohash_code, geohash_file, water_gdf)
                 
                 if export_geojson:
                     save_as_geojson(geohash_code, gdf, output_path)
@@ -137,12 +204,13 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet:
                     save_as_cbor(geohash_code, gdf, output_path)
 
     con.close()
-    print("🎉 Processing complete!")
+    logger.info("🎉 Processing complete!")
 
 def main():
-    parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles with export format options.")
+    parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, 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=True, 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")
@@ -153,9 +221,9 @@ def main():
         args.parquet = True  # Default to GeoParquet if no options are given
 
     try:
-        slice_and_split_geoparquet(args.input, args.output, args.parquet, args.geojson, args.cbor)
+        slice_and_split_geoparquet(args.input, args.output, args.water, args.parquet, args.geojson, args.cbor)
     except Exception as e:
-        print(f"❌ Error: {e}")
+        logger.error(f"❌ Error: {e}")
 
 if __name__ == "__main__":
     main()