Modernizing Spatial Workflows

# ═══════════════════════════════════════════════════════════════════════════════
# BENCHMARK 01: Legacy R Stack (sf + CSV/Shapefile)
# ═══════════════════════════════════════════════════════════════════════════════
# FIX (vs previous version):
#
#   force(expr) promise-caching bug:
#     In R, function arguments are lazily evaluated promises whose result is
#     CACHED after the first evaluation. The previous measure_rss() used
#     force(expr) inside a loop, which meant the expression only actually ran
#     on iteration 1 — subsequent iterations timed the cost of returning a
#     cached NULL value (~0 µs). For ITERATIONS > 1 this produced completely
#     wrong median timings.
#
#     Fix: measure_rss() now accepts a function (fn) instead of an expression,
#     and calls fn() on each iteration. Functions are re-invoked every call.
#     Call sites change from measure_rss({ ... }) to measure_rss(function() { ... }).
#
# Author: Pukar Bhandari
# ═══════════════════════════════════════════════════════════════════════════════

library(sf)
library(dplyr)
library(readr)
library(ps)

results_csv <- "_results/benchmark_results.csv"

ITERATIONS <- as.integer(
  Sys.getenv(
    "BENCHMARK_ITERATIONS",
    unset = ifelse(exists("ITERATIONS"), as.character(ITERATIONS), "1")
  )
)

# Helper: accepts a FUNCTION (not a bare expression) and calls it each iteration.
# Using fn() instead of force(expr) ensures the body re-executes every time.
measure_rss <- function(fn, iterations = 1L) {
  times <- numeric(iterations)
  mem_deltas <- numeric(iterations)

  for (i in seq_len(iterations)) {
    gc(verbose = FALSE)
    mem_before <- ps::ps_memory_info()[["rss"]]
    t_start <- proc.time()[["elapsed"]]

    fn() # Re-invoked on every iteration — no caching

    t_end <- proc.time()[["elapsed"]]
    times[i] <- t_end - t_start
    mem_deltas[i] <- (ps::ps_memory_info()[["rss"]] - mem_before) / 1024 / 1024
  }

  list(time_sec = median(times), mem_mb = max(mem_deltas))
}

log_result <- function(sys, lang, op, time_sec, mem_mb) {
  cat(
    sprintf("%s,%s,%s,%.4f,%.2f\n", sys, lang, op, time_sec, mem_mb),
    file = results_csv,
    append = TRUE
  )
}

addr_csv <- "_data/utah_addresses.csv"
county_shp <- "_data/utah_counties.shp"

message("────────────────────────────────────────────────────────────")
message("Running: BENCHMARK 01 - R (sf) Legacy Stack")
message("────────────────────────────────────────────────────────────")

# ───────────────────────────────────────────────────────────────────────────────
## 2. I/O BENCHMARK (Data Load) ----
# ───────────────────────────────────────────────────────────────────────────────

res_io <- measure_rss(
  function() {
    pts <- read_csv(
      addr_csv,
      col_types = cols(CountyID = col_character()),
      show_col_types = FALSE
    ) |>
      filter(!is.na(x), !is.na(y)) |> # drop UGRC points with no geometry
      st_as_sf(coords = c("x", "y"), crs = 4326)

    cnt <- st_read(county_shp, quiet = TRUE) |>
      st_transform(4326)
  },
  iterations = ITERATIONS
)

log_result("sf", "R", "Load Data", res_io$time_sec, res_io$mem_mb)
message(sprintf(
  "  Load Data: %.3fs | %.1f MB RSS delta",
  res_io$time_sec,
  res_io$mem_mb
))

# ───────────────────────────────────────────────────────────────────────────────
## 3. RELOAD DATA FOR ANALYSIS PHASE (not timed) ----
# ───────────────────────────────────────────────────────────────────────────────

pts <- read_csv(
  addr_csv,
  col_types = cols(CountyID = col_character()),
  show_col_types = FALSE
) |>
  filter(!is.na(x), !is.na(y)) |> # drop UGRC points with no geometry
  st_as_sf(coords = c("x", "y"), crs = 4326)

cnt <- st_read(county_shp, quiet = TRUE) |>
  st_transform(4326)

# ───────────────────────────────────────────────────────────────────────────────
## 4. ANALYSIS BENCHMARK (Spatial Join) ----
# ───────────────────────────────────────────────────────────────────────────────

res_analysis <- measure_rss(
  function() {
    res_attr <- pts |>
      st_drop_geometry() |>
      count(CountyID)

    res_spatial <- st_join(pts, cnt, join = st_within) |>
      st_drop_geometry() |>
      count(FIPS_ST)
  },
  iterations = ITERATIONS
)

log_result(
  "sf",
  "R",
  "Validation (Join)",
  res_analysis$time_sec,
  res_analysis$mem_mb
)
message(sprintf(
  "  Validation (Join): %.3fs | %.1f MB RSS delta",
  res_analysis$time_sec,
  res_analysis$mem_mb
))

# ───────────────────────────────────────────────────────────────────────────────
## 5. CLEANUP ----
# ───────────────────────────────────────────────────────────────────────────────

rm(pts, cnt, res_attr, res_spatial)

message("────────────────────────────────────────────────────────────")
message("BENCHMARK 01 Complete!")
message("────────────────────────────────────────────────────────────")

# ═══════════════════════════════════════════════════════════════════════════════
# BENCHMARK 02: Modern R Stack (DuckDB + Parquet + Spatial Index)
# ═══════════════════════════════════════════════════════════════════════════════
# FIX (vs previous version):
#
#   force(expr) promise-caching bug (same as 01_bench_sf.R):
#     measure_rss() now accepts a function and calls fn() each iteration
#     so the body re-executes every time. Call sites use function() { ... }.
#
#   Removed unused library(bench):
#     bench::mark() was replaced by measure_rss() but the import was left in.
#
# Previous fixes retained:
#   - ST_Intersects + polygon-first join order (RTREE acceleration)
#   - LOAD spatial only (INSTALL spatial is in 00_setup_data.R)
#   - RSS-based memory measurement
#   - Split Load Data and Build Index into separate timed phases
#
# Author: Pukar Bhandari
# ═══════════════════════════════════════════════════════════════════════════════

library(duckdb)
library(dplyr)
library(DBI)
library(readr)
library(ps)

results_csv <- "_results/benchmark_results.csv"

ITERATIONS <- as.integer(
  Sys.getenv(
    "BENCHMARK_ITERATIONS",
    unset = ifelse(exists("ITERATIONS"), as.character(ITERATIONS), "1")
  )
)

# Helper: accepts a FUNCTION, calls fn() on each iteration (no promise caching).
measure_rss <- function(fn, iterations = 1L) {
  times <- numeric(iterations)
  mem_deltas <- numeric(iterations)

  for (i in seq_len(iterations)) {
    gc(verbose = FALSE)
    mem_before <- ps::ps_memory_info()[["rss"]]
    t_start <- proc.time()[["elapsed"]]

    fn()

    t_end <- proc.time()[["elapsed"]]
    times[i] <- t_end - t_start
    mem_deltas[i] <- (ps::ps_memory_info()[["rss"]] - mem_before) / 1024 / 1024
  }

  list(time_sec = median(times), mem_mb = max(mem_deltas))
}

log_result <- function(sys, lang, op, time_sec, mem_mb) {
  cat(
    sprintf("%s,%s,%s,%.4f,%.2f\n", sys, lang, op, time_sec, mem_mb),
    file = results_csv,
    append = TRUE
  )
}

addr_pq <- "_data/utah_addresses.parquet"
county_pq <- "_data/utah_counties.parquet"

message("────────────────────────────────────────────────────────────")
message("Running: BENCHMARK 02 - R (DuckDB) Modern Stack")
message("────────────────────────────────────────────────────────────")

# ───────────────────────────────────────────────────────────────────────────────
## 2. I/O BENCHMARK A — Ingest (Data Load Only, No Indexing) ----
# ───────────────────────────────────────────────────────────────────────────────
#
# Measures: connect + LOAD spatial + CREATE TABLE from Parquet (WKB conversion).
# No RTREE index creation — that is a separate phase below.
# Fair comparison with sf/GeoPandas Load Data (pure I/O, no indexing).
# ───────────────────────────────────────────────────────────────────────────────

message("Phase 2a: Benchmarking data ingest (no indexing)...")

res_ingest <- measure_rss(
  function() {
    con_i <- dbConnect(duckdb(shared_home = TRUE))
    dbExecute(con_i, "LOAD spatial;")

    dbExecute(
      con_i,
      sprintf(
        "CREATE TABLE p AS
     SELECT * EXCLUDE(geometry), ST_GeomFromWKB(geometry) AS geom
     FROM '%s'",
        addr_pq
      )
    )
    dbExecute(
      con_i,
      sprintf(
        "CREATE TABLE c AS
     SELECT * EXCLUDE(geometry), ST_GeomFromWKB(geometry) AS geom
     FROM '%s'",
        county_pq
      )
    )
    dbDisconnect(con_i)
  },
  iterations = ITERATIONS
)

log_result("duckdb", "R", "Load Data", res_ingest$time_sec, res_ingest$mem_mb)
message(sprintf(
  "  Load Data: %.3fs | %.1f MB RSS delta",
  res_ingest$time_sec,
  res_ingest$mem_mb
))

# ───────────────────────────────────────────────────────────────────────────────
## 3. I/O BENCHMARK B — Build Index ----
# ───────────────────────────────────────────────────────────────────────────────
#
# Separated from Load Data to show the one-time RTREE index creation cost
# explicitly. This cost is amortized across all subsequent queries.
# ───────────────────────────────────────────────────────────────────────────────

message("Phase 2b: Benchmarking RTREE index creation...")

# Persistent connection for analysis phase — tables stay in memory
con <- dbConnect(duckdb(shared_home = TRUE))
dbExecute(con, "LOAD spatial;")
dbExecute(
  con,
  sprintf(
    "CREATE TABLE p AS
   SELECT * EXCLUDE(geometry), ST_GeomFromWKB(geometry) AS geom
   FROM '%s'",
    addr_pq
  )
)
dbExecute(
  con,
  sprintf(
    "CREATE TABLE c AS
   SELECT * EXCLUDE(geometry), ST_GeomFromWKB(geometry) AS geom
   FROM '%s'",
    county_pq
  )
)

res_index <- measure_rss(
  function() {
    # Drop and rebuild so each iteration is independently timed
    tryCatch(
      dbExecute(con, "DROP INDEX IF EXISTS idx_p_geom;"),
      error = function(e) NULL
    )
    tryCatch(
      dbExecute(con, "DROP INDEX IF EXISTS idx_c_geom;"),
      error = function(e) NULL
    )
    dbExecute(con, "CREATE INDEX idx_p_geom ON p USING RTREE (geom);")
    dbExecute(con, "CREATE INDEX idx_c_geom ON c USING RTREE (geom);")
  },
  iterations = ITERATIONS
)

log_result("duckdb", "R", "Build Index", res_index$time_sec, res_index$mem_mb)
message(sprintf(
  "  Build Index: %.3fs | %.1f MB RSS delta",
  res_index$time_sec,
  res_index$mem_mb
))

# ───────────────────────────────────────────────────────────────────────────────
## 4. WARMUP QUERY ----
# ───────────────────────────────────────────────────────────────────────────────

message("Running warmup query (not timed)...")

warmup_sql <- "
  SELECT c.FIPS_STR, COUNT(a.OBJECTID) AS point_count
  FROM c
  INNER JOIN p a ON ST_Intersects(a.geom, c.geom)
  GROUP BY c.FIPS_STR
"
invisible(dbGetQuery(con, warmup_sql))
message("Warmup complete. Starting timed benchmark...")

# ───────────────────────────────────────────────────────────────────────────────
## 5. ANALYSIS BENCHMARK (Spatial Join) ----
# ───────────────────────────────────────────────────────────────────────────────
#
# ST_Intersects (not ST_Within) — DuckDB's RTREE accelerates ST_Intersects.
# ST_Within in a JOIN falls back to BLOCKWISE_NL_JOIN (no index).
# For point geometries, ST_Intersects == ST_Within (points can't straddle a boundary).
#
# FROM c INNER JOIN p — polygon-first order (29 rows) lets DuckDB iterate
# polygons and use the RTREE on 1.2M points to retrieve candidates per polygon.
# ───────────────────────────────────────────────────────────────────────────────

sql <- "
  SELECT c.FIPS_STR, COUNT(a.OBJECTID) AS point_count
  FROM c
  INNER JOIN p a ON ST_Intersects(a.geom, c.geom)
  GROUP BY c.FIPS_STR
"

res_join <- measure_rss(
  function() {
    res <- dbGetQuery(con, sql)
  },
  iterations = ITERATIONS
)

log_result(
  "duckdb",
  "R",
  "Validation (Join)",
  res_join$time_sec,
  res_join$mem_mb
)
message(sprintf(
  "  Validation (Join): %.3fs | %.1f MB RSS delta",
  res_join$time_sec,
  res_join$mem_mb
))

# ───────────────────────────────────────────────────────────────────────────────
## 6. VALIDATION & EXPORT ----
# ───────────────────────────────────────────────────────────────────────────────

message("Generating validation statistics...")

sql_diff <- "
WITH attr_counts AS (
    SELECT CountyID, COUNT(*) AS count_attr
    FROM p
    GROUP BY CountyID
),
spatial_counts AS (
    SELECT c.FIPS_STR, COUNT(a.OBJECTID) AS count_spatial
    FROM c
    INNER JOIN p a ON ST_Intersects(a.geom, c.geom)
    GROUP BY c.FIPS_STR
)
SELECT
    COALESCE(a.CountyID, s.FIPS_STR)     AS CountyID,
    COALESCE(a.count_attr,    0)          AS count_attr,
    COALESCE(s.count_spatial, 0)          AS count_spatial,
    (COALESCE(a.count_attr, 0) - COALESCE(s.count_spatial, 0)) AS diff
FROM attr_counts a
FULL OUTER JOIN spatial_counts s ON a.CountyID = s.FIPS_STR
ORDER BY ABS(diff) DESC
LIMIT 10;
"

stats <- dbGetQuery(con, sql_diff)
write_csv(stats, "_results/validation_stats.csv")
message("Validation stats saved to: _results/validation_stats.csv")

# ───────────────────────────────────────────────────────────────────────────────
## 7. CLEANUP ----
# ───────────────────────────────────────────────────────────────────────────────

dbDisconnect(con)

message("────────────────────────────────────────────────────────────")
message("BENCHMARK 02 Complete!")
message("────────────────────────────────────────────────────────────")

# ═══════════════════════════════════════════════════════════════════════════════
# BENCHMARK 03: Legacy Python Stack (GeoPandas + CSV/Shapefile)
# ═══════════════════════════════════════════════════════════════════════════════
# Purpose: Establish baseline performance using traditional Python geospatial
#          tools and file formats (CSV for points, Shapefile for polygons)
#
# FIX NOTES (vs. original):
#
#   FIX 1 — Honest memory measurement via psutil RSS:
#     tracemalloc only tracks Python-managed allocations. GeoPandas uses
#     shapely geometries backed by GEOS (C library), which allocates in a
#     separate arena that tracemalloc cannot see. For 1.2M address points,
#     GEOS geometry objects can consume 300–500 MB of C-heap memory that
#     tracemalloc reports as zero. psutil.Process().memory_info().rss
#     captures the true process resident set size.
#
# Author: Pukar Bhandari
# ═══════════════════════════════════════════════════════════════════════════════

import os
import sys
import time
import statistics
import pandas as pd
import geopandas as gpd
import psutil       # FIX 1
from pathlib import Path

ITERATIONS = int(os.environ.get('BENCHMARK_ITERATIONS', '1'))

if os.name == 'nt':
    gdal_data_path = Path(sys.prefix) / "Library" / "share" / "gdal"
    if gdal_data_path.exists():
        os.environ['GDAL_DATA'] = str(gdal_data_path)

RESULTS_CSV = "_results/benchmark_results.csv"
ADDR_CSV    = "_data/utah_addresses.csv"
COUNTY_SHP  = "_data/utah_counties.shp"

def log_result(sys_name, lang, op, time_sec, peak_mem_mb):
    with open(RESULTS_CSV, "a") as f:
        f.write(f"{sys_name},{lang},{op},{time_sec:.4f},{peak_mem_mb:.2f}\n")

# FIX 1: RSS-based memory measurement
proc = psutil.Process(os.getpid())

def rss_mb():
    return proc.memory_info().rss / 1024 / 1024

print("─" * 60)
print("Running: BENCHMARK 03 - Python (GeoPandas) Legacy Stack")
print("─" * 60)

# ───────────────────────────────────────────────────────────────────────────────
## 2. I/O BENCHMARK (Data Load) ----
# ───────────────────────────────────────────────────────────────────────────────

load_times = []
load_mems  = []

for _ in range(ITERATIONS):
    mem_before = rss_mb()
    t_start    = time.perf_counter()

    df = pd.read_csv(ADDR_CSV, dtype={"CountyID": str}, low_memory=False)
    df = df.dropna(subset=["x", "y"])  # drop UGRC points with no geometry
    pts = gpd.GeoDataFrame(
        df,
        geometry=gpd.points_from_xy(df.x, df.y),
        crs="EPSG:4326"
    )
    cnt = gpd.read_file(COUNTY_SHP).to_crs("EPSG:4326")

    load_times.append(time.perf_counter() - t_start)
    load_mems.append(rss_mb() - mem_before)

log_result("geopandas", "Python", "Load Data",
           statistics.median(load_times),
           max(load_mems))
print(f"  Load Data: {statistics.median(load_times):.3f}s | {max(load_mems):.1f} MB RSS delta")

# ───────────────────────────────────────────────────────────────────────────────
## 3. ANALYSIS BENCHMARK (Spatial Join) ----
# ───────────────────────────────────────────────────────────────────────────────

# Load data once outside the timed loop (same pattern as 01_bench_sf.R)
df  = pd.read_csv(ADDR_CSV, dtype={"CountyID": str}, low_memory=False)
df = df.dropna(subset=["x", "y"])  # drop UGRC points with no geometry
pts = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df.x, df.y), crs="EPSG:4326")
cnt = gpd.read_file(COUNTY_SHP).to_crs("EPSG:4326")

join_times = []
join_mems  = []

for _ in range(ITERATIONS):
    mem_before = rss_mb()
    t_start    = time.perf_counter()

    res_attr    = pts.groupby("CountyID").size()
    joined      = gpd.sjoin(pts, cnt, how="inner", predicate="within")
    res_spatial = joined.groupby("FIPS_ST").size()

    join_times.append(time.perf_counter() - t_start)
    join_mems.append(rss_mb() - mem_before)

log_result("geopandas", "Python", "Validation (Join)",
           statistics.median(join_times),
           max(join_mems))
print(f"  Validation (Join): {statistics.median(join_times):.3f}s | {max(join_mems):.1f} MB RSS delta")

# ───────────────────────────────────────────────────────────────────────────────
## 4. CLEANUP ----
# ───────────────────────────────────────────────────────────────────────────────

del pts, cnt, joined, res_attr, res_spatial

print("─" * 60)
print("BENCHMARK 03 Complete!")
print("─" * 60)

# ═══════════════════════════════════════════════════════════════════════════════
# BENCHMARK 04: Modern Python Stack (DuckDB + Parquet + Spatial Index)
# ═══════════════════════════════════════════════════════════════════════════════
# Purpose: Demonstrate DuckDB's efficiency for large-scale spatial operations
#          using columnar storage (Parquet) and RTREE spatial indexing.
#
# FIX NOTES (vs. original):
#
#   FIX 1 — ST_Intersects instead of ST_Within in JOIN:
#     DuckDB's RTREE index only accelerates bounding-box queries. ST_Within in
#     a JOIN condition causes a BLOCKWISE_NL_JOIN (no index). ST_Intersects IS
#     RTREE-accelerated for point-polygon (a point can't straddle a boundary,
#     so ST_Intersects == ST_Within for points). Verify with EXPLAIN.
#
#   FIX 2 — Table order: FROM c INNER JOIN p (polygons as left/build side):
#     Lets DuckDB iterate 29 polygons and use the RTREE on 1.2M points to
#     retrieve candidates per polygon. Opposite order forces full point scan.
#
#   FIX 3 — install_extension called once in setup, not in benchmark loop:
#     install_extension() performs a network version-check even when cached.
#     This inflates DuckDB's "Load Data" time with overhead unrelated to I/O.
#     The spatial extension is installed once in 00_setup_data.R (R side) or
#     in a one-time setup block here.
#
#   FIX 4 — Honest memory measurement via psutil RSS:
#     tracemalloc only tracks Python-managed objects. DuckDB's C++ engine
#     allocates in its own arena, invisible to tracemalloc. psutil.Process()
#     .memory_info().rss captures the actual process resident set size.
#
#   FIX 5 — Split Load phase into Ingest + Build Index:
#     Separates the one-time indexing cost from pure I/O for a fair
#     comparison with geopandas/sf (which have no explicit index step).
#
# Author: Pukar Bhandari
# ═══════════════════════════════════════════════════════════════════════════════

# ───────────────────────────────────────────────────────────────────────────────
## 1. SETUP & CONFIGURATION ----
# ───────────────────────────────────────────────────────────────────────────────

import os
import sys
import time
import statistics
import duckdb
import psutil       # FIX 4: for honest RSS-based memory measurement
from pathlib import Path

ITERATIONS = int(os.environ.get('BENCHMARK_ITERATIONS', '1'))

RESULTS_CSV = "_results/benchmark_results.csv"
ADDR_PQ    = "_data/utah_addresses.parquet"
COUNTY_PQ  = "_data/utah_counties.parquet"

def log_result(sys_name, lang, op, time_sec, peak_mem_mb):
    with open(RESULTS_CSV, "a") as f:
        f.write(f"{sys_name},{lang},{op},{time_sec:.4f},{peak_mem_mb:.2f}\n")

# FIX 4: RSS-based memory measurement
proc = psutil.Process(os.getpid())

def rss_mb():
    return proc.memory_info().rss / 1024 / 1024

print("─" * 60)
print("Running: BENCHMARK 04 - Python (DuckDB) Modern Stack")
print("─" * 60)

# ───────────────────────────────────────────────────────────────────────────────
## 2. ONE-TIME SETUP — Install Spatial Extension (not timed) ----
# ───────────────────────────────────────────────────────────────────────────────
#
# FIX 3: Install the spatial extension ONCE outside the benchmark loop.
# Subsequent benchmark iterations use LOAD spatial (fast, no network check).
# ───────────────────────────────────────────────────────────────────────────────

print("Installing spatial extension (one-time, not timed)...")
_setup_con = duckdb.connect()
_setup_con.install_extension("spatial")
_setup_con.close()
print("Spatial extension ready.")

# ───────────────────────────────────────────────────────────────────────────────
## 3. I/O BENCHMARK A — Ingest (Data Load Only, No Indexing) ----
# ───────────────────────────────────────────────────────────────────────────────
#
# FIX 3: Only LOAD spatial (fast), no install. Measures only Parquet ingest.
# FIX 4: RSS delta instead of tracemalloc.
# FIX 5: Index creation is a separate timed phase.
# ───────────────────────────────────────────────────────────────────────────────

ingest_times = []
ingest_mems  = []

for _ in range(ITERATIONS):
    mem_before = rss_mb()
    t_start    = time.perf_counter()

    con_i = duckdb.connect()
    con_i.load_extension("spatial")   # FIX 3: LOAD only

    con_i.execute(
        f"CREATE TABLE p AS "
        f"SELECT * EXCLUDE(geometry), ST_GeomFromWKB(geometry) AS geom "
        f"FROM '{ADDR_PQ}'"
    )
    con_i.execute(
        f"CREATE TABLE c AS "
        f"SELECT * EXCLUDE(geometry), ST_GeomFromWKB(geometry) AS geom "
        f"FROM '{COUNTY_PQ}'"
    )
    con_i.close()

    ingest_times.append(time.perf_counter() - t_start)
    ingest_mems.append(rss_mb() - mem_before)

log_result("duckdb", "Python", "Load Data",
           statistics.median(ingest_times),
           max(ingest_mems))
print(f"  Load Data: {statistics.median(ingest_times):.3f}s | {max(ingest_mems):.1f} MB RSS delta")

# ───────────────────────────────────────────────────────────────────────────────
## 4. I/O BENCHMARK B — Build Index ----
# ───────────────────────────────────────────────────────────────────────────────
#
# FIX 5: Index creation measured separately as a standalone operation.
# ───────────────────────────────────────────────────────────────────────────────

# Persistent connection for analysis phase
con = duckdb.connect()
con.load_extension("spatial")
con.execute(
    f"CREATE TABLE p AS "
    f"SELECT * EXCLUDE(geometry), ST_GeomFromWKB(geometry) AS geom "
    f"FROM '{ADDR_PQ}'"
)
con.execute(
    f"CREATE TABLE c AS "
    f"SELECT * EXCLUDE(geometry), ST_GeomFromWKB(geometry) AS geom "
    f"FROM '{COUNTY_PQ}'"
)

index_times = []
index_mems  = []

for _ in range(ITERATIONS):
    # Drop and rebuild so each iteration is independent
    try:
        con.execute("DROP INDEX IF EXISTS idx_p_geom;")
        con.execute("DROP INDEX IF EXISTS idx_c_geom;")
    except Exception:
        pass

    mem_before = rss_mb()
    t_start    = time.perf_counter()

    con.execute("CREATE INDEX idx_p_geom ON p USING RTREE (geom);")
    con.execute("CREATE INDEX idx_c_geom ON c USING RTREE (geom);")

    index_times.append(time.perf_counter() - t_start)
    index_mems.append(rss_mb() - mem_before)

log_result("duckdb", "Python", "Build Index",
           statistics.median(index_times),
           max(index_mems))
print(f"  Build Index: {statistics.median(index_times):.3f}s | {max(index_mems):.1f} MB RSS delta")

# ───────────────────────────────────────────────────────────────────────────────
## 5. WARMUP QUERY ----
# ───────────────────────────────────────────────────────────────────────────────

print("Running warmup query (not timed)...")

# FIX 1 + FIX 2: ST_Intersects (RTREE-accelerated) + c as build side (29 rows)
warmup_query = """
    SELECT c.FIPS_STR, COUNT(a.OBJECTID) AS point_count
    FROM c
    INNER JOIN p a ON ST_Intersects(a.geom, c.geom)
    GROUP BY c.FIPS_STR
"""
_ = con.execute(warmup_query).fetchdf()
print("Warmup complete. Starting timed benchmark...")

# ───────────────────────────────────────────────────────────────────────────────
## 6. ANALYSIS BENCHMARK (Spatial Join) ----
# ───────────────────────────────────────────────────────────────────────────────
#
# FIX 1: ST_Intersects instead of ST_Within.
#   For point geometries, these produce identical results. But ST_Intersects
#   is RTREE-accelerated while ST_Within falls back to nested-loop scan.
#   Check the query plan with:
#       print(con.execute("EXPLAIN " + query).fetchall())
#   Look for RTREE_INDEX_SCAN (good) vs BLOCKWISE_NL_JOIN (slow).
#
# FIX 2: FROM c INNER JOIN p a — polygon-first join order.
#   With 29 polygons on the left (build) side, DuckDB iterates each polygon
#   and uses the RTREE index on p (1.2M points) to retrieve candidate points
#   within each polygon's bounding box, then does exact intersection check.
# ───────────────────────────────────────────────────────────────────────────────

query = """
    SELECT c.FIPS_STR, COUNT(a.OBJECTID) AS point_count
    FROM c
    INNER JOIN p a ON ST_Intersects(a.geom, c.geom)
    GROUP BY c.FIPS_STR
"""

join_times = []
join_mems  = []

for _ in range(ITERATIONS):
    mem_before = rss_mb()
    t_start    = time.perf_counter()

    result = con.execute(query).fetchdf()

    join_times.append(time.perf_counter() - t_start)
    join_mems.append(rss_mb() - mem_before)

log_result("duckdb", "Python", "Validation (Join)",
           statistics.median(join_times),
           max(join_mems))
print(f"  Validation (Join): {statistics.median(join_times):.3f}s | {max(join_mems):.1f} MB RSS delta")

# ───────────────────────────────────────────────────────────────────────────────
## 7. CLEANUP ----
# ───────────────────────────────────────────────────────────────────────────────

con.close()

print("─" * 60)
print("BENCHMARK 04 Complete!")
print("─" * 60)

# ═══════════════════════════════════════════════════════════════════════════════
# BENCHMARK 05: Modern R Stack (Apache SedonaDB + Parquet)
# ═══════════════════════════════════════════════════════════════════════════════
# Architecture notes:
#   - SedonaDB is an embedded Rust/DataFusion engine — no JVM, no Spark.
#   - No explicit INSTALL/LOAD step — spatial functions ship with the package.
#   - No explicit index creation — DataFusion handles spatial indexing internally.
#   - sd_read_parquet() + sd_to_view() is LAZY: no data is scanned until
#     sd_collect() fires. "Load Data" times view registration only (~0.01s).
#     Actual Parquet scan happens during "Validation (Join)".
#
# Fixes applied:
#   FIX 1 — fn() pattern in measure_rss():
#     R promises cache after first evaluation — force(expr) in a loop only runs
#     the body on iteration 1. fn() re-executes every iteration.
#
#   FIX 2 — No "SELECT * EXCLUDE":
#     DataFusion does not support DuckDB's EXCLUDE extension.
#     All columns named explicitly using schemas confirmed from Parquet files.
#
#   FIX 3 — Double-quoted identifiers throughout:
#     DataFusion is CASE-SENSITIVE. Unquoted identifiers are lowercased at parse
#     time (OBJECTID → objectid → schema error). All column references use
#     double-quotes to preserve original case.
#
#   FIX 4 — ps::ps_memory_info()[["rss"]] not $rss:
#     On Windows, ps_memory_info() returns a named numeric vector (not a list).
#     $ accessor throws "invalid for atomic vectors". [["rss"]] works on both.
#
# Author: Pukar Bhandari
# ═══════════════════════════════════════════════════════════════════════════════

library(sedonadb)
library(dplyr)
library(readr)
library(ps)

results_csv <- "_results/benchmark_results.csv"

ITERATIONS <- as.integer(
  Sys.getenv(
    "BENCHMARK_ITERATIONS",
    unset = ifelse(exists("ITERATIONS"), as.character(ITERATIONS), "1")
  )
)

# FIX 1 + FIX 4
measure_rss <- function(fn, iterations = 1L) {
  times <- numeric(iterations)
  mem_deltas <- numeric(iterations)

  for (i in seq_len(iterations)) {
    gc(verbose = FALSE)
    mem_before <- ps::ps_memory_info()[["rss"]]
    t_start <- proc.time()[["elapsed"]]

    fn()

    t_end <- proc.time()[["elapsed"]]
    times[i] <- t_end - t_start
    mem_deltas[i] <- (ps::ps_memory_info()[["rss"]] - mem_before) / 1024 / 1024
  }

  list(time_sec = median(times), mem_mb = max(mem_deltas))
}

log_result <- function(sys, lang, op, time_sec, mem_mb) {
  cat(
    sprintf("%s,%s,%s,%.4f,%.2f\n", sys, lang, op, time_sec, mem_mb),
    file = results_csv,
    append = TRUE
  )
}

addr_pq <- "_data/utah_addresses.parquet"
county_pq <- "_data/utah_counties.parquet"

message("────────────────────────────────────────────────────────────")
message("Running: BENCHMARK 05 - R (SedonaDB) Modern Stack")
message("────────────────────────────────────────────────────────────")

# ───────────────────────────────────────────────────────────────────────────────
## 2. I/O BENCHMARK — Load Data ----
# ───────────────────────────────────────────────────────────────────────────────
#
# FIX 2 + FIX 3: Explicit quoted column lists from actual Parquet schemas.
# Schemas confirmed from live Parquet files:
#
#   Address (UGRC UtahAddressPoints):
#     OBJECTID, AddSystem, UTAddPtID, FullAdd, AddNum, AddNumSuffix, PrefixDir,
#     StreetName, StreetType, SuffixDir, LandmarkName, Building, UnitType,
#     UnitID, City, ZipCode, CountyID, State, PtLocation, PtType, Structure,
#     ParcelID, AddSource, LoadDate, USNG, geometry
#
#   Counties (UGRC UtahCountyBoundaries):
#     OBJECTID, COUNTYNBR, ENTITYNBR, ENTITYYR, NAME, FIPS, STATEPLANE,
#     POP_LASTCENSUS, POP_CURRESTIMATE, GlobalID, FIPS_STR, COLOR4,
#     Shape__Area, Shape__Length, CLASS, geometry
#
# LAZY EVALUATION NOTE:
#   sd_read_parquet() + sd_to_view() registers a view definition — no data read.
#   "Load Data" timing is ~0.01s (view registration overhead only).
#   Actual Parquet scan fires during "Validation (Join)" via sd_collect().
# ───────────────────────────────────────────────────────────────────────────────

res_io <- measure_rss(
  function() {
    sd_read_parquet(addr_pq) |> sd_to_view("addr_raw", overwrite = TRUE)
    sd_read_parquet(county_pq) |> sd_to_view("county_raw", overwrite = TRUE)

    sd_sql(
      '
    SELECT
      "OBJECTID", "AddSystem", "UTAddPtID", "FullAdd", "AddNum",
      "AddNumSuffix", "PrefixDir", "StreetName", "StreetType", "SuffixDir",
      "LandmarkName", "Building", "UnitType", "UnitID", "City", "ZipCode",
      "CountyID", "State", "PtLocation", "PtType", "Structure", "ParcelID",
      "AddSource", "LoadDate", "USNG",
      ST_GeomFromWKB(geometry) AS geom
    FROM addr_raw
  '
    ) |>
      sd_to_view("p", overwrite = TRUE)

    sd_sql(
      '
    SELECT
      "OBJECTID", "COUNTYNBR", "ENTITYNBR", "ENTITYYR", "NAME", "FIPS",
      "STATEPLANE", "POP_LASTCENSUS", "POP_CURRESTIMATE", "GlobalID",
      "FIPS_STR", "COLOR4", "Shape__Area", "Shape__Length", "CLASS",
      ST_GeomFromWKB(geometry) AS geom
    FROM county_raw
  '
    ) |>
      sd_to_view("c", overwrite = TRUE)
  },
  iterations = ITERATIONS
)

log_result("sedonadb", "R", "Load Data", res_io$time_sec, res_io$mem_mb)
message(sprintf(
  "  Load Data: %.3fs | %.1f MB RSS delta",
  res_io$time_sec,
  res_io$mem_mb
))

# ───────────────────────────────────────────────────────────────────────────────
## 3. WARMUP QUERY ----
# ───────────────────────────────────────────────────────────────────────────────

message("Running warmup query (not timed)...")

invisible(
  sd_sql(
    '
    SELECT c."FIPS_STR", COUNT(*) AS point_count
    FROM c
    INNER JOIN p a ON ST_Intersects(a.geom, c.geom)
    GROUP BY c."FIPS_STR"
  '
  ) |>
    sd_collect()
)

message("Warmup complete. Starting timed benchmark...")

# ───────────────────────────────────────────────────────────────────────────────
## 4. ANALYSIS BENCHMARK — Spatial Join ----
# ───────────────────────────────────────────────────────────────────────────────
#
# COUNT(*) used instead of COUNT(a."OBJECTID") — semantically equivalent for
# counting matched rows, and avoids any residual case-sensitivity concerns.
# ───────────────────────────────────────────────────────────────────────────────

res_join <- measure_rss(
  function() {
    sd_sql(
      '
    SELECT c."FIPS_STR", COUNT(*) AS point_count
    FROM c
    INNER JOIN p a ON ST_Intersects(a.geom, c.geom)
    GROUP BY c."FIPS_STR"
  '
    ) |>
      sd_collect()
  },
  iterations = ITERATIONS
)

log_result(
  "sedonadb",
  "R",
  "Validation (Join)",
  res_join$time_sec,
  res_join$mem_mb
)
message(sprintf(
  "  Validation (Join): %.3fs | %.1f MB RSS delta",
  res_join$time_sec,
  res_join$mem_mb
))

# ───────────────────────────────────────────────────────────────────────────────
## 5. CLEANUP ----
# ───────────────────────────────────────────────────────────────────────────────

sd_drop_view("p")
sd_drop_view("c")
sd_drop_view("addr_raw")
sd_drop_view("county_raw")

message("────────────────────────────────────────────────────────────")
message("BENCHMARK 05 Complete!")
message("────────────────────────────────────────────────────────────")

# ═══════════════════════════════════════════════════════════════════════════════
# BENCHMARK 06: Modern Python Stack (Apache SedonaDB + Parquet)
# ═══════════════════════════════════════════════════════════════════════════════
# Architecture notes:
#   - SedonaDB Python is the same embedded Rust/DataFusion engine as the R
#     package — same query planner, same performance characteristics.
#   - Install: pip install "apache-sedona[db]"
#   - API: sedona.db.connect() returns a connection with a .sql() method
#     that returns a DataFrame-like result.
#   - Lazy evaluation: read_parquet() + create_view() registers a view
#     definition — no data scanned until .collect() fires.
#
# Case sensitivity note (same as R version):
#   DataFusion is CASE-SENSITIVE. Column names must be double-quoted in SQL
#   to preserve original case from the Parquet schema.
#
# Author: Pukar Bhandari
# ═══════════════════════════════════════════════════════════════════════════════

import os
import time
import statistics
import psutil
import sedona.db

ITERATIONS = int(os.environ.get("BENCHMARK_ITERATIONS", "1"))

RESULTS_CSV = "_results/benchmark_results.csv"
ADDR_PQ     = "_data/utah_addresses.parquet"
COUNTY_PQ   = "_data/utah_counties.parquet"

def log_result(sys_name, lang, op, time_sec, peak_mem_mb):
    with open(RESULTS_CSV, "a") as f:
        f.write(f"{sys_name},{lang},{op},{time_sec:.4f},{peak_mem_mb:.2f}\n")

proc = psutil.Process(os.getpid())

def rss_mb():
    return proc.memory_info().rss / 1024 / 1024

print("─" * 60)
print("Running: BENCHMARK 06 - Python (SedonaDB) Modern Stack")
print("─" * 60)

# ───────────────────────────────────────────────────────────────────────────────
## 2. CONNECT ----
# ───────────────────────────────────────────────────────────────────────────────
# Single persistent connection — SedonaDB manages its own session state.
# ───────────────────────────────────────────────────────────────────────────────

sd = sedona.db.connect()

# ───────────────────────────────────────────────────────────────────────────────
## 3. I/O BENCHMARK — Load Data ----
# ───────────────────────────────────────────────────────────────────────────────
#
# read_parquet() + create_view() is LAZY — no data scanned until .collect().
# "Load Data" timing measures view registration only (~0.01s).
# Actual Parquet scan happens during "Validation (Join)" when .collect() fires.
#
# Column schemas confirmed from live Parquet files (same as R version):
#
#   Address (UGRC UtahAddressPoints):
#     OBJECTID, AddSystem, UTAddPtID, FullAdd, AddNum, AddNumSuffix, PrefixDir,
#     StreetName, StreetType, SuffixDir, LandmarkName, Building, UnitType,
#     UnitID, City, ZipCode, CountyID, State, PtLocation, PtType, Structure,
#     ParcelID, AddSource, LoadDate, USNG, geometry
#
#   Counties (UGRC UtahCountyBoundaries):
#     OBJECTID, COUNTYNBR, ENTITYNBR, ENTITYYR, NAME, FIPS, STATEPLANE,
#     POP_LASTCENSUS, POP_CURRESTIMATE, GlobalID, FIPS_STR, COLOR4,
#     Shape__Area, Shape__Length, CLASS, geometry
# ───────────────────────────────────────────────────────────────────────────────

load_times = []
load_mems  = []

for _ in range(ITERATIONS):
    mem_before = rss_mb()
    t_start    = time.perf_counter()

    # Register raw Parquet views (lazy — no I/O yet)
    sd.read_parquet(ADDR_PQ).create_view("addr_raw",   replace=True)
    sd.read_parquet(COUNTY_PQ).create_view("county_raw", replace=True)

    # Geometry-converted views — all identifiers double-quoted (case-sensitive)
    sd.sql("""
        SELECT
          "OBJECTID", "AddSystem", "UTAddPtID", "FullAdd", "AddNum",
          "AddNumSuffix", "PrefixDir", "StreetName", "StreetType", "SuffixDir",
          "LandmarkName", "Building", "UnitType", "UnitID", "City", "ZipCode",
          "CountyID", "State", "PtLocation", "PtType", "Structure", "ParcelID",
          "AddSource", "LoadDate", "USNG",
          ST_GeomFromWKB(geometry) AS geom
        FROM addr_raw
    """).create_view("p", replace=True)

    sd.sql("""
        SELECT
          "OBJECTID", "COUNTYNBR", "ENTITYNBR", "ENTITYYR", "NAME", "FIPS",
          "STATEPLANE", "POP_LASTCENSUS", "POP_CURRESTIMATE", "GlobalID",
          "FIPS_STR", "COLOR4", "Shape__Area", "Shape__Length", "CLASS",
          ST_GeomFromWKB(geometry) AS geom
        FROM county_raw
    """).create_view("c", replace=True)

    load_times.append(time.perf_counter() - t_start)
    load_mems.append(rss_mb() - mem_before)

log_result("sedonadb", "Python", "Load Data",
           statistics.median(load_times), max(load_mems))
print(f"  Load Data: {statistics.median(load_times):.3f}s | {max(load_mems):.1f} MB RSS delta")

# ───────────────────────────────────────────────────────────────────────────────
## 4. WARMUP QUERY ----
# ───────────────────────────────────────────────────────────────────────────────

print("Running warmup query (not timed)...")

sd.sql("""
    SELECT c."FIPS_STR", COUNT(*) AS point_count
    FROM c
    INNER JOIN p a ON ST_Intersects(a.geom, c.geom)
    GROUP BY c."FIPS_STR"
""").collect()

print("Warmup complete. Starting timed benchmark...")

# ───────────────────────────────────────────────────────────────────────────────
## 5. ANALYSIS BENCHMARK — Spatial Join ----
# ───────────────────────────────────────────────────────────────────────────────
#
# Because views are lazy, .collect() here triggers the ACTUAL Parquet scan
# in addition to the join. "Validation (Join)" includes full I/O + compute.
# ───────────────────────────────────────────────────────────────────────────────

query = """
    SELECT c."FIPS_STR", COUNT(*) AS point_count
    FROM c
    INNER JOIN p a ON ST_Intersects(a.geom, c.geom)
    GROUP BY c."FIPS_STR"
"""

join_times = []
join_mems  = []

for _ in range(ITERATIONS):
    mem_before = rss_mb()
    t_start    = time.perf_counter()

    result = sd.sql(query).collect()

    join_times.append(time.perf_counter() - t_start)
    join_mems.append(rss_mb() - mem_before)

log_result("sedonadb", "Python", "Validation (Join)",
           statistics.median(join_times), max(join_mems))
print(f"  Validation (Join): {statistics.median(join_times):.3f}s | {max(join_mems):.1f} MB RSS delta")

# ───────────────────────────────────────────────────────────────────────────────
## 6. CLEANUP ----
# ───────────────────────────────────────────────────────────────────────────────

sd.sql("DROP VIEW IF EXISTS p")
sd.sql("DROP VIEW IF EXISTS c")
sd.sql("DROP VIEW IF EXISTS addr_raw")
sd.sql("DROP VIEW IF EXISTS county_raw")

print("─" * 60)
print("BENCHMARK 06 Complete!")
print("─" * 60)