notebooks/kerchunker.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "571c8e10-7e57-4710-8064-9e296d4aa55d",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "## Generating Kerchunk sidecar for ATL03 files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "645deb16-b926-4591-8726-a96e624176a7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install git+https://github.com/fsspec/kerchunk\n",
    "\n",
    "# restart the kernel after this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "43efc03a-28fe-456c-9cfd-a23296b9f081",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kerchunk.hdf import SingleHdf5ToZarr\n",
    "import fsspec\n",
    "from pathlib import Path\n",
    "\n",
    "import os\n",
    "import ujson\n",
    "import dask\n",
    "\n",
    "# from pqdm.processes import pqdm\n",
    "from pqdm.threads import pqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8752edc6-b6ba-4a39-9955-177fe066d8e7",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "### Using the example from the Kerchunk documentation "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "27113370-a633-4dcd-8142-0a36a8db8140",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['nasa-cryo-scratch/h5cloud/original/ATL03_20181120182818_08110112_006_02.h5',\n",
       " 'nasa-cryo-scratch/h5cloud/original/ATL03_20190219140808_08110212_006_02.h5',\n",
       " 'nasa-cryo-scratch/h5cloud/original/ATL03_20200217204710_08110612_006_01.h5',\n",
       " 'nasa-cryo-scratch/h5cloud/original/ATL03_20211114142614_08111312_006_01.h5',\n",
       " 'nasa-cryo-scratch/h5cloud/original/ATL03_20230211164520_08111812_006_01.h5']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fs_read_files = fsspec.filesystem('s3') #S3 file system to manage ERA5 files\n",
    "flist = fs_read_files.glob('s3://nasa-cryo-scratch/h5cloud/original/*')\n",
    "\n",
    "def gen_json(file_url):\n",
    "    # default_fill_cache=False avoids caching data in between file chunks to lowers memory usage.\n",
    "    so = dict(mode='rb', default_fill_cache=False, default_cache_type='first') # args to fs.open()\n",
    "    fs = fsspec.filesystem('s3')\n",
    "    fs_local = fsspec.filesystem('')  #local file system to save final jsons to\n",
    "    with fs.open(file_url, **so) as infile:\n",
    "        print(f\"processing:{file_url} \")\n",
    "        h5chunks = SingleHdf5ToZarr(infile, file_url, inline_threshold=300)\n",
    "        # inline threshold adjusts the Size below which binary blocks are included directly in the output\n",
    "        # a higher inline threshold can result in a larger json file but faster loading time\n",
    "        variable = file_url.split('/')[-1].split('.')[0]\n",
    "        month = file_url.split('/')[2]\n",
    "        outf = f'{month}_{variable}.json' #file name to save json to\n",
    "        with fs_local.open(f\"./kerchunked/{outf}\", 'wb') as f:\n",
    "            f.write(ujson.dumps(h5chunks.translate()).encode());\n",
    "flist"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a965c87b-4784-4902-abd2-5b72cf26b0a4",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "### TODO: efficiently parallelize the kerchunking\n",
    "\n",
    "threads do not buys us much when we run these tasks in Python next to the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e9653ca-536d-4772-843c-426914dc14bb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "result = pqdm(flist, gen_json, n_jobs=len(flist))"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}