H3_Pandas_UDF.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Notebook to demo H3 aggregation\n",
    "\n",
    "```\n",
    "conda install -c conda-forge h3-py\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spark_esri import spark_start, spark_stop\n",
    "\n",
    "import h3\n",
    "import pandas as pd\n",
    "from pyspark.sql.functions import pandas_udf, avg\n",
    "from pyspark.sql.types import StringType"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark_stop()\n",
    "\n",
    "config = {\n",
    "    \"spark.driver.memory\": \"16G\",\n",
    "    \"spark.executor.memory\": \"16G\"\n",
    "}\n",
    "spark = spark_start(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pandas_udf(returnType=StringType())\n",
    "def geo_to_h3(lat: pd.Series, lon: pd.Series) -> pd.Series:\n",
    "    return pd.Series([h3.geo_to_h3(l1, l2, 9) for l1, l2 in zip(lat, lon)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------+-------------------+-------------------+\n",
      "|             h3|              a_lon|              a_lat|\n",
      "+---------------+-------------------+-------------------+\n",
      "|89cf15c4c3bffff|  -74.6647747954433|-52.693300866272196|\n",
      "|89e174d50a3ffff| 58.531446088677086| -67.82375331375933|\n",
      "|895b73d1a23ffff| 171.00574881945653|  7.919585098570835|\n",
      "|893ce2043dbffff|  90.66817709650394| 24.984775028380795|\n",
      "|8904057ab6bffff| 146.73747622017055|  75.67234637381009|\n",
      "|89c56d8308fffff|-23.063688559739205| -34.22443333277086|\n",
      "|8902ce9034fffff|-103.59266598948172|   70.4082288178719|\n",
      "|89f29533553ffff|-161.66214924158498| -88.38420002814281|\n",
      "|89ef3518297ffff| -48.33792136010945| -86.00040764997252|\n",
      "|89de4661e53ffff| -57.48443605445647| -49.91041810329152|\n",
      "|89ae204d427ffff|  93.12491177532615|-23.471222013891364|\n",
      "|89af850c647ffff|  87.27259031136703|-27.165742048820135|\n",
      "|8993922804fffff| -113.6801979879013| -6.730946183111428|\n",
      "|89e5456e52bffff| 103.35911593001634|  -69.4653749367924|\n",
      "|89df8672e8fffff| -42.25594463252801| -61.96962470641755|\n",
      "|89e82b74c4fffff| -91.21301482492153|-61.532931734530266|\n",
      "|8992c0b8557ffff|-106.13517860454974|-18.158913212565523|\n",
      "|89f243050b7ffff|-114.52691762732421| -73.52193205422306|\n",
      "|89a15645a0bffff|-137.54576487790925|-23.326674395062113|\n",
      "|896ad326aafffff| 28.469342724738596|-2.1941499353932414|\n",
      "+---------------+-------------------+-------------------+\n",
      "only showing top 20 rows\n",
      "\n",
      "Wall time: 10.6 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "spark\\\n",
    "    .range(1_000_000)\\\n",
    "    .selectExpr(\"id\", \"rand()*360D-180D lon\", \"rand()*180D-90D lat\")\\\n",
    "    .withColumn(\"h3\", geo_to_h3(\"lat\", \"lon\"))\\\n",
    "    .groupby(\"h3\")\\\n",
    "    .agg(avg(\"lon\").alias(\"a_lon\"), avg(\"lat\").alias(\"a_lat\"))\\\n",
    "    .show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ArcGISPro",
   "language": "Python",
   "name": "python3"
  },
  "language_info": {
   "file_extension": ".py",
   "name": "python",
   "version": "3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}