src/notebooks/benchmark_zarr_parquet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "80137bf7",
   "metadata": {},
   "source": [
    "# Benchmark Zarr vs. Parquet reach queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "20de8eb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import timeit\n",
    "\n",
    "import xarray as xr\n",
    "import fsspec\n",
    "import dask.dataframe as dd\n",
    "from scipy import stats\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dfdbb436-4881-4347-a936-dd1a086aebd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# These reach ids are within a HUC around Cobb's Creek in West Philly.\n",
    "avail_reach_ids = [4495680, 4495656, 4494480, 4489136, 4489138, 4496602]\n",
    "\n",
    "# Set these to the location of the Zarr and Parquet data samples.\n",
    "zarr_path = 's3://research-lf-dev/noaa/reformat-sample/streamflow-zarr/'\n",
    "parq_path = 's3://research-lf-dev/noaa/reformat-sample/streamflow-parquet/'\n",
    "\n",
    "# Set this to location of AWS credentials file if it's not in the default place. This is useful \n",
    "# if running on MSPC. In that case, just put this in the root directory of the JupyterHub file explorer.\n",
    "# os.environ['AWS_SHARED_CREDENTIALS_FILE'] = './config'\n",
    "\n",
    "# Number of repeats to use for benchmarking.\n",
    "repeat = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9c24ac20-c395-4fc5-86d6-04a3dfb9c415",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_zarr_streamflow():\n",
    "    ds = xr.open_zarr(fsspec.get_mapper(zarr_path, anon=False, requester_pays=True))\n",
    "    sub_ds = ds.sel(feature_id=avail_reach_ids)\n",
    "    return sub_ds.streamflow.values\n",
    "\n",
    "def get_parq_streamflow():\n",
    "    df = dd.read_parquet(parq_path)\n",
    "    sub_df = df['streamflow'].loc[avail_reach_ids].compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f861c2c6-f4bc-4009-ace7-c4430ccef565",
   "metadata": {},
   "outputs": [],
   "source": [
    "zarr_times = timeit.repeat(get_zarr_streamflow, number=1, repeat=repeat)\n",
    "parq_times = timeit.repeat(get_parq_streamflow, number=1, repeat=repeat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7cbd4a0b-6e30-4aa6-8932-26d6243a8fc6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Querying Zarr takes 2.61x longer than Parquet.\n"
     ]
    }
   ],
   "source": [
    "zarr_mean = np.array(zarr_times).mean()\n",
    "parq_mean = np.array(parq_times).mean()\n",
    "\n",
    "print(f'Querying Zarr takes {zarr_mean / parq_mean:.2f}x longer than Parquet.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c81bdb16",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}