Generate H3 / Geohash Cells for Fast Geo Filtering

				
					#!/usr/bin/env python3

"""
radius_demo.py

End-to-end radius demo (Part 2 Step 3 + Putting It All Together):

  1) Create geo_points if needed
  2) Generate random points around a center
  3) Compute each point's H3 cell and insert
  4) Compute :cells for a radius search (grid_disk for h3 v4+, k_ring for older)
  5) Print a nicely formatted SQL ARRAY[...] you can paste into YSQL
  6) Compute bounding box for a chosen radius
  7) Run:
       - coarse cell filter count
       - full radius query count (cells + bbox + exact distance)

Requirements:
  pip install h3 psycopg2-binary
"""

import math
import random
import h3
import psycopg2


# ----------------------------
# Tunables
# ----------------------------
CENTER_LAT = 40.4406
CENTER_LON = -79.9959

H3_RES = 9
K_RING = 2

NUM_POINTS = 2000
SPREAD_DEG = 0.10  # synthetic spread (degrees)

RADIUS_METERS = 2000  # used for bounding box + exact distance

DB_HOST = "127.0.0.1"
DB_PORT = 5433
DB_NAME = "yugabyte"
DB_USER = "yugabyte"
DB_PASSWORD = None  # set if needed


# ----------------------------
# Helpers
# ----------------------------
def rand_point(center_lat: float, center_lon: float, spread_deg: float):
    return (
        center_lat + (random.random() - 0.5) * spread_deg,
        center_lon + (random.random() - 0.5) * spread_deg,
    )


def connect():
    kwargs = dict(
        host=DB_HOST,
        port=DB_PORT,
        dbname=DB_NAME,
        user=DB_USER,
    )
    if DB_PASSWORD:
        kwargs["password"] = DB_PASSWORD
    return psycopg2.connect(**kwargs)


def h3_disk(cell: str, k: int):
    """
    Return the set of cells within k steps of `cell`.

    h3-py v4+:  grid_disk(cell, k)
    older:      k_ring(cell, k)
    """
    if hasattr(h3, "grid_disk"):
        return h3.grid_disk(cell, k)
    if hasattr(h3, "k_ring"):
        return h3.k_ring(cell, k)
    raise RuntimeError("Unsupported h3 library: missing grid_disk/k_ring APIs")


def pretty_sql_array(items, per_line=3):
    """
    Produce a readable SQL array literal:

    ARRAY[
      'a','b','c',
      'd','e','f'
    ]
    """
    escaped = [s.replace("'", "''") for s in items]
    lines = []
    for i in range(0, len(escaped), per_line):
        chunk = escaped[i:i + per_line]
        suffix = "," if i + per_line < len(escaped) else ""
        lines.append("    " + ",".join(f"'{x}'" for x in chunk) + suffix)
    return "ARRAY[\n" + "\n".join(lines) + "\n  ]"


def bounding_box(center_lat, center_lon, radius_m):
    """
    Approximate bounding box for a radius in meters.
    Good enough for coarse filtering.

    lat delta: meters / 111_320
    lon delta: meters / (111_320 * cos(lat))
    """
    lat_rad = math.radians(center_lat)
    dlat = radius_m / 111_320.0
    dlon = radius_m / (111_320.0 * max(math.cos(lat_rad), 1e-12))
    return (
        center_lat - dlat,
        center_lat + dlat,
        center_lon - dlon,
        center_lon + dlon,
    )


# ----------------------------
# Main
# ----------------------------
def main():
    # Compute radius cells (these become :cells in Part 1)
    center_cell = h3.latlng_to_cell(CENTER_LAT, CENTER_LON, H3_RES)
    cells = sorted(h3_disk(center_cell, K_RING))

    # Compute bounding box for the chosen radius
    min_lat, max_lat, min_lon, max_lon = bounding_box(CENTER_LAT, CENTER_LON, RADIUS_METERS)

    print("Radius cell expansion (H3):")
    print(f"  center lat/lon          : {CENTER_LAT}, {CENTER_LON}")
    print(f"  H3 resolution           : {H3_RES}")
    print(f"  k (ring distance)       : {K_RING}")
    print(f"  center cell             : {center_cell}")
    print(f"  number of cells         : {len(cells)}")
    print()

    print("SQL-ready array (paste into YSQL):")
    print(pretty_sql_array(cells))
    print()

    print("Bounding box (approx):")
    print(f"  radius_meters           : {RADIUS_METERS}")
    print(f"  lat BETWEEN             : {min_lat:.6f} AND {max_lat:.6f}")
    print(f"  lon BETWEEN             : {min_lon:.6f} AND {max_lon:.6f}")
    print()

    # Connect to YugabyteDB and load demo data
    conn = connect()
    conn.autocommit = True

    with conn.cursor() as cur:
        # Create table + index (idempotent)
        cur.execute("""
          CREATE TABLE IF NOT EXISTS geo_points (
            id   BIGSERIAL PRIMARY KEY,
            lat  DOUBLE PRECISION NOT NULL,
            lon  DOUBLE PRECISION NOT NULL,
            cell TEXT NOT NULL
          );
        """)
        cur.execute("CREATE INDEX IF NOT EXISTS geo_points_cell_idx ON geo_points (cell);")

        # Load synthetic points
        cur.execute("TRUNCATE TABLE geo_points;")

        rows = []
        for _ in range(NUM_POINTS):
            lat, lon = rand_point(CENTER_LAT, CENTER_LON, SPREAD_DEG)
            cell = h3.latlng_to_cell(lat, lon, H3_RES)
            rows.append((lat, lon, cell))

        cur.executemany(
            "INSERT INTO geo_points (lat, lon, cell) VALUES (%s, %s, %s)",
            rows
        )

        print(f"Loaded {NUM_POINTS} points into geo_points (spread_deg={SPREAD_DEG})")
        print()

        # Stage 1: coarse cell filter only
        cur.execute(
            "SELECT count(*) FROM geo_points WHERE cell = ANY(%s);",
            (cells,)
        )
        coarse_count = cur.fetchone()[0]

        # Stage 2+3: full radius query (cells + bbox + exact distance)
        cur.execute(
            """
            SELECT count(*)
            FROM geo_points p
            WHERE p.cell = ANY(%s)
              AND p.lat BETWEEN %s AND %s
              AND p.lon BETWEEN %s AND %s
              AND (
                6371000 * acos(
                  cos(radians(%s)) * cos(radians(p.lat)) *
                  cos(radians(p.lon) - radians(%s)) +
                  sin(radians(%s)) * sin(radians(p.lat))
                )
              ) <= %s;
            """,
            (
                cells,
                min_lat, max_lat,
                min_lon, max_lon,
                CENTER_LAT, CENTER_LON, CENTER_LAT,
                RADIUS_METERS,
            )
        )
        full_count = cur.fetchone()[0]

        print("Validation counts:")
        print(f"  coarse cell filter      : {coarse_count} / {NUM_POINTS}")
        print(f"  full radius query       : {full_count} / {NUM_POINTS}")
        print()

    conn.close()


if __name__ == "__main__":
    main()