merged and faster slicer
This commit is contained in:
randogoth
parent
189f135840
commit
2e06b2c8bf
6 changed files with 90 additions and 193 deletions
154
slice.py
154
slice.py
|
|
@ -3,12 +3,11 @@ import duckdb
|
|||
import geohash
|
||||
import geopandas as gpd
|
||||
import pandas as pd
|
||||
import json
|
||||
import cbor2
|
||||
from pathlib import Path
|
||||
from shapely import wkb
|
||||
from shapely.geometry import box
|
||||
from typing import List, Dict
|
||||
from tqdm import tqdm # For progress bar
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
# Configure logging
|
||||
import logging
|
||||
|
|
@ -31,39 +30,40 @@ def get_geohashes_from_bbox(min_x, min_y, max_x, max_y) -> List[str]:
|
|||
|
||||
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."""
|
||||
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)
|
||||
combined_data = pd.concat([existing_data, new_data], ignore_index=True)
|
||||
|
||||
# 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")]
|
||||
|
||||
if "feature_id" in combined_data.columns:
|
||||
combined_data = combined_data.drop_duplicates(subset="feature_id", keep="last")
|
||||
# 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 = combined_data.reset_index(drop=True)
|
||||
combined_data.to_parquet(file_path, index=False)
|
||||
logger.info(f"✅ Appended new data to {file_path}")
|
||||
except Exception as e:
|
||||
logger.warning(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)
|
||||
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."""
|
||||
geojson_file = output_path / f"{geohash_code}.geojson"
|
||||
geojson_dict = json.loads(gdf.to_json())
|
||||
|
||||
with open(geojson_file, "w", encoding="utf-8") as f:
|
||||
json.dump(geojson_dict, f)
|
||||
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}")
|
||||
|
|
@ -81,14 +81,11 @@ def ensure_crs_consistency(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
|
||||
return False
|
||||
|
||||
existing_gdf = ensure_crs_consistency(existing_gdf)
|
||||
|
||||
|
|
@ -99,13 +96,13 @@ def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.G
|
|||
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
|
||||
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
|
||||
# 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
|
||||
|
|
@ -120,22 +117,22 @@ def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.G
|
|||
|
||||
# Save back to parquet
|
||||
updated_gdf.to_parquet(geohash_file, index=False)
|
||||
logger.info(f"Updated geohash file {geohash_file} with water polygons.")
|
||||
return True # Water polygons added
|
||||
|
||||
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)
|
||||
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.")
|
||||
if not water_path.exists():
|
||||
raise FileNotFoundError(f"File '{water_file}' not found.")
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
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))
|
||||
# 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()
|
||||
|
|
@ -161,56 +158,65 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: str
|
|||
geohash_mapping[geohash_code].append(rowid)
|
||||
|
||||
# Process each geohash tile
|
||||
for geohash_code, rowids in geohash_mapping.items():
|
||||
geohash_bbox = geohash.bbox(geohash_code)
|
||||
total_tiles = len(geohash_mapping)
|
||||
tiles_with_water = 0
|
||||
|
||||
# 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()
|
||||
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)
|
||||
|
||||
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"])
|
||||
# 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:
|
||||
gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326")
|
||||
# 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)
|
||||
|
||||
# Export to desired formats
|
||||
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)
|
||||
|
||||
if export_cbor:
|
||||
save_as_cbor(geohash_code, gdf, output_path)
|
||||
# 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()
|
||||
logger.info("🎉 Processing complete!")
|
||||
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, add water polygons, and export in desired formats.")
|
||||
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=True, help="Path to the water.parquet file.")
|
||||
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")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue