sample_mosaic_playground.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***Note:*** If using Spark < 3.4 (current version of Mosaic is using Spark 3.2.1), version of Pandas needs to be downgraded to avoid `'DataFrame' object has no attribute 'iteritems'` error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -qU pandas==1.5.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from importlib.metadata import version\n",
    "\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql import functions as F\n",
    "\n",
    "import mosaic as mos\n",
    "import mosaic.api as api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_up(use_jts=False):\n",
    "\n",
    "    library_location = f\"{mos.__path__[0]}/lib/mosaic-{version('databricks-mosaic')}-jar-with-dependencies.jar\"\n",
    "    if not os.path.exists(library_location):\n",
    "        library_location = f\"{mos.__path__[0]}/lib/mosaic-{version('databricks-mosaic')}-SNAPSHOT-jar-with-dependencies.jar\"\n",
    "\n",
    "    spark = (\n",
    "        SparkSession.builder.master(\"local\")\n",
    "        .config(\"spark.driver.host\", \"localhost\")\n",
    "        .config(\"spark.jars\", library_location)\n",
    "        .getOrCreate()\n",
    "    )\n",
    "    spark.conf.set(\"spark.databricks.labs.mosaic.jar.autoattach\", \"false\")\n",
    "    spark.sparkContext.setLogLevel(\"warn\")\n",
    "    if use_jts:\n",
    "        spark.conf.set(\"spark.databricks.labs.mosaic.geometry.api\", \"JTS\")\n",
    "    api.enable_mosaic(spark)\n",
    "\n",
    "    return spark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark = set_up()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create some dataframe here\n",
    "df = (\n",
    "    spark.createDataFrame([{'line_segment': 'LINESTRING (-83.488541 42.220045, -83.4924813 42.2198315)'}])\n",
    "    .withColumn('h3_index', mos.grid_geometrykringexplode('line_segment', F.lit(13), F.lit(2)))\n",
    ")\n",
    "df.printSchema()\n",
    "df.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***Note:*** `mosaic_kepler` magic requires `dbutis` to show the Kepler viewer in the same cell, otherwise it will output the results as HTML string. A small temporary modification to the `python/mosaic/utils/notebook_utils.py` file will write that HTML output to a file which can be displayed in IFrame, as shown the cell below next. Add the following code to that file just above the `print(html)`:\n",
    "```\n",
    "with open('./mosaic_kepler_view.html', 'w') as f:\n",
    "    f.write(html)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%mosaic_kepler\n",
    "df h3_index h3 100_000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import IFrame\n",
    "\n",
    "IFrame(src='./mosaic_kepler_view.html', width=1000, height=600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}