icechunk-python/notebooks/demo-dummy-data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d38f2066-411f-4c74-bf6f-aaa54b3472f5",
   "metadata": {},
   "source": [
    "# Icechunk with dummy data\n",
    "\n",
    "This demo illustrates how to use Icechunk as a Zarr store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2abaa80e-f07f-4e6d-b322-0ed280ec77e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "import numpy as np\n",
    "import zarr\n",
    "from icechunk import Repository, in_memory_storage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a382ed5-cd95-481a-b692-c1010c492cff",
   "metadata": {},
   "source": [
    "## Create a new Zarr store backed by Icechunk\n",
    "\n",
    "This example uses an in-memory store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4f40240b-eb4b-408b-8fd9-bb4e5a60a34d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<icechunk.repository.Repository at 0x118a19fd0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "repo = Repository.create(\n",
    "    storage=in_memory_storage()\n",
    ")\n",
    "repo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dc0d6b9-92be-42f0-9cc6-6930954deb3f",
   "metadata": {},
   "source": [
    "This dictionary will contain array paths and data that were written to Icechunk, so that we can check correctness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8fa3197a-0674-431f-9dc1-c59fab055cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "expected = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cba285dc-d153-47f7-b854-4eb67b09995a",
   "metadata": {},
   "source": [
    "These two utility functions generate and write dummy array data to a group."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9690e567-2494-4421-bf47-d9e442c4975f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_array_chunks(size: int, dtype=np.int32):\n",
    "    # dim sizes\n",
    "    nz = 64\n",
    "    nt = 128\n",
    "    nx = ny = int(math.sqrt(size / nz / nt))\n",
    "\n",
    "    # chunk sizes\n",
    "    ct = 2\n",
    "    cz = 8\n",
    "    cx = max(nx // 3, 1)\n",
    "    cy = max(ny // 2, 1)\n",
    "    chunks = (cx, cy, cz, ct)\n",
    "    shape = (nx, ny, nz, nt)\n",
    "\n",
    "    array = np.arange(nx * ny * nz * nt, dtype=dtype).reshape(shape)\n",
    "\n",
    "    return array, chunks\n",
    "\n",
    "\n",
    "def create_array(*, group, name, size, dtype, fill_value) -> np.ndarray:\n",
    "    dims = (\"x\", \"y\", \"z\", \"t\")\n",
    "    attrs = {\"description\": \"icechunk test data\"}\n",
    "\n",
    "    array, chunks = generate_array_chunks(size=size, dtype=dtype)\n",
    "\n",
    "    arr = group.create_array(\n",
    "        name=name,\n",
    "        shape=array.shape,\n",
    "        dtype=dtype,\n",
    "        fill_value=fill_value,\n",
    "        chunks=chunks,\n",
    "        dimension_names=dims,\n",
    "        attributes=attrs,\n",
    "        overwrite=True,\n",
    "    )\n",
    "\n",
    "    arr[:] = array\n",
    "\n",
    "    return array"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68a37c53-0c5c-44dd-a691-f37541c264d2",
   "metadata": {},
   "source": [
    "## A versioned transactional Zarr store\n",
    "\n",
    "### Open the root group, write an attribute, commit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f67ab262",
   "metadata": {},
   "outputs": [],
   "source": [
    "session = repo.writable_session(\"main\")\n",
    "store = session.store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "34a6780e-379d-47ec-bc2a-b599cfab105a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'foo': 'foo'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "root_group = zarr.group(store=store, overwrite=True)\n",
    "root_group.attrs[\"foo\"] = \"foo\"\n",
    "dict(root_group.attrs)  # check that it was written"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6289c6c8-41e7-40a4-a451-074f72e12c98",
   "metadata": {},
   "source": [
    "Commit that change"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d43060dd-6678-45f0-91ed-6786dea6cfa7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ACT3H2K88PB63QXF8EM0'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "first_commit = session.commit(\"wrote a root group attribute\")\n",
    "first_commit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f10f79b2-c44c-4842-84ad-f301934464f6",
   "metadata": {},
   "source": [
    "### Add a array to the root group\n",
    "\n",
    "We save the created array in `expected` to check that the write was correct (later)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ccd0b67d",
   "metadata": {},
   "outputs": [],
   "source": [
    "session = repo.writable_session(\"main\")\n",
    "root_group = zarr.open_group(session.store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "bcaf1dec-65de-4572-ac05-a470ce45e100",
   "metadata": {},
   "outputs": [],
   "source": [
    "expected[\"root-foo\"] = create_array(\n",
    "    group=root_group,\n",
    "    name=\"root-foo\",\n",
    "    size=1 * 1024 * 256,\n",
    "    dtype=np.int32,\n",
    "    fill_value=-1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "aae0f3e8-2db7-437a-8d67-11f07aa47d14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(('root-foo', <Array <icechunk.store.IcechunkStore object at 0x118ce2ed0>/root-foo shape=(5, 5, 64, 128) dtype=int32>),)\n"
     ]
    }
   ],
   "source": [
    "print(root_group.members())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "59254f22-48f2-4629-aec8-546cb12e8f02",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'description': 'icechunk test data'}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dict(root_group[\"root-foo\"].attrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "78b45ec7-ead8-46c5-b553-476abbd2bca4",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_group[\"root-foo\"].attrs[\"update\"] = \"new attr\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2399312c-d53f-443f-8be1-b8702ba6513e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'HFAB7VA15EC54TXB1HSG'"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "second_commit = session.commit(\"added array, updated attr\")\n",
    "second_commit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "edad201d-d9b3-4825-887a-1e6b3bf07e57",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(root_group[\"root-foo\"].attrs) == 2\n",
    "assert len(root_group.members()) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2b6a3c6-9518-4ec4-921f-5303d4e851c7",
   "metadata": {},
   "source": [
    "### Writing and Committing when not on `HEAD` will fail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d904f719-98cf-4f51-8e9a-1631dcb3fcba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "store error: cannot write to read-only store\n"
     ]
    }
   ],
   "source": [
    "session = repo.readonly_session(snapshot_id=first_commit)\n",
    "root_group = zarr.open_group(session.store, mode=\"r\")\n",
    "\n",
    "try:\n",
    "    root_group.attrs[\"update\"] = \"new attr 2\"\n",
    "    session.commit(\"new attr 2\")\n",
    "except ValueError as e:\n",
    "    print(e)\n",
    "else:\n",
    "    raise ValueError(\"should have failed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b135e56-d2c6-43dd-9222-7ecd251b1a7b",
   "metadata": {},
   "source": [
    "### Checkout `HEAD`, make a change, and commit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "d31009db-8f99-48f1-b7bb-3f66875575cc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'CV7SH3YAJSFMNM1BGD3G'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "session = repo.writable_session(\"main\")\n",
    "root_group = zarr.open_group(session.store)\n",
    "root_group[\"root-foo\"].attrs[\"update\"] = \"new attr 2\"\n",
    "third_commit = session.commit(\"new attr 2\")\n",
    "third_commit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "03f8d62b-d8a7-452c-b086-340bfcb76d50",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'RJE8ETHY3GJE8APJC6SG'"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "session = repo.writable_session(\"main\")\n",
    "root_group = zarr.open_group(session.store)\n",
    "\n",
    "root_group.attrs[\"update\"] = \"new attr 2\"\n",
    "fourth_commit = session.commit(\"rewrote array\")\n",
    "fourth_commit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "985f6512-1b90-452d-ad51-955f7ca78f1e",
   "metadata": {},
   "source": [
    "### Create a hierarchy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "aee87354-4c44-4428-a4bf-d38d99b7e608",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'root-foo': dtype('int32')}"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{k: v.dtype for k, v in expected.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f389f3f9-03d5-4625-9856-145e065785f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'9HE2NMYV2RXR0NBYFZNG'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "session = repo.writable_session(\"main\")\n",
    "store = session.store\n",
    "root_group = zarr.open_group(store)\n",
    "\n",
    "newgroup = zarr.group(store=store, path=\"group1/\")\n",
    "expected[\"group1/foo1\"] = create_array(\n",
    "    group=newgroup, name=\"foo1\", dtype=np.float32, size=1 * 1024 * 128, fill_value=-1234\n",
    ")\n",
    "expected[\"group1/foo2\"] = create_array(\n",
    "    group=newgroup, name=\"foo2\", dtype=np.float16, size=1 * 1024 * 64, fill_value=-1234\n",
    ")\n",
    "newgroup = zarr.group(store=store, path=\"group2/\")\n",
    "expected[\"group2/foo3\"] = create_array(\n",
    "    group=newgroup, name=\"foo3\", dtype=np.int64, size=1 * 1024 * 32, fill_value=-1234\n",
    ")\n",
    "fifth_commit = session.commit(\"added groups and arrays\")\n",
    "fifth_commit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2edb658-d9ff-43d3-b028-2aebcf08c86a",
   "metadata": {},
   "source": [
    "### Overwrite an array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "bc9d1ef4-2c06-4147-ad4d-9e8051ac4ea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "session = repo.writable_session(\"main\")\n",
    "store = session.store\n",
    "root_group = zarr.open_group(store)\n",
    "\n",
    "expected[\"root-foo\"] = create_array(\n",
    "    group=root_group,\n",
    "    name=\"root-foo\",\n",
    "    size=1 * 1024 * 128,\n",
    "    dtype=np.int32,\n",
    "    fill_value=-1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "4264bbfa-4193-45e9-bc82-932f488bff28",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'4HY5BRC305VX36HHVCWG'"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "session.commit(\"overwrote root-foo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30217cff-d99e-4397-9c1a-a01d1c39c157",
   "metadata": {},
   "source": [
    "### Examine the hierarchy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "895faf9f-c1ec-4b9b-9676-f6b1745d73de",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(('group1',\n",
       "  <Group <icechunk.store.IcechunkStore object at 0x11e47bfd0>/group1>),\n",
       " ('group2',\n",
       "  <Group <icechunk.store.IcechunkStore object at 0x11e47bfd0>/group2>),\n",
       " ('root-foo',\n",
       "  <Array <icechunk.store.IcechunkStore object at 0x11e47bfd0>/root-foo shape=(4, 4, 64, 128) dtype=int32>))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "root_group.members()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "14c5afc8-640f-464f-8ee0-b8631e0aacc7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(('foo1',\n",
       "  <Array <icechunk.store.IcechunkStore object at 0x11e47bfd0>/group1/foo1 shape=(4, 4, 64, 128) dtype=float32>),\n",
       " ('foo2',\n",
       "  <Array <icechunk.store.IcechunkStore object at 0x11e47bfd0>/group1/foo2 shape=(2, 2, 64, 128) dtype=float16>))"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "root_group[\"group1\"].members()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "1fc3f29a-5915-4c66-bfed-5b75389e44e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(('foo3',\n",
       "  <Array <icechunk.store.IcechunkStore object at 0x11e47bfd0>/group2/foo3 shape=(2, 2, 64, 128) dtype=int64>),)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "root_group[\"group2\"].members()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ff3d952-a2e9-4065-aa3e-239bc0d04c45",
   "metadata": {},
   "source": [
    "### Append"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "fb608382-04e8-4deb-8e2e-3f130845cf8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Array <icechunk.store.IcechunkStore object at 0x11ed54590>/group2/foo3 shape=(2, 2, 64, 128) dtype=int64>\n",
      "<Array <icechunk.store.IcechunkStore object at 0x11ed54590>/group2/foo3 shape=(4, 2, 64, 128) dtype=int64>\n",
      "[    0 16384]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'20P32F6NVC6ZPYDX4NEG'"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "session = repo.writable_session(\"main\")\n",
    "root_group = zarr.open_group(session.store)\n",
    "\n",
    "array = root_group[\"group2/foo3\"]\n",
    "print(array)\n",
    "\n",
    "array.resize((array.shape[0] * 2, *array.shape[1:]))\n",
    "print(array)\n",
    "array[array.shape[0] // 2 :, ...] = expected[\"group2/foo3\"]\n",
    "print(array[2:, 0, 0, 0])\n",
    "expected[\"group2/foo3\"] = np.concatenate([expected[\"group2/foo3\"]] * 2, axis=0)\n",
    "\n",
    "session.commit(\"appended to group2/foo3\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}