bench_lsdb_io.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ec2b5a92-8462-4da8-8ce2-abddd149a447",
   "metadata": {},
   "source": [
    "# Benchamrk LSDB for IO-bound analysis\n",
    "\n",
    "The NB investigates LSDB/Dask overhead in parquet data reading and filtering over pure pyarrow analysis.\n",
    "\n",
    "Cells are self-consistent, so you can run them in any order. It is recommended to run everything few times to make the data \"hot\"."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c37b0f39-5087-4715-978a-faf2a1c0a088",
   "metadata": {},
   "source": [
    "### Naive LSDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "794b5140-a692-4754-8d75-778a760d8ea1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.26 s, sys: 1.2 s, total: 8.45 s\n",
      "Wall time: 25.9 s\n",
      "(4370011, 7)\n"
     ]
    }
   ],
   "source": [
    "import lsdb\n",
    "from dask.distributed import Client\n",
    "\n",
    "\n",
    "# Load Gaia DR3 path on PSC\n",
    "PATH = '/ocean/projects/phy210048p/shared/hipscat/catalogs/gaia_dr3'\n",
    "catalog = lsdb.read_hipscat(\n",
    "    PATH,\n",
    "    # All the following kwargs are passed to pd.read_parquet()\n",
    "    engine='pyarrow',  # should be the default\n",
    "    columns=[\n",
    "        \"source_id\",\n",
    "        \"ra\",\n",
    "        \"dec\",\n",
    "        \"phot_g_mean_mag\",\n",
    "        \"phot_proc_mode\",\n",
    "        \"azero_gspphot\",\n",
    "        \"classprob_dsc_combmod_star\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Get Dask dataframe\n",
    "df = catalog._ddf\n",
    "\n",
    "# Filter it\n",
    "df = df.query(\n",
    "    \"15.0 <= phot_g_mean_mag\"\n",
    "    + \"& phot_g_mean_mag <= 16.0\"\n",
    "    + \"& phot_proc_mode == 0\"\n",
    "    + \"& azero_gspphot < 0.1\"\n",
    "    + \"& classprob_dsc_combmod_star >= 0.99\"\n",
    ")\n",
    "\n",
    "# Run computations with Dask Client\n",
    "with Client(n_workers=16) as client:\n",
    "    %time result = df.compute()\n",
    "    \n",
    "print(result.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "603d7cbf-3cef-420f-b88d-3b2255d2dfa3",
   "metadata": {},
   "source": [
    "### Filter with `lsdb.read_hipscat(filters)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "38f08aa4-476d-4071-b52c-c3e94abeb602",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.86 s, sys: 1.02 s, total: 7.88 s\n",
      "Wall time: 29.5 s\n",
      "(4370011, 7)\n"
     ]
    }
   ],
   "source": [
    "import lsdb\n",
    "from dask.distributed import Client\n",
    "\n",
    "\n",
    "# Load Gaia DR3 path on PSC\n",
    "PATH = '/ocean/projects/phy210048p/shared/hipscat/catalogs/gaia_dr3'\n",
    "catalog = lsdb.read_hipscat(\n",
    "    PATH,\n",
    "    # All the following kwargs are passed to pd.read_parquet()\n",
    "    engine='pyarrow',  # should be the default\n",
    "    columns=[\n",
    "        \"source_id\",\n",
    "        \"ra\",\n",
    "        \"dec\",\n",
    "        \"phot_g_mean_mag\",\n",
    "        \"phot_proc_mode\",\n",
    "        \"azero_gspphot\",\n",
    "        \"classprob_dsc_combmod_star\",\n",
    "    ],\n",
    "    filters=[[  # We need nested list, because inner list does AND\n",
    "        (\"phot_g_mean_mag\", \">=\", 15.0),\n",
    "        (\"phot_g_mean_mag\", \"<=\", 16.0),\n",
    "        (\"phot_proc_mode\", \"==\", 0),\n",
    "        (\"azero_gspphot\", \"<\", 0.1),\n",
    "        (\"classprob_dsc_combmod_star\", \">=\", 0.99),\n",
    "    ]],\n",
    ")\n",
    "\n",
    "# Get Dask dataframe\n",
    "df = catalog._ddf\n",
    "\n",
    "# Run computations with Dask Client\n",
    "with Client(n_workers=16) as client:\n",
    "    %time result = df.compute()\n",
    "    \n",
    "print(result.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c23c178-8120-4243-bec0-4a0b28ff3476",
   "metadata": {},
   "source": [
    "### Use `pyarrow` directly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "116b62a2-628a-4f4d-920c-25f16eb45762",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 59s, sys: 1min 49s, total: 3min 49s\n",
      "Wall time: 38.1 s\n",
      "(4370011, 7)\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import pyarrow.dataset\n",
    "import pyarrow.parquet\n",
    "from pyarrow.dataset import field\n",
    "\n",
    "\n",
    "# Location on PSC\n",
    "PATH = Path('/ocean/projects/phy210048p/shared/hipscat/catalogs/gaia_dr3')\n",
    "\n",
    "# Load HiPSCat as a pyarrow dataset\n",
    "schema = pyarrow.parquet.read_schema(PATH / \"_common_metadata\")\n",
    "dataset = pyarrow.dataset.parquet_dataset(PATH / \"_metadata\", partitioning=\"hive\", schema=schema)\n",
    "\n",
    "# Read to pandas DF\n",
    "# We need a function here, so we can do %time in a single line\n",
    "def read_df():\n",
    "    return dataset.to_table(\n",
    "        columns=[\n",
    "            \"source_id\",\n",
    "            \"ra\",\n",
    "            \"dec\",\n",
    "            \"phot_g_mean_mag\",\n",
    "            \"phot_proc_mode\",\n",
    "            \"azero_gspphot\",\n",
    "            \"classprob_dsc_combmod_star\",\n",
    "        ],\n",
    "        filter=(\n",
    "            (field(\"phot_g_mean_mag\") >= 15.0)\n",
    "            & (field(\"phot_g_mean_mag\") <= 16.0)\n",
    "            & (field(\"phot_proc_mode\") == 0)\n",
    "            & (field(\"azero_gspphot\") < 0.1)\n",
    "            & (field(\"classprob_dsc_combmod_star\") >= 0.99)\n",
    "        ),\n",
    "    )\n",
    "\n",
    "%time df = read_df()\n",
    "\n",
    "print(df.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "859e0677-8960-4f34-82e3-7d8b26f9d799",
   "metadata": {},
   "source": [
    "## Appendix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f347c64-bf63-42ec-bf8c-e5ced719fc5d",
   "metadata": {},
   "source": [
    "### Naive LSDB, but with no columns selected\n",
    "\n",
    "We still do `dtype_backend='pyarrow'` to not have `to_pyarrow_string` overhead, see this issue:\n",
    "\n",
    "https://github.com/astronomy-commons/lsdb/issues/89"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "da5a4d11-699f-4497-8678-1a202a248d06",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 2s, sys: 14.4 s, total: 1min 17s\n",
      "Wall time: 5min 34s\n",
      "(4370011, 155)\n"
     ]
    }
   ],
   "source": [
    "import lsdb\n",
    "from dask.distributed import Client\n",
    "\n",
    "\n",
    "# Load Gaia DR3 path on PSC\n",
    "PATH = '/ocean/projects/phy210048p/shared/hipscat/catalogs/gaia_dr3'\n",
    "catalog = lsdb.read_hipscat(\n",
    "    PATH,\n",
    "    # All the following kwargs are passed to pd.read_parquet()\n",
    "    engine='pyarrow',  # should be the default\n",
    "    dtype_backend='pyarrow',\n",
    ")\n",
    "\n",
    "# Get Dask dataframe\n",
    "df = catalog._ddf\n",
    "\n",
    "# Filter it\n",
    "df = df.query(\n",
    "    \"15.0 <= phot_g_mean_mag\"\n",
    "    + \"& phot_g_mean_mag <= 16.0\"\n",
    "    + \"& phot_proc_mode == 0\"\n",
    "    + \"& azero_gspphot < 0.1\"\n",
    "    + \"& classprob_dsc_combmod_star >= 0.99\"\n",
    ")\n",
    "\n",
    "# Run computations with Dask Client\n",
    "with Client(n_workers=16) as client:\n",
    "    %time result = df.compute()\n",
    "    \n",
    "print(result.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcbbda1c-339e-497b-b266-8d073aba872a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tape-venv-python3.9",
   "language": "python",
   "name": "tape-venv-python3.9"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}