init

slice.py

@@ -0,0 +1,155 @@
+import argparse
+import duckdb
+import geohash
+import geopandas as gpd
+import pandas as pd
+import pyarrow.parquet as pq
+import pyarrow as pa
+from pathlib import Path
+from shapely import wkb
+from typing import List, Dict
+
+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):
+    """Appends new data to an existing Parquet file while ensuring correct merging."""
+    if file_path.exists():
+        try:
+            # ✅ Load existing data
+            existing_data = gpd.read_parquet(file_path)
+            
+            # ✅ Merge with new data
+            combined_data = pd.concat([existing_data, new_data], ignore_index=True)
+
+            # ✅ Drop duplicates *only if feature_id exists*
+            if "feature_id" in combined_data.columns:
+                combined_data = combined_data.drop_duplicates(subset="feature_id", keep="last")
+
+            # ✅ Reset index before saving
+            combined_data = combined_data.reset_index(drop=True)
+
+            # ✅ Save back to the same file
+            combined_data.to_parquet(file_path, index=False)
+            print(f"✅ Appended new data to {file_path}")
+
+        except Exception as e:
+            print(f"❌ Error while merging {file_path}: {e}")
+
+    else:
+        # If the file doesn't exist, create it
+        new_data.to_parquet(file_path, index=False)
+        print(f"✅ Created new file: {file_path}")
+
+def slice_and_split_geoparquet(input_file: str, output_dir: str):
+    """
+    Splits the GeoParquet file into multiple files based on 4-character Geohash boxes.
+    Appends new data to existing Geohash Parquet files while preventing duplicate geometries.
+    """
+    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)
+
+    print(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...")
+
+    con = duckdb.connect()
+    con.execute("INSTALL spatial; LOAD spatial;")
+
+    # Create a temporary table of all non POINT geometries
+    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'")
+
+    # Compute bounding box for each geometry
+    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()
+
+    # Compute all intersecting geohashes
+    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)
+
+    # Dictionary to store in-memory results before writing
+    geohash_data: Dict[str, gpd.GeoDataFrame] = {}
+
+    for geohash_code, rowids in geohash_mapping.items():
+        geohash_bbox = geohash.bbox(geohash_code)
+
+        # Clip geometries inside this geohash box
+        filtered_data = con.execute(f"""
+            SELECT feature_id, rowid, 
+                   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:
+            # ✅ Remove non-WKB values
+            filtered_data = filtered_data[filtered_data["clipped_geom"].apply(lambda x: isinstance(x, (bytes, bytearray)))]
+
+            if not filtered_data.empty:
+                # ✅ Convert WKB to Shapely geometries safely
+                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)
+
+                # Remove rows where geometry conversion failed
+                filtered_data = filtered_data.dropna(subset=["geometry"])
+
+                if not filtered_data.empty:
+                    # Convert to GeoPandas GeoDataFrame
+                    gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326")
+
+                    # Store in memory
+                    geohash_data[geohash_code] = gdf
+
+    # Step 4: Write each geohash's data **only once** and prevent duplicates
+    for geohash_code, gdf in geohash_data.items():
+        geohash_file = output_path / f"{geohash_code}.parquet"
+        append_to_parquet(geohash_file, gdf)
+        print(f"✅ Updated: {geohash_file}")
+
+    con.close()
+    print("🎉 Processing complete!")
+
+def main():
+    parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles.")
+    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 Geohash tiles.")
+    
+    args = parser.parse_args()
+
+    try:
+        slice_and_split_geoparquet(args.input, args.output)
+    except Exception as e:
+        print(f"❌ Error: {e}")
+
+if __name__ == "__main__":
+    main()