merged and faster slicer

This commit is contained in:
randogoth 2025-03-04 16:17:56 +00:00
parent 189f135840
commit 2e06b2c8bf
6 changed files with 90 additions and 193 deletions

3
.gitignore vendored
View file

@ -1,4 +1,5 @@
*.parquet *.parquet
*.pbf *.pbf
*.png *.png
*.pyc *.pyc
data/

View file

@ -4,14 +4,14 @@ extract input output:
download region output: download region output:
quackosm --geom-filter-geocode {{region}} --osm-extract-source Geofabrik --osm-tags-filter-file tags.json --compact --output {{output}} quackosm --geom-filter-geocode {{region}} --osm-extract-source Geofabrik --osm-tags-filter-file tags.json --compact --output {{output}}
convert input output export_flags='--parquet': convert input water='water.parquet' output='geohash' export_flags='--parquet':
python3 slice.py --input {{input}} --output {{output}} {{export_flags}} python3 slice.py -i {{input}} -w {{water}} -o {{output}} {{export_flags}}
water input output export_flags='--parquet': water coast='coastline.parquet' water='water.parquet' output='geohash':
python3 water.py -i {{input}} -o {{output}} -c coastline.parquet {{export_flags}} python3 water.py {{coast}} {{water}} {{output}}
shp input output: shp input output:
ogr2ogr -of 'Parquet' {{output}} {{input}} ogr2ogr -s_srs EPSG:4326 -t_srs EPSG:4326 -a_srs EPSG:4326 -of 'Parquet' {{output}} {{input}}
plot file: plot file:
#!/usr/bin/env bash #!/usr/bin/env bash
@ -24,7 +24,7 @@ process file:
echo "Extracting Data from OpenStreetMaps PBF" echo "Extracting Data from OpenStreetMaps PBF"
quackosm {{file}} --osm-tags-filter-file tags.json --compact --output ${stem}.parquet quackosm {{file}} --osm-tags-filter-file tags.json --compact --output ${stem}.parquet
echo "Generating geohash parquet files" echo "Generating geohash parquet files"
python3 slice.py --input ${stem}.parquet --output geohash python3 slice.py --input ${stem}.parquet --water water.parquet --output geohash --parquet
echo "Cleaning up" echo "Cleaning up"
rm ${stem}.parquet rm ${stem}.parquet

View file

@ -126,10 +126,10 @@ just download '<name of the region>' <extracted_data.parquet>
The original PBF file remains saved in the `files/` folder. The original PBF file remains saved in the `files/` folder.
### Split Parquet File to GeoHash Parquet Files ### Split Parquet File to GeoHash Parquet Files and add Sea Polygons
```bash ```bash
just convert <extracted_data.parquet> <geohash_parquet_folder> just convert <extracted_data.parquet> <water_polygons.parquet> <geohash_parquet_folder>
``` ```
### Plot Parquet File as PNG ### Plot Parquet File as PNG

154
slice.py
View file

@ -3,12 +3,11 @@ import duckdb
import geohash import geohash
import geopandas as gpd import geopandas as gpd
import pandas as pd import pandas as pd
import json
import cbor2
from pathlib import Path from pathlib import Path
from shapely import wkb from shapely import wkb
from shapely.geometry import box 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 # Configure logging
import 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) return sorted(geohashes)
def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame): 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.""" """Appends new data to an existing Parquet file while ensuring correct merging and avoiding duplicates."""
if file_path.exists(): if file_path.exists():
try: try:
existing_data = gpd.read_parquet(file_path) 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: # Ensure no duplicates based on feature_id (for land geometries)
combined_data = combined_data.drop_duplicates(subset="feature_id", keep="last") 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) combined_data.to_parquet(file_path, index=False)
logger.info(f"✅ Appended new data to {file_path}")
except Exception as e: except Exception as e:
logger.warning(f"Error while merging {file_path}: {e}") logger.warning(f"Error while merging {file_path}: {e}")
else: else:
new_data.to_parquet(file_path, index=False) 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): def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
"""Saves a GeoDataFrame as a GeoJSON file.""" """Saves a GeoDataFrame as a GeoJSON file."""
geojson_file = output_path / f"{geohash_code}.geojson" geojson_file = output_path / f"{geohash_code}.geojson"
geojson_dict = json.loads(gdf.to_json()) gdf.to_file(geojson_file, driver="GeoJSON")
with open(geojson_file, "w", encoding="utf-8") as f:
json.dump(geojson_dict, f)
logger.info(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): def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
"""Saves a GeoDataFrame as a CBOR file (compact binary format).""" """Saves a GeoDataFrame as a CBOR file (compact binary format)."""
import cbor2
cbor_file = output_path / f"{geohash_code}.cbor" cbor_file = output_path / f"{geohash_code}.cbor"
geojson_dict = json.loads(gdf.to_json()) geojson_dict = json.loads(gdf.to_json())
with open(cbor_file, "wb") as f: with open(cbor_file, "wb") as f:
cbor2.dump(geojson_dict, f) cbor2.dump(geojson_dict, f)
logger.info(f"✅ Saved CBOR: {cbor_file}") 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): def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.GeoDataFrame):
"""Adds water polygons to a geohash tile.""" """Adds water polygons to a geohash tile."""
logger.info(f"Processing {geohash_file}...")
# Load existing geohash file
try: try:
existing_gdf = gpd.read_parquet(geohash_file) existing_gdf = gpd.read_parquet(geohash_file)
except Exception as e: except Exception as e:
logger.warning(f"Failed to read {geohash_file}. Skipping. Error: {e}") logger.warning(f"Failed to read {geohash_file}. Skipping. Error: {e}")
return return False
existing_gdf = ensure_crs_consistency(existing_gdf) 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() water_in_tile = water_gdf[water_gdf.intersects(bbox)].copy()
if water_in_tile.empty: if water_in_tile.empty:
logger.info(f"No water found for {geohash_code}. Skipping water addition.") return False # No water polygons added
return
# Clip water polygons to the geohash tile boundary # Clip water polygons to the geohash tile boundary
water_in_tile["geometry"] = water_in_tile.intersection(bbox) 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) water_in_tile["tags"] = [{"natural": "water", "water": "sea"}] * len(water_in_tile)
# Ensure columns match before merging # 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 # Save back to parquet
updated_gdf.to_parquet(geohash_file, index=False) 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): 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 adds water polygons.""" """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles and optionally adds water polygons."""
input_path, output_path, water_path = Path(input_file), Path(output_dir), Path(water_file) input_path, output_path = Path(input_file), Path(output_dir)
if not input_path.exists(): if not input_path.exists():
raise FileNotFoundError(f"File '{input_file}' not found.") 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) output_path.mkdir(parents=True, exist_ok=True)
logger.info(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...") # Load water dataset if provided
water_gdf = None
# Load water dataset if water_file:
logger.info(f"Loading water dataset from {water_file}...") water_path = Path(water_file)
water_gdf = ensure_crs_consistency(gpd.read_parquet(water_path)) 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 # Connect to DuckDB and load spatial extension
con = duckdb.connect() 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) geohash_mapping[geohash_code].append(rowid)
# Process each geohash tile # Process each geohash tile
for geohash_code, rowids in geohash_mapping.items(): total_tiles = len(geohash_mapping)
geohash_bbox = geohash.bbox(geohash_code) tiles_with_water = 0
# Fetch and clip geometries for this geohash tile with tqdm(total=total_tiles, desc="Processing geohash tiles") as pbar:
filtered_data = con.execute(f""" for geohash_code, rowids in geohash_mapping.items():
SELECT feature_id, tags, geohash_bbox = geohash.bbox(geohash_code)
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: # Fetch and clip geometries for this geohash tile
# Convert WKB geometries to Shapely geometries filtered_data = con.execute(f"""
filtered_data["geometry"] = filtered_data["clipped_geom"].apply( SELECT feature_id, tags,
lambda x: wkb.loads(bytes(x)) if isinstance(x, (bytes, bytearray)) else None ST_AsWKB(ST_Intersection(geometry,
) ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]},
filtered_data.drop(columns=["clipped_geom"], inplace=True) {geohash_bbox["e"]}, {geohash_bbox["n"]})))
AS clipped_geom
# Drop rows with invalid geometries FROM geoparquet
filtered_data = filtered_data.dropna(subset=["geometry"]) 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: 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 # Drop rows with invalid geometries
if export_parquet: filtered_data = filtered_data.dropna(subset=["geometry"])
geohash_file = output_path / f"{geohash_code}.parquet"
append_to_parquet(geohash_file, gdf) if not filtered_data.empty:
add_water_to_geohash(geohash_code, geohash_file, water_gdf) gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326")
if export_geojson: # Export to GeoParquet
save_as_geojson(geohash_code, gdf, output_path) if export_parquet:
geohash_file = output_path / f"{geohash_code}.parquet"
if export_cbor: append_to_parquet(geohash_file, gdf, add_water=(water_gdf is not None))
save_as_cbor(geohash_code, gdf, output_path) 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() 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(): 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("-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("-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("--parquet", action="store_true", help="Export as GeoParquet")
parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON") parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON")
parser.add_argument("--cbor", action="store_true", help="Export as CBOR") parser.add_argument("--cbor", action="store_true", help="Export as CBOR")

BIN
u33d.png

Binary file not shown.

Before

Width:  |  Height:  |  Size: 4 MiB

110
water.py
View file

@ -1,110 +0,0 @@
import os
import argparse
import pandas as pd
import geopandas as gpd
import geohash
import logging
from shapely.geometry import box
# Configure logging
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)
logger = logging.getLogger(__name__)
def load_geohash_files(folder):
"""Finds all existing geohash parquet files in the folder."""
geohash_files = {}
for filename in os.listdir(folder):
if filename.endswith(".parquet"):
geohash_code = filename[:-8] # Remove '.parquet' extension
geohash_files[geohash_code] = os.path.join(folder, filename)
logger.info(f"Found {len(geohash_files)} geohash parquet files in {folder}.")
return geohash_files
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 process_geohash_files(water_gdf, geohash_files):
"""Processes each geohash file: adds water polygons from water.parquet."""
# Precompute bounding boxes for all geohash codes
geohash_bboxes = {geohash_code: geohash_bbox(geohash_code) for geohash_code in geohash_files.keys()}
# Spatial index for faster spatial queries
water_sindex = water_gdf.sindex
for geohash_code, file_path in geohash_files.items():
logger.info(f"Processing {file_path}...")
# Load existing geohash file
try:
existing_gdf = gpd.read_parquet(file_path)
except Exception as e:
logger.warning(f"Failed to read {file_path}. Skipping. Error: {e}")
continue
existing_gdf = ensure_crs_consistency(existing_gdf)
# Get the precomputed bounding box for this geohash
bbox = geohash_bboxes[geohash_code]
# Use spatial index to find intersecting water polygons
possible_matches_index = list(water_sindex.intersection(bbox.bounds))
possible_matches = water_gdf.iloc[possible_matches_index]
water_in_tile = possible_matches[possible_matches.intersects(bbox)].copy()
if water_in_tile.empty:
logger.info(f"No water found for {geohash_code}. Skipping water addition.")
continue
# 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(file_path, index=False)
logger.info(f"Updated geohash file {file_path} with water polygons.")
def main():
parser = argparse.ArgumentParser(description="Add water polygons to geohash parquet files from a global water dataset.")
parser.add_argument("water", help="Path to the water.parquet file.")
parser.add_argument("geohash_folder", help="Path to the folder containing <geohash>.parquet files.")
args = parser.parse_args()
# Load water dataset
logger.info(f"Loading water dataset from {args.water}...")
water_gdf = ensure_crs_consistency(gpd.read_parquet(args.water))
# Get existing geohash parquet files
geohash_files = load_geohash_files(args.geohash_folder)
if not geohash_files:
logger.warning("No geohash parquet files found. Exiting.")
return
# Process geohash files
process_geohash_files(water_gdf, geohash_files)
logger.info("Processing complete. All water polygons added.")
if __name__ == "__main__":
main()