4-iterate-over-blocks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "97c7412a-a324-4efc-af30-4ec79c24dda5",
   "metadata": {},
   "source": [
    "<img src=\"images/dask_horizontal.svg\"\n",
    "     width=\"300px\"\n",
    "     alt=\"Dask logo\">\n",
    "     \n",
    "# Iterate over Blocks\n",
    "\n",
    "You can manually apply a function to each **block** by iterating over the **blocks** in the array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f20c4add-0ef8-48f2-b4a6-38c870747875",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import dask.array as da\n",
    "\n",
    "arr = da.random.random(size=(1_000, 1_000), chunks=(250, 500))\n",
    "\n",
    "result = 0\n",
    "for block in arr.blocks:\n",
    "    result += block.sum()\n",
    "    \n",
    "result.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31c8d09c-8947-4362-8116-33ef85c7d56d",
   "metadata": {},
   "source": [
    "For dataframes this looks very similar. Here we iterate over **partitions** in a list comprehension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3416b31d-2992-4abf-86d4-fa894f5ae8b0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import dask\n",
    "\n",
    "ddf = dask.datasets.timeseries()\n",
    "\n",
    "result = sum(partition.groupby(\"name\").sum() for partition in ddf.partitions)\n",
    "result.visualize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "690069b1-b46a-4165-b1f6-b0d9fb477413",
   "metadata": {},
   "outputs": [],
   "source": [
    "result.visualize(optimize_graph=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "368715c4-f704-49f8-b2c9-6d6e708647b5",
   "metadata": {},
   "source": [
    "## Potentially better graph\n",
    "\n",
    "Another approach to this problem would be to add neighbors together repeatedly until you only have one left."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a18885b-b2ef-4dd0-8bb8-595fdfc0f5b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = [partition.groupby(\"name\").sum() for partition in ddf.partitions]\n",
    "    \n",
    "while len(outputs) > 1:\n",
    "    start = len(outputs) % 2  # 1 if odd, 0 if even\n",
    "    outputs = [\n",
    "        *outputs[:start],\n",
    "        *[outputs[i] + outputs[i + 1] for i in range(start, len(outputs), 2)]\n",
    "    ]\n",
    "\n",
    "output = outputs[0]\n",
    "output.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51967029-7817-4025-873a-f5095cc296af",
   "metadata": {},
   "source": [
    "That graph optimizes to the same as our original."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aca49ba1-cf27-4161-917b-7e7a87ff8359",
   "metadata": {},
   "outputs": [],
   "source": [
    "output.visualize(optimize_graph=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}