diff --git a/.gitignore b/.gitignore index 976e199..62c5106 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,5 @@ __pycache__/ *.pyc target/ .data/ +.bench/ +benchmark-results.jsonl diff --git a/bin/benchmark-pipeline b/bin/benchmark-pipeline new file mode 100755 index 0000000..57e4e54 --- /dev/null +++ b/bin/benchmark-pipeline @@ -0,0 +1,290 @@ +#!/usr/bin/env bash +# OpenInterstate pipeline benchmark harness. +# +# Downloads a small PBF (Rhode Island by default), runs the full pipeline +# from scratch, and records per-stage wall-clock times to a JSON log. +# +# Usage: +# bin/benchmark-pipeline [--runs N] [--pbf-url URL] [--label LABEL] [--log FILE] +# +# Requires: docker (compose plugin) + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="${OI_REPO_ROOT:-$(cd "$SCRIPT_DIR/.." && pwd)}" + +# --- defaults ---------------------------------------------------------------- + +RI_PBF_URL="https://download.geofabrik.de/north-america/us/rhode-island-latest.osm.pbf" +PBF_URL="$RI_PBF_URL" +RUNS=1 +LABEL="" +LOG_FILE="$REPO_ROOT/benchmark-results.jsonl" +DB_PORT="${OI_BENCH_DB_PORT:-55433}" +DB_NAME="osm" +DB_USER="osm" +DB_PASSWORD="osm_dev" +IMPORT_CACHE_MB="${OI_IMPORT_CACHE_MB:-2048}" +BENCH_WORK_DIR="${OI_BENCH_WORK_DIR:-$REPO_ROOT/.bench}" + +# --- parse args -------------------------------------------------------------- + +while [[ $# -gt 0 ]]; do + case "$1" in + --runs) RUNS="$2"; shift 2 ;; + --pbf-url) PBF_URL="$2"; shift 2 ;; + --label) LABEL="$2"; shift 2 ;; + --log) LOG_FILE="$2"; shift 2 ;; + --cache-mb) IMPORT_CACHE_MB="$2"; shift 2 ;; + --work-dir) BENCH_WORK_DIR="$2"; shift 2 ;; + -h|--help) + sed -n '2,/^$/s/^# \?//p' "$0" + exit 0 + ;; + *) echo "unknown arg: $1" >&2; exit 1 ;; + esac +done + +# --- helpers ----------------------------------------------------------------- + +log() { echo "[$(date '+%H:%M:%S')] $*" >&2; } +die() { echo "ERROR: $*" >&2; exit 1; } + +epoch_ms() { + if command -v gdate >/dev/null 2>&1; then + gdate +%s%3N + elif date +%s%3N >/dev/null 2>&1; then + date +%s%3N + else + python3 -c 'import time; print(int(time.time()*1000))' + fi +} + +time_stage() { + local name="$1"; shift + local t0 t1 elapsed_ms + log "▶ $name" + t0="$(epoch_ms)" + "$@" + t1="$(epoch_ms)" + elapsed_ms=$(( t1 - t0 )) + STAGE_TIMES+=("\"$name\": $elapsed_ms") + log "✓ $name ${elapsed_ms}ms ($(( elapsed_ms / 1000 ))s)" +} + +DB_CONTAINER="" +cleanup() { + if [[ -n "$DB_CONTAINER" ]]; then + docker rm -f "$DB_CONTAINER" >/dev/null 2>&1 || true + fi +} +trap cleanup EXIT + +runner() { + docker run --rm \ + --network=host \ + -v "$REPO_ROOT:/workspace" \ + -v "$BENCH_WORK_DIR:/bench" \ + -w /workspace \ + -e PGPASSWORD="$DB_PASSWORD" \ + openinterstate-bench-runner "$@" +} + +# --- pre-flight checks ------------------------------------------------------- + +command -v docker >/dev/null 2>&1 || die "docker is required" + +mkdir -p "$BENCH_WORK_DIR"/{source,state,releases,cache} + +# Build runner image (cached after first build) +log "Ensuring runner image" +docker build -q -t openinterstate-bench-runner \ + -f "$REPO_ROOT/docker/runner/Dockerfile" "$REPO_ROOT" >/dev/null + +# Pre-build the Rust derive binary so compilation doesn't pollute stage timing +log "Pre-building openinterstate-derive (release)" +runner cargo build --locked --release -p openinterstate-derive >/dev/null 2>&1 + +# Download PBF (cached) +SOURCE_PBF="$BENCH_WORK_DIR/source/$(basename "$PBF_URL")" +if [[ ! -f "$SOURCE_PBF" ]]; then + log "Downloading PBF: $PBF_URL" + curl -L --fail --progress-bar -o "$SOURCE_PBF" "$PBF_URL" +else + log "Using cached PBF: $SOURCE_PBF" +fi +PBF_SIZE=$(stat -f '%z' "$SOURCE_PBF" 2>/dev/null || stat -c '%s' "$SOURCE_PBF") + +# --- run loop ---------------------------------------------------------------- + +for (( run=1; run<=RUNS; run++ )); do + log "=== Run $run/$RUNS ===" + STAGE_TIMES=() + + # Clean slate for each run + cleanup + DB_CONTAINER="" + + FILTERED_PBF="$BENCH_WORK_DIR/source/filtered.osm.pbf" + RELATION_CACHE="$BENCH_WORK_DIR/cache/interstate-relations.tsv" + DATABASE_URL="postgres://${DB_USER}:${DB_PASSWORD}@127.0.0.1:${DB_PORT}/${DB_NAME}" + + # --- stage: prefilter ----------------------------------------------------- + time_stage "prefilter" \ + runner osmium tags-filter \ + "/bench/source/$(basename "$SOURCE_PBF")" \ + n/highway=motorway_junction \ + "n/amenity=fuel,restaurant,fast_food,cafe,charging_station" \ + "n/tourism=hotel,motel" \ + n/shop=gas \ + n/cuisine \ + "n/highway=rest_area,services" \ + n/amenity=toilets \ + w/highway=construction \ + "w/highway=motorway,motorway_link,trunk,trunk_link,rest_area,services" \ + "w/amenity=fuel,restaurant,fast_food,cafe,charging_station" \ + "w/tourism=hotel,motel" \ + w/shop=gas \ + w/cuisine \ + --overwrite \ + -o "/bench/source/filtered.osm.pbf" + + # --- stage: extract relations --------------------------------------------- + time_stage "extract_relations" \ + runner python3 /workspace/tooling/extract_interstate_relations.py \ + --source-pbf "/bench/source/$(basename "$SOURCE_PBF")" \ + --output "/bench/cache/interstate-relations.tsv" + + # --- stage: start database ------------------------------------------------ + DB_CONTAINER="openinterstate-bench-db-$$-$run" + start_db_fn() { + docker run --detach --rm \ + --name "$DB_CONTAINER" \ + --shm-size 2g \ + -e POSTGRES_DB="$DB_NAME" \ + -e POSTGRES_USER="$DB_USER" \ + -e POSTGRES_PASSWORD="$DB_PASSWORD" \ + -p "${DB_PORT}:5432" \ + postgis/postgis:16-3.4 \ + postgres \ + -c shared_buffers=512MB \ + -c effective_cache_size=2GB \ + -c maintenance_work_mem=512MB \ + -c work_mem=32MB \ + -c max_wal_size=32GB \ + -c min_wal_size=8GB \ + -c checkpoint_timeout=60min \ + -c checkpoint_completion_target=0.9 \ + -c wal_compression=on \ + -c wal_level=minimal \ + -c max_wal_senders=0 \ + -c archive_mode=off \ + -c synchronous_commit=off \ + -c fsync=off \ + -c full_page_writes=off \ + -c autovacuum=off \ + -c effective_io_concurrency=200 \ + -c random_page_cost=1.1 \ + >/dev/null + } + time_stage "start_db" start_db_fn + + # Wait for postgres to accept connections (from inside and via TCP from runner) + log "Waiting for PostGIS..." + for (( i=1; i<=120; i++ )); do + if docker exec "$DB_CONTAINER" pg_isready -U "$DB_USER" -d "$DB_NAME" >/dev/null 2>&1; then + if runner bash -c "PGPASSWORD=$DB_PASSWORD psql -h host.docker.internal -p $DB_PORT -U $DB_USER -d $DB_NAME -Atc 'SELECT 1'" >/dev/null 2>&1; then + break + fi + fi + sleep 1 + done + + # Bootstrap + time_stage "bootstrap_schema" \ + runner bash -c " + PGPASSWORD=$DB_PASSWORD psql -h host.docker.internal -p $DB_PORT -U $DB_USER -d $DB_NAME \ + -v ON_ERROR_STOP=1 -c 'CREATE EXTENSION IF NOT EXISTS postgis;' && + PGPASSWORD=$DB_PASSWORD psql -h host.docker.internal -p $DB_PORT -U $DB_USER -d $DB_NAME \ + -v ON_ERROR_STOP=1 -f /workspace/schema/bootstrap.sql + " + + # --- stage: osm2pgsql import ---------------------------------------------- + time_stage "osm2pgsql_import" \ + runner bash -c " + PGPASSWORD=$DB_PASSWORD osm2pgsql \ + --slim \ + --create \ + --output=flex \ + --style=/workspace/schema/osm2pgsql/openinterstate.lua \ + --database=$DB_NAME \ + --host=host.docker.internal \ + --port=$DB_PORT \ + -U $DB_USER \ + --cache=$IMPORT_CACHE_MB \ + /bench/source/filtered.osm.pbf + " + + # --- stage: derive SQL ---------------------------------------------------- + time_stage "derive_sql" \ + runner bash -c " + PGPASSWORD=$DB_PASSWORD psql \ + -h host.docker.internal -p $DB_PORT -U $DB_USER -d $DB_NAME \ + -v ON_ERROR_STOP=1 \ + -f /workspace/schema/derive.sql + " + + # --- stage: rust derive (graph + corridors + routes) ---------------------- + time_stage "rust_derive" \ + runner bash -c " + ./target/release/openinterstate-derive \ + --database-url 'postgres://$DB_USER:$DB_PASSWORD@host.docker.internal:$DB_PORT/$DB_NAME' \ + --interstate-relation-cache /bench/cache/interstate-relations.tsv \ + all + " + + # --- stage: export release ------------------------------------------------ + RELEASE_ID="bench-run-$run" + time_stage "export_release" \ + runner bash -c " + python3 /workspace/tooling/export_release.py \ + --database-url 'postgres://$DB_USER:$DB_PASSWORD@host.docker.internal:$DB_PORT/$DB_NAME' \ + --release-id $RELEASE_ID \ + --output-dir /bench/releases/$RELEASE_ID \ + --state-dir /bench/state \ + --source-pbf-file /bench/source/$(basename "$SOURCE_PBF") \ + --import-pbf-file /bench/source/filtered.osm.pbf + " + + # --- collect results ------------------------------------------------------ + TOTAL_MS=0 + for entry in "${STAGE_TIMES[@]}"; do + ms="${entry##*: }" + TOTAL_MS=$(( TOTAL_MS + ms )) + done + + RESULT_JSON=$(cat <> "$LOG_FILE" + log "Total: ${TOTAL_MS}ms ($(( TOTAL_MS / 1000 ))s)" + log "Stages: $(IFS=,; echo "${STAGE_TIMES[*]}")" + + # Cleanup for next run + cleanup + DB_CONTAINER="" +done + +log "Results appended to $LOG_FILE" diff --git a/crates/derive/src/graph/mod.rs b/crates/derive/src/graph/mod.rs index 6eb54ac..303c4d7 100644 --- a/crates/derive/src/graph/mod.rs +++ b/crates/derive/src/graph/mod.rs @@ -66,69 +66,130 @@ pub async fn build_graph( corridor_entries.len() ); - // Write edges + // Write edges in batches using UNNEST for bulk insert + const EDGE_BATCH: usize = 5_000; let mut tx = pool.begin().await?; - for (i, edge) in edges.iter().enumerate() { + for (batch_idx, chunk) in edges.chunks(EDGE_BATCH).enumerate() { + let mut ids = Vec::with_capacity(chunk.len()); + let mut highways = Vec::with_capacity(chunk.len()); + let mut components = Vec::with_capacity(chunk.len()); + let mut start_nodes = Vec::with_capacity(chunk.len()); + let mut end_nodes = Vec::with_capacity(chunk.len()); + let mut lengths = Vec::with_capacity(chunk.len()); + let mut geom_wkts = Vec::with_capacity(chunk.len()); + let mut min_lats = Vec::with_capacity(chunk.len()); + let mut max_lats = Vec::with_capacity(chunk.len()); + let mut min_lons = Vec::with_capacity(chunk.len()); + let mut max_lons = Vec::with_capacity(chunk.len()); + let mut polylines = Vec::with_capacity(chunk.len()); + let mut source_ways = Vec::with_capacity(chunk.len()); + let mut directions: Vec> = Vec::with_capacity(chunk.len()); + + for edge in chunk { + ids.push(edge.id.as_str()); + highways.push(edge.highway.as_str()); + components.push(edge.component); + start_nodes.push(edge.start_node); + end_nodes.push(edge.end_node); + lengths.push(edge.length_m); + geom_wkts.push(edge.geom_wkt.as_str()); + min_lats.push(edge.min_lat); + max_lats.push(edge.max_lat); + min_lons.push(edge.min_lon); + max_lons.push(edge.max_lon); + polylines.push(edge.polyline_json.as_str()); + source_ways.push(edge.source_way_ids_json.as_str()); + directions.push(edge.direction.clone()); + } + sqlx::query( "INSERT INTO highway_edges \ (id, highway, component, start_node, end_node, length_m, \ geom, min_lat, max_lat, min_lon, max_lon, polyline_json, source_way_ids_json, direction) \ - VALUES ($1, $2, $3, $4, $5, $6, \ - ST_GeomFromText($7, 4326), $8, $9, $10, $11, $12, $13, $14) \ + SELECT id, highway, component, start_node, end_node, length_m, \ + ST_GeomFromText(geom_wkt, 4326), min_lat, max_lat, min_lon, max_lon, \ + polyline_json, source_way_ids_json, direction \ + FROM UNNEST($1::text[], $2::text[], $3::int[], $4::int8[], $5::int8[], $6::int[], \ + $7::text[], $8::float8[], $9::float8[], $10::float8[], $11::float8[], \ + $12::text[], $13::text[], $14::text[]) \ + AS t(id, highway, component, start_node, end_node, length_m, \ + geom_wkt, min_lat, max_lat, min_lon, max_lon, polyline_json, source_way_ids_json, direction) \ ON CONFLICT (id) DO NOTHING", ) - .bind(&edge.id) - .bind(&edge.highway) - .bind(edge.component) - .bind(edge.start_node) - .bind(edge.end_node) - .bind(edge.length_m) - .bind(&edge.geom_wkt) - .bind(edge.min_lat) - .bind(edge.max_lat) - .bind(edge.min_lon) - .bind(edge.max_lon) - .bind(&edge.polyline_json) - .bind(&edge.source_way_ids_json) - .bind(&edge.direction) + .bind(&ids) + .bind(&highways) + .bind(&components) + .bind(&start_nodes) + .bind(&end_nodes) + .bind(&lengths) + .bind(&geom_wkts) + .bind(&min_lats) + .bind(&max_lats) + .bind(&min_lons) + .bind(&max_lons) + .bind(&polylines) + .bind(&source_ways) + .bind(&directions) .execute(&mut *tx) .await?; - if (i + 1) % 10_000 == 0 { - tracing::info!(" edges: {}/{}", i + 1, edges.len()); + let done = (batch_idx + 1) * EDGE_BATCH; + if done % 10_000 < EDGE_BATCH { + tracing::info!(" edges: {}/{}", done.min(edges.len()), edges.len()); } } tx.commit().await?; - // Write corridor entries + // Write corridor entries in batches using UNNEST + const CORRIDOR_BATCH: usize = 5_000; let mut tx = pool.begin().await?; - let mut inserted = 0_usize; - for (i, entry) in corridor_entries.iter().enumerate() { - let result = sqlx::query( + for (batch_idx, chunk) in corridor_entries.chunks(CORRIDOR_BATCH).enumerate() { + let mut exit_ids = Vec::with_capacity(chunk.len()); + let mut c_highways = Vec::with_capacity(chunk.len()); + let mut c_components = Vec::with_capacity(chunk.len()); + let mut c_node_ids = Vec::with_capacity(chunk.len()); + let mut c_directions: Vec> = Vec::with_capacity(chunk.len()); + + for entry in chunk { + exit_ids.push(entry.exit_id.as_str()); + c_highways.push(entry.highway.as_str()); + c_components.push(entry.component); + c_node_ids.push(entry.node_id); + c_directions.push(entry.direction.clone()); + } + + sqlx::query( "INSERT INTO exit_corridors (exit_id, highway, graph_component, graph_node, direction) \ - VALUES ($1, $2, $3, $4, $5) \ + SELECT exit_id, highway, graph_component, graph_node, direction \ + FROM UNNEST($1::text[], $2::text[], $3::int[], $4::int8[], $5::text[]) \ + AS t(exit_id, highway, graph_component, graph_node, direction) \ ON CONFLICT (exit_id, highway) DO UPDATE SET \ graph_component = EXCLUDED.graph_component, \ graph_node = EXCLUDED.graph_node, \ direction = EXCLUDED.direction", ) - .bind(&entry.exit_id) - .bind(&entry.highway) - .bind(entry.component) - .bind(entry.node_id) - .bind(&entry.direction) + .bind(&exit_ids) + .bind(&c_highways) + .bind(&c_components) + .bind(&c_node_ids) + .bind(&c_directions) .execute(&mut *tx) .await?; - if result.rows_affected() > 0 { - inserted += 1; - } - if (i + 1) % 10_000 == 0 { - tracing::info!(" corridor entries: {}/{}", i + 1, corridor_entries.len()); + let done = (batch_idx + 1) * CORRIDOR_BATCH; + if done % 10_000 < CORRIDOR_BATCH { + tracing::info!( + " corridor entries: {}/{}", + done.min(corridor_entries.len()), + corridor_entries.len() + ); } } tx.commit().await?; - tracing::info!(" Inserted {} exit_corridors entries", inserted); + tracing::info!( + " Inserted {} exit_corridors entries", + corridor_entries.len() + ); Ok(edges.len()) } diff --git a/crates/derive/src/graph/relation_corridors.rs b/crates/derive/src/graph/relation_corridors.rs index 6e1ce9d..d6bbc68 100644 --- a/crates/derive/src/graph/relation_corridors.rs +++ b/crates/derive/src/graph/relation_corridors.rs @@ -2353,47 +2353,120 @@ async fn write_corridors( .execute(&mut *tx) .await?; - for draft in drafts { - sqlx::query( - "INSERT INTO corridors \ - (corridor_id, highway, canonical_direction, root_relation_id, geometry_json, source_way_ids_json) \ - VALUES ($1, $2, $3, $4, $5, $6)", - ) - .bind(draft.corridor_id) - .bind(&draft.highway) - .bind(&draft.canonical_direction) - .bind(draft.root_relation_id) - .bind(&draft.geometry_json) - .bind(serde_json::to_string(&draft.source_way_ids).unwrap_or_else(|_| "[]".to_string())) - .execute(&mut *tx) - .await?; + // Batch insert corridors using UNNEST + { + let mut c_ids = Vec::with_capacity(drafts.len()); + let mut c_highways = Vec::with_capacity(drafts.len()); + let mut c_directions = Vec::with_capacity(drafts.len()); + let mut c_relation_ids = Vec::with_capacity(drafts.len()); + let mut c_geom_jsons = Vec::with_capacity(drafts.len()); + let mut c_source_ways = Vec::with_capacity(drafts.len()); + + for draft in drafts.iter() { + c_ids.push(draft.corridor_id); + c_highways.push(draft.highway.as_str()); + c_directions.push(draft.canonical_direction.as_str()); + c_relation_ids.push(draft.root_relation_id); + c_geom_jsons.push(draft.geometry_json.as_str()); + c_source_ways.push( + serde_json::to_string(&draft.source_way_ids) + .unwrap_or_else(|_| "[]".to_string()), + ); + } + let c_source_ways_refs: Vec<&str> = c_source_ways.iter().map(|s| s.as_str()).collect(); - for (idx, exit) in draft.exits.iter().enumerate() { + const CORRIDOR_BATCH: usize = 500; + for (i, chunk_start) in (0..drafts.len()).step_by(CORRIDOR_BATCH).enumerate() { + let chunk_end = (chunk_start + CORRIDOR_BATCH).min(drafts.len()); + sqlx::query( + "INSERT INTO corridors \ + (corridor_id, highway, canonical_direction, root_relation_id, geometry_json, source_way_ids_json) \ + SELECT corridor_id, highway, canonical_direction, root_relation_id, geometry_json, source_way_ids_json \ + FROM UNNEST($1::int[], $2::text[], $3::text[], $4::int8[], $5::text[], $6::text[]) \ + AS t(corridor_id, highway, canonical_direction, root_relation_id, geometry_json, source_way_ids_json)", + ) + .bind(&c_ids[chunk_start..chunk_end]) + .bind(&c_highways[chunk_start..chunk_end]) + .bind(&c_directions[chunk_start..chunk_end]) + .bind(&c_relation_ids[chunk_start..chunk_end]) + .bind(&c_geom_jsons[chunk_start..chunk_end]) + .bind(&c_source_ways_refs[chunk_start..chunk_end]) + .execute(&mut *tx) + .await?; + + let _ = i; // suppress unused + } + } + + // Batch insert corridor exits using UNNEST + { + let mut ce_corridor_ids = Vec::new(); + let mut ce_indices = Vec::new(); + let mut ce_exit_ids = Vec::new(); + let mut ce_refs: Vec> = Vec::new(); + let mut ce_names: Vec> = Vec::new(); + let mut ce_lats = Vec::new(); + let mut ce_lons = Vec::new(); + + for draft in drafts.iter() { + for (idx, exit) in draft.exits.iter().enumerate() { + ce_corridor_ids.push(draft.corridor_id); + ce_indices.push(idx as i32); + ce_exit_ids.push(exit.exit_id.as_str()); + ce_refs.push(exit.ref_val.clone()); + ce_names.push(exit.name.clone()); + ce_lats.push(exit.lat); + ce_lons.push(exit.lon); + } + } + + const EXIT_BATCH: usize = 5_000; + for chunk_start in (0..ce_corridor_ids.len()).step_by(EXIT_BATCH) { + let chunk_end = (chunk_start + EXIT_BATCH).min(ce_corridor_ids.len()); sqlx::query( "INSERT INTO corridor_exits \ (corridor_id, corridor_index, exit_id, ref, name, lat, lon) \ - VALUES ($1, $2, $3, $4, $5, $6, $7)", + SELECT corridor_id, corridor_index, exit_id, ref_val, name, lat, lon \ + FROM UNNEST($1::int[], $2::int[], $3::text[], $4::text[], $5::text[], $6::float8[], $7::float8[]) \ + AS t(corridor_id, corridor_index, exit_id, ref_val, name, lat, lon)", ) - .bind(draft.corridor_id) - .bind(idx as i32) - .bind(&exit.exit_id) - .bind(&exit.ref_val) - .bind(&exit.name) - .bind(exit.lat) - .bind(exit.lon) + .bind(&ce_corridor_ids[chunk_start..chunk_end]) + .bind(&ce_indices[chunk_start..chunk_end]) + .bind(&ce_exit_ids[chunk_start..chunk_end]) + .bind(&ce_refs[chunk_start..chunk_end]) + .bind(&ce_names[chunk_start..chunk_end]) + .bind(&ce_lats[chunk_start..chunk_end]) + .bind(&ce_lons[chunk_start..chunk_end]) .execute(&mut *tx) .await?; } } + // Batch update highway_edges corridor_id using UNNEST let mut edges_updated = 0usize; - for draft in drafts { - for edge_id in &draft.edge_ids { - let result = sqlx::query("UPDATE highway_edges SET corridor_id = $2 WHERE id = $1") - .bind(edge_id) - .bind(draft.corridor_id) - .execute(&mut *tx) - .await?; + { + let mut upd_edge_ids = Vec::new(); + let mut upd_corridor_ids = Vec::new(); + + for draft in drafts.iter() { + for edge_id in &draft.edge_ids { + upd_edge_ids.push(edge_id.as_str()); + upd_corridor_ids.push(draft.corridor_id); + } + } + + const UPDATE_BATCH: usize = 5_000; + for chunk_start in (0..upd_edge_ids.len()).step_by(UPDATE_BATCH) { + let chunk_end = (chunk_start + UPDATE_BATCH).min(upd_edge_ids.len()); + let result = sqlx::query( + "UPDATE highway_edges SET corridor_id = t.corridor_id \ + FROM UNNEST($1::text[], $2::int[]) AS t(id, corridor_id) \ + WHERE highway_edges.id = t.id", + ) + .bind(&upd_edge_ids[chunk_start..chunk_end]) + .bind(&upd_corridor_ids[chunk_start..chunk_end]) + .execute(&mut *tx) + .await?; edges_updated += result.rows_affected() as usize; } } diff --git a/crates/derive/src/main.rs b/crates/derive/src/main.rs index 6c40643..520aa90 100644 --- a/crates/derive/src/main.rs +++ b/crates/derive/src/main.rs @@ -73,7 +73,7 @@ async fn run(cli: Cli) -> anyhow::Result<()> { async fn connect_pool(database_url: &str) -> anyhow::Result { Ok(sqlx::postgres::PgPoolOptions::new() - .max_connections(20) + .max_connections(4) .connect(database_url) .await?) } diff --git a/crates/derive/src/routes.rs b/crates/derive/src/routes.rs index 9f2c208..95f951e 100644 --- a/crates/derive/src/routes.rs +++ b/crates/derive/src/routes.rs @@ -109,52 +109,121 @@ pub async fn build_reference_routes( .execute(&mut *tx) .await?; - for route in &routes { + // Batch insert routes using UNNEST + const ROUTE_BATCH: usize = 500; + for chunk in routes.chunks(ROUTE_BATCH) { + let mut r_ids = Vec::with_capacity(chunk.len()); + let mut r_highways = Vec::with_capacity(chunk.len()); + let mut r_dir_codes = Vec::with_capacity(chunk.len()); + let mut r_dir_labels = Vec::with_capacity(chunk.len()); + let mut r_display_names = Vec::with_capacity(chunk.len()); + let mut r_corridor_ids = Vec::with_capacity(chunk.len()); + let mut r_variant_ranks = Vec::with_capacity(chunk.len()); + let mut r_distances = Vec::with_capacity(chunk.len()); + let mut r_durations = Vec::with_capacity(chunk.len()); + let mut r_intervals = Vec::with_capacity(chunk.len()); + let mut r_point_counts = Vec::with_capacity(chunk.len()); + let mut r_start_lats = Vec::with_capacity(chunk.len()); + let mut r_start_lons = Vec::with_capacity(chunk.len()); + let mut r_end_lats = Vec::with_capacity(chunk.len()); + let mut r_end_lons = Vec::with_capacity(chunk.len()); + let mut r_min_lats = Vec::with_capacity(chunk.len()); + let mut r_max_lats = Vec::with_capacity(chunk.len()); + let mut r_min_lons = Vec::with_capacity(chunk.len()); + let mut r_max_lons = Vec::with_capacity(chunk.len()); + let mut r_waypoints = Vec::with_capacity(chunk.len()); + + for route in chunk { + r_ids.push(route.id.as_str()); + r_highways.push(route.highway.as_str()); + r_dir_codes.push(route.direction_code.as_str()); + r_dir_labels.push(route.direction_label.as_str()); + r_display_names.push(route.display_name.as_str()); + r_corridor_ids.push(route.corridor_id); + r_variant_ranks.push(route.variant_rank); + r_distances.push(route.distance_m); + r_durations.push(route.duration_s); + r_intervals.push(route.interval_s); + r_point_counts.push(route.point_count); + r_start_lats.push(route.start_lat); + r_start_lons.push(route.start_lon); + r_end_lats.push(route.end_lat); + r_end_lons.push(route.end_lon); + r_min_lats.push(route.min_lat); + r_max_lats.push(route.max_lat); + r_min_lons.push(route.min_lon); + r_max_lons.push(route.max_lon); + r_waypoints.push(route.waypoints_json.as_str()); + } + sqlx::query( "INSERT INTO reference_routes \ (id, highway, direction_code, direction_label, display_name, corridor_id, variant_rank, \ distance_m, duration_s, interval_s, point_count, \ start_lat, start_lon, end_lat, end_lon, \ min_lat, max_lat, min_lon, max_lon, waypoints_json) \ - VALUES \ - ($1::uuid, $2, $3, $4, $5, $6, $7, \ - $8, $9, $10, $11, \ - $12, $13, $14, $15, \ - $16, $17, $18, $19, $20)", + SELECT id::uuid, highway, direction_code, direction_label, display_name, \ + corridor_id, variant_rank, distance_m, duration_s, interval_s, point_count, \ + start_lat, start_lon, end_lat, end_lon, \ + min_lat, max_lat, min_lon, max_lon, waypoints_json \ + FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[], $5::text[], \ + $6::int[], $7::int[], $8::float8[], $9::float8[], $10::float8[], $11::int[], \ + $12::float8[], $13::float8[], $14::float8[], $15::float8[], \ + $16::float8[], $17::float8[], $18::float8[], $19::float8[], $20::text[]) \ + AS t(id, highway, direction_code, direction_label, display_name, \ + corridor_id, variant_rank, distance_m, duration_s, interval_s, point_count, \ + start_lat, start_lon, end_lat, end_lon, \ + min_lat, max_lat, min_lon, max_lon, waypoints_json)", ) - .bind(&route.id) - .bind(&route.highway) - .bind(&route.direction_code) - .bind(&route.direction_label) - .bind(&route.display_name) - .bind(route.corridor_id) - .bind(route.variant_rank) - .bind(route.distance_m) - .bind(route.duration_s) - .bind(route.interval_s) - .bind(route.point_count) - .bind(route.start_lat) - .bind(route.start_lon) - .bind(route.end_lat) - .bind(route.end_lon) - .bind(route.min_lat) - .bind(route.max_lat) - .bind(route.min_lon) - .bind(route.max_lon) - .bind(&route.waypoints_json) + .bind(&r_ids) + .bind(&r_highways) + .bind(&r_dir_codes) + .bind(&r_dir_labels) + .bind(&r_display_names) + .bind(&r_corridor_ids) + .bind(&r_variant_ranks) + .bind(&r_distances) + .bind(&r_durations) + .bind(&r_intervals) + .bind(&r_point_counts) + .bind(&r_start_lats) + .bind(&r_start_lons) + .bind(&r_end_lats) + .bind(&r_end_lons) + .bind(&r_min_lats) + .bind(&r_max_lats) + .bind(&r_min_lons) + .bind(&r_max_lons) + .bind(&r_waypoints) .execute(&mut *tx) .await?; } - for anchor in &anchors { + // Batch insert anchors using UNNEST (critical — can be 100k+ rows) + const ANCHOR_BATCH: usize = 10_000; + for chunk in anchors.chunks(ANCHOR_BATCH) { + let mut a_route_ids = Vec::with_capacity(chunk.len()); + let mut a_indices = Vec::with_capacity(chunk.len()); + let mut a_lats = Vec::with_capacity(chunk.len()); + let mut a_lons = Vec::with_capacity(chunk.len()); + + for anchor in chunk { + a_route_ids.push(anchor.route_id.as_str()); + a_indices.push(anchor.anchor_index); + a_lats.push(anchor.lat); + a_lons.push(anchor.lon); + } + sqlx::query( "INSERT INTO reference_route_anchors (route_id, anchor_index, lat, lon) \ - VALUES ($1::uuid, $2, $3, $4)", + SELECT route_id::uuid, anchor_index, lat, lon \ + FROM UNNEST($1::text[], $2::int[], $3::float8[], $4::float8[]) \ + AS t(route_id, anchor_index, lat, lon)", ) - .bind(&anchor.route_id) - .bind(anchor.anchor_index) - .bind(anchor.lat) - .bind(anchor.lon) + .bind(&a_route_ids) + .bind(&a_indices) + .bind(&a_lats) + .bind(&a_lons) .execute(&mut *tx) .await?; } diff --git a/schema/derive.sql b/schema/derive.sql index 4be2b01..5e4e4e9 100644 --- a/schema/derive.sql +++ b/schema/derive.sql @@ -220,6 +220,12 @@ TRUNCATE reference_route_anchors, reference_routes; +-- Drop GiST indexes before bulk loading exits/pois so inserts don't pay +-- per-row index maintenance. Recreated before the spatial join below. +DROP INDEX IF EXISTS exits_geom_idx; +DROP INDEX IF EXISTS pois_geom_idx; +DROP INDEX IF EXISTS exit_poi_candidates_exit_idx; + WITH rest_anchor_exits AS ( SELECT DISTINCT en.node_id FROM osm2pgsql_v2_exits_nodes en @@ -335,6 +341,16 @@ FROM ( ) src; +-- Rebuild GiST indexes after bulk load (single-pass build is faster than +-- incremental maintenance during inserts). +CREATE INDEX exits_geom_idx ON exits USING GIST (geom); +CREATE INDEX pois_geom_idx ON pois USING GIST (geom); + +-- Refresh planner statistics after bulk-loading exits and pois so the +-- spatial join below uses accurate row estimates and index selectivity. +ANALYZE exits; +ANALYZE pois; + INSERT INTO exit_poi_candidates (exit_id, poi_id, category, distance_m, rank) SELECT exit_id, poi_id, category, distance_m, rank FROM ( @@ -428,6 +444,10 @@ USING mismatched m WHERE c.exit_id = m.exit_id AND c.poi_id = m.poi_id; +-- Rebuild exit_poi_candidates index after bulk load + directional pruning. +CREATE INDEX exit_poi_candidates_exit_idx + ON exit_poi_candidates (exit_id, category, rank); + -- highway_edges, exit_corridors, corridors, and corridor_exits are rebuilt -- by openinterstate-derive steps that run after this SQL in the pipeline: -- 1. --build-graph-only -> highway_edges + exit_corridors