test_divisions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e572e970-6e3e-4e1c-ba79-a146d7c936ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tape import Ensemble, ColumnMapper\n",
    "import dask.dataframe as dd\n",
    "import dask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfc1d051-51fb-4742-bb8d-89011cf0c67e",
   "metadata": {},
   "source": [
    "# Testing Dask Divisions\n",
    "This notebook is aimed at trying to produce a number of example dataframes, representing possible TAPE datasets, and see the behavior of Dask's Divisions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "81151767-ae28-48cc-82e4-bd22e7b606db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2023.10.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dask.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca0431d4-c72a-4423-98ce-f9c4e89d29fc",
   "metadata": {},
   "source": [
    "## Case 1: Unique Integer Indices\n",
    "\n",
    "Case 1 Findings:\n",
    "* No unexpected issues with a unique integer index\n",
    "* Divisions being known in one dataframe but not both is problematic, Dask doesn't default to not using divisions and instead complains about the lack of divisions in the other dataframe\n",
    "* Having different divisions between divisions between frames is fine, but the number of partitions in any resulting operation will be increased. For TAPE, this means that operations between Object and Source may explode Object's number of partitions if Source is not repartitioned\n",
    "* The Sort flag will sort (shocking) and also populate divisions information, The sorted flag will populate divisions based on the min and max values, and will fail if the dataset is not actually sorted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "98c0efe5-149f-4961-b9ef-dc994e376e7e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 4, 7, 9)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><strong>Dask DataFrame Structure:</strong></div>\n",
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>time</th>\n",
       "      <th>flux</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=3</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "<div>Dask Name: fix-overlap, 3 graph layers</div>"
      ],
      "text/plain": [
       "Dask DataFrame Structure:\n",
       "                  time     flux\n",
       "npartitions=3                  \n",
       "1              float64  float64\n",
       "4                  ...      ...\n",
       "7                  ...      ...\n",
       "9                  ...      ...\n",
       "Dask Name: fix-overlap, 3 graph layers"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 2, 3, 4, 5, 6, 7, 8, 9],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df1 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=True, sort=False)\n",
    "print(df1.divisions)\n",
    "df1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6f71dbe0-2ee1-4b55-a631-08f04fa14db6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dask DataFrame Structure:\n",
      "                  time     flux\n",
      "npartitions=1                  \n",
      "2              float64  float64\n",
      "2                  ...      ...\n",
      "Dask Name: loc, 4 graph layers\n",
      "Dask DataFrame Structure:\n",
      "                  time     flux\n",
      "npartitions=2                  \n",
      "2              float64  float64\n",
      "4                  ...      ...\n",
      "4                  ...      ...\n",
      "Dask Name: loc, 4 graph layers\n"
     ]
    }
   ],
   "source": [
    "# Verify Access\n",
    "\n",
    "print(df1.loc[2]) # Retrieve a row\n",
    "print(df1.loc[2:4]) #Retrieve a range of rows across partition boundaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1ac51cad-0426-46aa-b11c-e882b4795e61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 4, 7, 9)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><strong>Dask DataFrame Structure:</strong></div>\n",
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>time</th>\n",
       "      <th>flux</th>\n",
       "      <th>err</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=3</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "<div>Dask Name: assign, 8 graph layers</div>"
      ],
      "text/plain": [
       "Dask DataFrame Structure:\n",
       "                  time     flux      err\n",
       "npartitions=3                           \n",
       "1              float64  float64  float64\n",
       "4                  ...      ...      ...\n",
       "7                  ...      ...      ...\n",
       "9                  ...      ...      ...\n",
       "Dask Name: assign, 8 graph layers"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Assign a new column from a series with equal partitions\n",
    "\n",
    "new_row = {\n",
    "        \"id\": [1, 2, 3, 4, 5, 6, 7, 8, 9],\n",
    "        \"err\": [1.0, 2.0, 1.0, 3.0, 2.0, 3.0, 4.0, 5.0, 6.0],\n",
    "        }\n",
    "\n",
    "err = dd.from_dict(new_row, npartitions=3).set_index(\"id\", sorted=True, sort=False)[\"err\"]\n",
    "\n",
    "print(err.divisions)\n",
    "df1.assign(err=err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "953b7315-0708-4bb4-93f3-691f8257ce2e",
   "metadata": {},
   "source": [
    "The following cell errors, which is expected. Despite having a series and dataframe that can align row-by-row, the lack of divisions information makes this un-doable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2f9fabee-7168-4604-8643-38dc90f85da9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(None, None, None, None)\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Not all divisions are known, can't align partitions. Please use `set_index` to set the index.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[7], line 12\u001b[0m\n\u001b[1;32m      9\u001b[0m err \u001b[38;5;241m=\u001b[39m dd\u001b[38;5;241m.\u001b[39mfrom_dict(new_row, npartitions\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m3\u001b[39m)\u001b[38;5;241m.\u001b[39mset_index(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mid\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28msorted\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, sort\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124merr\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28mprint\u001b[39m(err\u001b[38;5;241m.\u001b[39mdivisions)\n\u001b[0;32m---> 12\u001b[0m \u001b[43mdf1\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43massign\u001b[49m\u001b[43m(\u001b[49m\u001b[43merr\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43merr\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/core.py:5393\u001b[0m, in \u001b[0;36mDataFrame.assign\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m   5389\u001b[0m     \u001b[38;5;66;03m# Figure out columns of the output\u001b[39;00m\n\u001b[1;32m   5390\u001b[0m     df2 \u001b[38;5;241m=\u001b[39m data\u001b[38;5;241m.\u001b[39m_meta_nonempty\u001b[38;5;241m.\u001b[39massign(\n\u001b[1;32m   5391\u001b[0m         \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m_extract_meta({k: kwargs[k]}, nonempty\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m   5392\u001b[0m     )\n\u001b[0;32m-> 5393\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[43melemwise\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmethods\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43massign\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mpairs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmeta\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdf2\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   5395\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m data\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/core.py:6498\u001b[0m, in \u001b[0;36melemwise\u001b[0;34m(op, meta, out, transform_divisions, *args, **kwargs)\u001b[0m\n\u001b[1;32m   6494\u001b[0m args \u001b[38;5;241m=\u001b[39m _maybe_from_pandas(args)\n\u001b[1;32m   6496\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdask\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdataframe\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmulti\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _maybe_align_partitions\n\u001b[0;32m-> 6498\u001b[0m args \u001b[38;5;241m=\u001b[39m \u001b[43m_maybe_align_partitions\u001b[49m\u001b[43m(\u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   6499\u001b[0m dasks \u001b[38;5;241m=\u001b[39m [arg \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(arg, (_Frame, Scalar, Array))]\n\u001b[1;32m   6500\u001b[0m dfs \u001b[38;5;241m=\u001b[39m [df \u001b[38;5;28;01mfor\u001b[39;00m df \u001b[38;5;129;01min\u001b[39;00m dasks \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(df, _Frame)]\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/multi.py:176\u001b[0m, in \u001b[0;36m_maybe_align_partitions\u001b[0;34m(args)\u001b[0m\n\u001b[1;32m    174\u001b[0m divisions \u001b[38;5;241m=\u001b[39m dfs[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mdivisions\n\u001b[1;32m    175\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mall\u001b[39m(df\u001b[38;5;241m.\u001b[39mdivisions \u001b[38;5;241m==\u001b[39m divisions \u001b[38;5;28;01mfor\u001b[39;00m df \u001b[38;5;129;01min\u001b[39;00m dfs):\n\u001b[0;32m--> 176\u001b[0m     dfs2 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28miter\u001b[39m(\u001b[43malign_partitions\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mdfs\u001b[49m\u001b[43m)\u001b[49m[\u001b[38;5;241m0\u001b[39m])\n\u001b[1;32m    177\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m [a \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(a, _Frame) \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mnext\u001b[39m(dfs2) \u001b[38;5;28;01mfor\u001b[39;00m a \u001b[38;5;129;01min\u001b[39;00m args]\n\u001b[1;32m    178\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m args\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/multi.py:130\u001b[0m, in \u001b[0;36malign_partitions\u001b[0;34m(*dfs)\u001b[0m\n\u001b[1;32m    128\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdfs contains no DataFrame and Series\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    129\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mall\u001b[39m(df\u001b[38;5;241m.\u001b[39mknown_divisions \u001b[38;5;28;01mfor\u001b[39;00m df \u001b[38;5;129;01min\u001b[39;00m dfs1):\n\u001b[0;32m--> 130\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    131\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mNot all divisions are known, can\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mt align \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    132\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpartitions. Please use `set_index` \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    133\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mto set the index.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    134\u001b[0m     )\n\u001b[1;32m    136\u001b[0m divisions \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(unique(merge_sorted(\u001b[38;5;241m*\u001b[39m[df\u001b[38;5;241m.\u001b[39mdivisions \u001b[38;5;28;01mfor\u001b[39;00m df \u001b[38;5;129;01min\u001b[39;00m dfs1])))\n\u001b[1;32m    137\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(divisions) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:  \u001b[38;5;66;03m# single value for index\u001b[39;00m\n",
      "\u001b[0;31mValueError\u001b[0m: Not all divisions are known, can't align partitions. Please use `set_index` to set the index."
     ]
    }
   ],
   "source": [
    "# Assign a new column from a series with unknown partitions\n",
    "# Errors as expected\n",
    "\n",
    "new_row = {\n",
    "        \"id\": [1, 2, 3, 4, 5, 6, 7, 8, 9],\n",
    "        \"err\": [1.0, 2.0, 1.0, 3.0, 2.0, 3.0, 4.0, 5.0, 6.0],\n",
    "        }\n",
    "\n",
    "err = dd.from_dict(new_row, npartitions=3).set_index(\"id\", sorted=False, sort=False)[\"err\"]\n",
    "\n",
    "print(err.divisions)\n",
    "df1.assign(err=err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf378f20-b723-464f-8525-5591e616eaa3",
   "metadata": {},
   "source": [
    "### Case 1.1: Operation with unknown divisions\n",
    "\n",
    "In the case that divisions are unknown in both frames, the operation is able to be performed, and divisions are not set in the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "22e71721-0e01-4158-9efe-f166a9c7a5bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(None, None, None, None)\n",
      "(None, None, None, None)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><strong>Dask DataFrame Structure:</strong></div>\n",
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>time</th>\n",
       "      <th>flux</th>\n",
       "      <th>err</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=3</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "<div>Dask Name: assign, 6 graph layers</div>"
      ],
      "text/plain": [
       "Dask DataFrame Structure:\n",
       "                  time     flux      err\n",
       "npartitions=3                           \n",
       "               float64  float64  float64\n",
       "                   ...      ...      ...\n",
       "                   ...      ...      ...\n",
       "                   ...      ...      ...\n",
       "Dask Name: assign, 6 graph layers"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 2, 3, 4, 5, 6, 7, 8, 9],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df1_1 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=False, sort=False)\n",
    "print(df1_1.divisions)\n",
    "\n",
    "# Assign a new column from a series with unknown partitions\n",
    "# Errors as expected\n",
    "\n",
    "new_row = {\n",
    "        \"id\": [1, 2, 3, 4, 5, 6, 7, 8, 9],\n",
    "        \"err\": [1.0, 2.0, 1.0, 3.0, 2.0, 3.0, 4.0, 5.0, 6.0],\n",
    "        }\n",
    "\n",
    "err = dd.from_dict(new_row, npartitions=3).set_index(\"id\", sorted=False, sort=False)[\"err\"]\n",
    "\n",
    "print(err.divisions)\n",
    "df1_1.assign(err=err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "605687a4-2837-4e90-8e77-ed92c02e0cfb",
   "metadata": {},
   "source": [
    "### Case 1.2: Sorting Instead of Sorted\n",
    "\n",
    "The `sorted` flag lets dask know that the dataset is sorted, allowing it to compute divisions based off the minimum and maximum values. Alternatively, we can set `sort` to have it sort the dataset. This will sort the index and compute divisions, it does find slightly different boundaries as a result here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "618bafb0-e22a-42ff-a46b-d36eed3285f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 3, 6, 9)\n",
      "(1, 3, 6, 9)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><strong>Dask DataFrame Structure:</strong></div>\n",
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>time</th>\n",
       "      <th>flux</th>\n",
       "      <th>err</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=3</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "<div>Dask Name: assign, 16 graph layers</div>"
      ],
      "text/plain": [
       "Dask DataFrame Structure:\n",
       "                  time     flux      err\n",
       "npartitions=3                           \n",
       "1              float64  float64  float64\n",
       "3                  ...      ...      ...\n",
       "6                  ...      ...      ...\n",
       "9                  ...      ...      ...\n",
       "Dask Name: assign, 16 graph layers"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 9, 3, 4, 5, 6, 7, 8, 2], #Ids have been swapped\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df1_2 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=False, sort=True)\n",
    "print(df1_2.divisions)\n",
    "\n",
    "# Assign a new column from a series with unknown partitions\n",
    "# Errors as expected\n",
    "\n",
    "new_row = {\n",
    "        \"id\": [1, 9, 3, 4, 5, 6, 7, 8, 2],\n",
    "        \"err\": [1.0, 2.0, 1.0, 3.0, 2.0, 3.0, 4.0, 5.0, 6.0],\n",
    "        }\n",
    "\n",
    "err = dd.from_dict(new_row, npartitions=3).set_index(\"id\", sorted=False, sort=True)[\"err\"]\n",
    "\n",
    "print(err.divisions)\n",
    "res = df1_2.assign(err=err)\n",
    "res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e36b8ea-2ab3-4680-97fc-821e2c91b70b",
   "metadata": {},
   "source": [
    "### Case 1.3: Having Unequal Divisions\n",
    "\n",
    "Often, divisions between two dataframes will be different. These operations seem fully viable, but can introduce some partition bloat."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1b83219d-e7c8-47e8-808a-89f6b0f3f7b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 3, 6, 9)\n",
      "(1, 4, 7, 9)\n",
      "(1, 3, 4, 6, 7, 9)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><strong>Dask DataFrame Structure:</strong></div>\n",
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>time</th>\n",
       "      <th>flux</th>\n",
       "      <th>err</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=5</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "<div>Dask Name: assign, 14 graph layers</div>"
      ],
      "text/plain": [
       "Dask DataFrame Structure:\n",
       "                  time     flux      err\n",
       "npartitions=5                           \n",
       "1              float64  float64  float64\n",
       "3                  ...      ...      ...\n",
       "...                ...      ...      ...\n",
       "7                  ...      ...      ...\n",
       "9                  ...      ...      ...\n",
       "Dask Name: assign, 14 graph layers"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 9, 3, 4, 5, 6, 7, 8, 2], #Ids have been swapped\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df1_3 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=False, sort=True)\n",
    "print(df1_3.divisions)\n",
    "\n",
    "# Assign a new column from a series with unknown partitions\n",
    "# Errors as expected\n",
    "\n",
    "new_row = {\n",
    "        \"id\": [1, 2, 3, 4, 5, 6, 7, 8, 9], #Ids have not been swapped\n",
    "        \"err\": [1.0, 2.0, 1.0, 3.0, 2.0, 3.0, 4.0, 5.0, 6.0],\n",
    "        }\n",
    "\n",
    "err = dd.from_dict(new_row, npartitions=3).set_index(\"id\", sorted=True, sort=False)[\"err\"]\n",
    "\n",
    "print(err.divisions)\n",
    "res = df1_3.assign(err=err)\n",
    "print(res.divisions) # Increased number of partitions\n",
    "res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8222d09-1a69-46d0-a269-adb70f24fa10",
   "metadata": {},
   "source": [
    "### Case 1.4: Can we game `sorted`?\n",
    "\n",
    "When sorted is set, it just calculates divisions based on the min and max within each partition, can this produce unintended behavior? The below example suggests it errors out instead. The state of the `sort` flag does not appear to affect this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "adc6d86e-06ff-4150-b203-879563714870",
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "('Partitions are not sorted ascending by id', 'In your dataset the (min, max, len) values of id for each partition are : [(1, 9, 3), (4, 6, 3), (2, 8, 3)]')",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[11], line 8\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Dataframe setup\u001b[39;00m\n\u001b[1;32m      2\u001b[0m rows \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m      3\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mid\u001b[39m\u001b[38;5;124m\"\u001b[39m: [\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m9\u001b[39m, \u001b[38;5;241m3\u001b[39m, \u001b[38;5;241m4\u001b[39m, \u001b[38;5;241m5\u001b[39m, \u001b[38;5;241m6\u001b[39m, \u001b[38;5;241m7\u001b[39m, \u001b[38;5;241m8\u001b[39m, \u001b[38;5;241m2\u001b[39m], \u001b[38;5;66;03m#Ids have been swapped\u001b[39;00m\n\u001b[1;32m      4\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtime\u001b[39m\u001b[38;5;124m\"\u001b[39m: [\u001b[38;5;241m10.1\u001b[39m, \u001b[38;5;241m10.2\u001b[39m, \u001b[38;5;241m10.2\u001b[39m, \u001b[38;5;241m11.1\u001b[39m, \u001b[38;5;241m11.2\u001b[39m, \u001b[38;5;241m11.3\u001b[39m, \u001b[38;5;241m11.4\u001b[39m, \u001b[38;5;241m15.0\u001b[39m, \u001b[38;5;241m15.1\u001b[39m],\n\u001b[1;32m      5\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mflux\u001b[39m\u001b[38;5;124m\"\u001b[39m: [\u001b[38;5;241m1.0\u001b[39m, \u001b[38;5;241m2.0\u001b[39m, \u001b[38;5;241m5.0\u001b[39m, \u001b[38;5;241m3.0\u001b[39m, \u001b[38;5;241m1.0\u001b[39m, \u001b[38;5;241m2.0\u001b[39m, \u001b[38;5;241m3.0\u001b[39m, \u001b[38;5;241m4.0\u001b[39m, \u001b[38;5;241m5.0\u001b[39m],\n\u001b[1;32m      6\u001b[0m         }\n\u001b[0;32m----> 8\u001b[0m df1_4 \u001b[38;5;241m=\u001b[39m \u001b[43mdd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrows\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnpartitions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mset_index\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mid\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43msorted\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msort\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28mprint\u001b[39m(df1_4\u001b[38;5;241m.\u001b[39mdivisions)\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/core.py:5272\u001b[0m, in \u001b[0;36mDataFrame.set_index\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m   5269\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m pre_sorted:\n\u001b[1;32m   5270\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdask\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdataframe\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mshuffle\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m set_sorted_index\n\u001b[0;32m-> 5272\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mset_sorted_index\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   5273\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mother\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdrop\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdrop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdivisions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdivisions\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m   5274\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   5275\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   5276\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdask\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdataframe\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mshuffle\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m set_index\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/shuffle.py:1129\u001b[0m, in \u001b[0;36mset_sorted_index\u001b[0;34m(df, index, drop, divisions, **kwargs)\u001b[0m\n\u001b[1;32m   1118\u001b[0m result \u001b[38;5;241m=\u001b[39m map_partitions(\n\u001b[1;32m   1119\u001b[0m     M\u001b[38;5;241m.\u001b[39mset_index,\n\u001b[1;32m   1120\u001b[0m     df,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1125\u001b[0m     transform_divisions\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m   1126\u001b[0m )\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m divisions:\n\u001b[0;32m-> 1129\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcompute_and_set_divisions\u001b[49m\u001b[43m(\u001b[49m\u001b[43mresult\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1130\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(divisions) \u001b[38;5;241m!=\u001b[39m \u001b[38;5;28mlen\u001b[39m(df\u001b[38;5;241m.\u001b[39mdivisions):\n\u001b[1;32m   1131\u001b[0m     msg \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m   1132\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mWhen doing `df.set_index(col, sorted=True, divisions=...)`, \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1133\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdivisions indicates known splits in the index column. In this \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1138\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`df.set_index(col, sorted=True).repartition(divisions=divisions)`\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1139\u001b[0m     )\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/shuffle.py:1097\u001b[0m, in \u001b[0;36mcompute_and_set_divisions\u001b[0;34m(df, **kwargs)\u001b[0m\n\u001b[1;32m   1096\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_and_set_divisions\u001b[39m(df: DataFrame, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m DataFrame:\n\u001b[0;32m-> 1097\u001b[0m     mins, maxes, lens \u001b[38;5;241m=\u001b[39m \u001b[43m_compute_partition_stats\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mindex\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mallow_overlap\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1098\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(mins) \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mlen\u001b[39m(df\u001b[38;5;241m.\u001b[39mdivisions) \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m   1099\u001b[0m         df\u001b[38;5;241m.\u001b[39m_divisions \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(mins) \u001b[38;5;241m+\u001b[39m (maxes[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m],)\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/shuffle.py:1067\u001b[0m, in \u001b[0;36m_compute_partition_stats\u001b[0;34m(column, allow_overlap, **kwargs)\u001b[0m\n\u001b[1;32m   1062\u001b[0m non_empty_maxes \u001b[38;5;241m=\u001b[39m [m \u001b[38;5;28;01mfor\u001b[39;00m m, length \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(maxes, lens) \u001b[38;5;28;01mif\u001b[39;00m length \u001b[38;5;241m!=\u001b[39m \u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m   1063\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m   1064\u001b[0m     \u001b[38;5;28msorted\u001b[39m(non_empty_mins) \u001b[38;5;241m!=\u001b[39m non_empty_mins\n\u001b[1;32m   1065\u001b[0m     \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28msorted\u001b[39m(non_empty_maxes) \u001b[38;5;241m!=\u001b[39m non_empty_maxes\n\u001b[1;32m   1066\u001b[0m ):\n\u001b[0;32m-> 1067\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m   1068\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPartitions are not sorted ascending by \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcolumn\u001b[38;5;241m.\u001b[39mname\u001b[38;5;250m \u001b[39m\u001b[38;5;129;01mor\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mthe index\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   1069\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mIn your dataset the (min, max, len) values of \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcolumn\u001b[38;5;241m.\u001b[39mname\u001b[38;5;250m \u001b[39m\u001b[38;5;129;01mor\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mthe index\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1070\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfor each partition are : \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlist\u001b[39m(\u001b[38;5;28mzip\u001b[39m(mins,\u001b[38;5;250m \u001b[39mmaxes,\u001b[38;5;250m \u001b[39mlens))\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   1071\u001b[0m     )\n\u001b[1;32m   1072\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m allow_overlap \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28many\u001b[39m(\n\u001b[1;32m   1073\u001b[0m     a \u001b[38;5;241m<\u001b[39m\u001b[38;5;241m=\u001b[39m b \u001b[38;5;28;01mfor\u001b[39;00m a, b \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(non_empty_mins[\u001b[38;5;241m1\u001b[39m:], non_empty_maxes[:\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n\u001b[1;32m   1074\u001b[0m ):\n\u001b[1;32m   1075\u001b[0m     warnings\u001b[38;5;241m.\u001b[39mwarn(\n\u001b[1;32m   1076\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPartitions have overlapping values, so divisions are non-unique.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1077\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUse `set_index(sorted=True)` with no `divisions` to allow dask to fix the overlap. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1080\u001b[0m         \u001b[38;5;167;01mUserWarning\u001b[39;00m,\n\u001b[1;32m   1081\u001b[0m     )\n",
      "\u001b[0;31mValueError\u001b[0m: ('Partitions are not sorted ascending by id', 'In your dataset the (min, max, len) values of id for each partition are : [(1, 9, 3), (4, 6, 3), (2, 8, 3)]')"
     ]
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 9, 3, 4, 5, 6, 7, 8, 2], #Ids have been swapped\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df1_4 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=True, sort=True)\n",
    "print(df1_4.divisions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1f07b46-841d-43b1-8f0d-54cf61938cea",
   "metadata": {},
   "source": [
    "## Case 2: Non-unique Integer Indices\n",
    "\n",
    "The most relevant TAPE case.\n",
    "\n",
    "Findings:\n",
    "* Surprisingly, Dask appears to make sure that duplicate Ids are co-partitioned, both when sorting or when it's already sorted it will actually consider the indices on where it draws the partition boundaries\n",
    "* There's a bug/issue where `sorted=True` will not fail when non-unique integer indices are not sorted, as in Case 1.4. It will even draw non-sensical division bounds. So we should be careful about users setting sorted=True on datasets that aren't actually sorted, perhaps an argument for checking if it's sorted within TAPE first and erroring out."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66b75d0a-eda4-4cfa-9693-a1ceda9793a1",
   "metadata": {},
   "source": [
    "### Case 2.1: No spillage of Ids across partitions\n",
    "\n",
    "This is interesting, as it looks like Dask will assign partitions in such a way to enforce no ID spillage. Even to the extent where it will leave one partition empty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3d732176-0f49-4f14-9f30-43eef4265a13",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 2, 3, 4)\n",
      "    time  flux\n",
      "id            \n",
      "1   10.1   1.0\n",
      "1   10.2   2.0\n",
      "1   10.2   5.0\n",
      "    time  flux\n",
      "id            \n",
      "2   11.1   3.0\n",
      "2   11.2   1.0\n",
      "    time  flux\n",
      "id            \n",
      "3   11.3   2.0\n",
      "3   11.4   3.0\n",
      "4   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 1, 1, 2, 2, 3, 3, 4, 4],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df2_1 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=True, sort=False)\n",
    "print(df2_1.divisions)\n",
    "for i in range(3):\n",
    "    print(df2_1.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "40dc3b47-b29a-4e37-90f2-9bf8163d0932",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    time  flux\n",
      "id            \n",
      "2   11.1   3.0\n",
      "2   11.2   1.0\n",
      "    time  flux\n",
      "id            \n",
      "2   11.1   3.0\n",
      "2   11.2   1.0\n",
      "3   11.3   2.0\n",
      "3   11.4   3.0\n",
      "4   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "print(df2_1.loc[2].compute()) # Retrieve a row\n",
    "print(df2_1.loc[2:4].compute()) #Retrieve a range of rows across partition boundaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "fce3ca09-8da3-4a1a-9f92-4203f6941947",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 1, 3, 4)\n",
      "Empty DataFrame\n",
      "Columns: [time, flux]\n",
      "Index: []\n",
      "    time  flux\n",
      "id            \n",
      "1   10.1   1.0\n",
      "1   10.2   2.0\n",
      "1   10.2   5.0\n",
      "1   11.1   3.0\n",
      "2   11.2   1.0\n",
      "2   11.3   2.0\n",
      "    time  flux\n",
      "id            \n",
      "3   11.4   3.0\n",
      "4   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "# Dask really tries to avoid spillage of IDs, this creates an empty first partition\n",
    "rows = {\n",
    "        \"id\": [1, 1, 1, 1, 2, 2, 3, 4, 4],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df2_1 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=True, sort=False)\n",
    "print(df2_1.divisions)\n",
    "for i in range(3):\n",
    "    print(df2_1.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb55def7-7e18-47df-a664-7bfc8d19336c",
   "metadata": {},
   "source": [
    "### Case 2.2: Does `sort` do anything surprising?\n",
    "\n",
    "Seems to sort as expected, avoids spillage of IDs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "491a1eb9-7c6f-4db2-a4d1-04d22f61aec9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 2, 3, 4)\n",
      "    time  flux\n",
      "id            \n",
      "1   11.2   1.0\n",
      "1   10.1   1.0\n",
      "    time  flux\n",
      "id            \n",
      "2   10.2   2.0\n",
      "2   15.0   4.0\n",
      "2   11.1   3.0\n",
      "    time  flux\n",
      "id            \n",
      "3   11.4   3.0\n",
      "3   10.2   5.0\n",
      "4   11.3   2.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 2, 3, 2, 1, 4, 3, 2, 4],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df2_2= dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=False, sort=True)\n",
    "print(df2_2.divisions)\n",
    "for i in range(3):\n",
    "    print(df2_2.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "661ff7f7-7b3c-4d5d-8cd1-69da8141f11c",
   "metadata": {},
   "source": [
    "### Case 2.3: A weird failure case?\n",
    "\n",
    "When `sorted` is True, this produces divisions but does not actually sort the data on them. This causes downstream operations like loc to just not work as expected. This appears to only fail when the dataset is not actually sorted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b7cb1e85-366d-4d16-a8e9-6beac92bfaa8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 1, 2, 4)\n",
      "    time  flux\n",
      "id            \n",
      "1   10.1   1.0\n",
      "2   10.2   2.0\n",
      "3   10.2   5.0\n",
      "    time  flux\n",
      "id            \n",
      "2   11.1   3.0\n",
      "1   11.2   1.0\n",
      "4   11.3   2.0\n",
      "    time  flux\n",
      "id            \n",
      "3   11.4   3.0\n",
      "2   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 2, 3, 2, 1, 4, 3, 2, 4],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df2_3 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=True, sort=False)\n",
    "print(df2_3.divisions)\n",
    "for i in range(3):\n",
    "    print(df2_3.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4681b8a7-3969-4873-a544-4f90d70ae970",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    time  flux\n",
      "id            \n",
      "2   15.0   4.0\n",
      "    time  flux\n",
      "id            \n",
      "2   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "print(df2_3.loc[2].compute()) # Retrieve a row\n",
    "print(df2_3.loc[2:4].compute()) #Retrieve a range of rows across partition boundaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c2ca7c7f-ddfa-4b5b-9c00-84ef1f569e32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 2, 3, 4)\n",
      "    time  flux\n",
      "id            \n",
      "1   10.1   1.0\n",
      "1   10.2   2.0\n",
      "    time  flux\n",
      "id            \n",
      "2   10.2   5.0\n",
      "2   11.1   3.0\n",
      "2   11.2   1.0\n",
      "2   11.3   2.0\n",
      "    time  flux\n",
      "id            \n",
      "3   11.4   3.0\n",
      "3   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 1, 2, 2, 2, 2, 3, 3, 4],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df2_3 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=True, sort=True)\n",
    "print(df2_3.divisions)\n",
    "for i in range(3):\n",
    "    print(df2_3.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "377e2768-6bda-4768-bcca-2e751a5261ca",
   "metadata": {},
   "source": [
    "## Case 3: Non-Integer Indices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98e77370-a3b2-4f04-adfb-a4cef2abe062",
   "metadata": {},
   "source": [
    "### Case 3.1: Unique Non-Integer Indices\n",
    "\n",
    "Actually seems to work well with strings! It's able to sort them and generate divisions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "52c5db75-83dc-4b65-a3c2-7a6e461ecf3a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('a', 'b2', 'g', 'r')\n",
      "    time  flux\n",
      "id            \n",
      "a   10.1   1.0\n",
      "b1  10.2   2.0\n",
      "    time  flux\n",
      "id            \n",
      "b2  10.2   5.0\n",
      "d   11.1   3.0\n",
      "e   11.2   1.0\n",
      "    time  flux\n",
      "id            \n",
      "g   11.4   3.0\n",
      "h   15.0   4.0\n",
      "i   15.1   5.0\n",
      "r   11.3   2.0\n"
     ]
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [\"a\", \"b1\", \"b2\", \"d\", \"e\", \"r\", \"g\", \"h\", \"i\"],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df3_1 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=False, sort=True)\n",
    "print(df3_1.divisions)\n",
    "for i in range(3):\n",
    "    print(df3_1.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "d4a54783-e25a-4f3b-8583-a7aa7f9215e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    time  flux\n",
      "id            \n",
      "a   10.1   1.0\n",
      "    time  flux\n",
      "id            \n",
      "b1  10.2   2.0\n",
      "b2  10.2   5.0\n",
      "d   11.1   3.0\n"
     ]
    }
   ],
   "source": [
    "print(df3_1.loc[\"a\"].compute()) # Retrieve a row\n",
    "print(df3_1.loc[\"b\":\"d\"].compute()) #Retrieve a range of rows across partition boundaries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78d3541b-f025-4e42-94f3-ffb7a4195b29",
   "metadata": {},
   "source": [
    "### Case 3.2: Non-Unique Non-Integer Indices\n",
    "\n",
    "Again, this appears to work well. It even appropriately errors out when attempting to use `sorted` with a non-sorted index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "1354eb4c-1457-49a2-8895-fd8e67012de0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('a', 'b1', 'g', 'g')\n",
      "    time  flux\n",
      "id            \n",
      "a   10.1   1.0\n",
      "a   10.2   2.0\n",
      "    time  flux\n",
      "id            \n",
      "b1  10.2   5.0\n",
      "b1  11.1   3.0\n",
      "b2  11.2   1.0\n",
      "b2  11.3   2.0\n",
      "    time  flux\n",
      "id            \n",
      "g   11.4   3.0\n",
      "g   15.0   4.0\n",
      "g   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "# Works well when sorted\n",
    "rows = {\n",
    "        \"id\": [\"a\", \"a\", \"b1\", \"b1\", \"b2\", \"b2\", \"g\", \"g\", \"g\"],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df3_2 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=True, sort=False)\n",
    "print(df3_2.divisions)\n",
    "for i in range(3):\n",
    "    print(df3_2.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "896dc811-cafc-4c3e-be40-f37b1253b6c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('a', 'b1', 'b2', 'g')\n",
      "    time  flux\n",
      "id            \n",
      "a   11.4   3.0\n",
      "a   15.0   4.0\n",
      "a   10.1   1.0\n",
      "a   10.2   2.0\n",
      "    time  flux\n",
      "id            \n",
      "b1  10.2   5.0\n",
      "b1  11.1   3.0\n",
      "    time  flux\n",
      "id            \n",
      "b2  11.2   1.0\n",
      "b2  11.3   2.0\n",
      "g   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "# Correctly produces the sorted on unsorted error\n",
    "rows = {\n",
    "        \"id\": [\"a\", \"a\", \"b1\", \"b1\", \"b2\", \"b2\", \"a\", \"a\", \"g\"],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df3_2 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=False, sort=True)\n",
    "print(df3_2.divisions)\n",
    "for i in range(3):\n",
    "    print(df3_2.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "dd1ff283-4fa5-4e00-b5e2-73e035904626",
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "('Partitions are not sorted ascending by id', \"In your dataset the (min, max, len) values of id for each partition are : [('a', 'b1', 3), ('b1', 'b2', 3), ('a', 'g', 3)]\")",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[28], line 8\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Correctly produces the sorted on unsorted error\u001b[39;00m\n\u001b[1;32m      2\u001b[0m rows \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m      3\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mid\u001b[39m\u001b[38;5;124m\"\u001b[39m: [\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ma\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ma\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb1\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb1\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb2\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb2\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ma\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ma\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mg\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[1;32m      4\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtime\u001b[39m\u001b[38;5;124m\"\u001b[39m: [\u001b[38;5;241m10.1\u001b[39m, \u001b[38;5;241m10.2\u001b[39m, \u001b[38;5;241m10.2\u001b[39m, \u001b[38;5;241m11.1\u001b[39m, \u001b[38;5;241m11.2\u001b[39m, \u001b[38;5;241m11.3\u001b[39m, \u001b[38;5;241m11.4\u001b[39m, \u001b[38;5;241m15.0\u001b[39m, \u001b[38;5;241m15.1\u001b[39m],\n\u001b[1;32m      5\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mflux\u001b[39m\u001b[38;5;124m\"\u001b[39m: [\u001b[38;5;241m1.0\u001b[39m, \u001b[38;5;241m2.0\u001b[39m, \u001b[38;5;241m5.0\u001b[39m, \u001b[38;5;241m3.0\u001b[39m, \u001b[38;5;241m1.0\u001b[39m, \u001b[38;5;241m2.0\u001b[39m, \u001b[38;5;241m3.0\u001b[39m, \u001b[38;5;241m4.0\u001b[39m, \u001b[38;5;241m5.0\u001b[39m],\n\u001b[1;32m      6\u001b[0m         }\n\u001b[0;32m----> 8\u001b[0m df3_2 \u001b[38;5;241m=\u001b[39m \u001b[43mdd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrows\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnpartitions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mset_index\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mid\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43msorted\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msort\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28mprint\u001b[39m(df3_2\u001b[38;5;241m.\u001b[39mdivisions)\n\u001b[1;32m     10\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m3\u001b[39m):\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/core.py:5272\u001b[0m, in \u001b[0;36mDataFrame.set_index\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m   5269\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m pre_sorted:\n\u001b[1;32m   5270\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdask\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdataframe\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mshuffle\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m set_sorted_index\n\u001b[0;32m-> 5272\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mset_sorted_index\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   5273\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mother\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdrop\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdrop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdivisions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdivisions\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m   5274\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   5275\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   5276\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdask\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdataframe\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mshuffle\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m set_index\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/shuffle.py:1129\u001b[0m, in \u001b[0;36mset_sorted_index\u001b[0;34m(df, index, drop, divisions, **kwargs)\u001b[0m\n\u001b[1;32m   1118\u001b[0m result \u001b[38;5;241m=\u001b[39m map_partitions(\n\u001b[1;32m   1119\u001b[0m     M\u001b[38;5;241m.\u001b[39mset_index,\n\u001b[1;32m   1120\u001b[0m     df,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1125\u001b[0m     transform_divisions\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m   1126\u001b[0m )\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m divisions:\n\u001b[0;32m-> 1129\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcompute_and_set_divisions\u001b[49m\u001b[43m(\u001b[49m\u001b[43mresult\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1130\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(divisions) \u001b[38;5;241m!=\u001b[39m \u001b[38;5;28mlen\u001b[39m(df\u001b[38;5;241m.\u001b[39mdivisions):\n\u001b[1;32m   1131\u001b[0m     msg \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m   1132\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mWhen doing `df.set_index(col, sorted=True, divisions=...)`, \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1133\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdivisions indicates known splits in the index column. In this \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1138\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`df.set_index(col, sorted=True).repartition(divisions=divisions)`\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1139\u001b[0m     )\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/shuffle.py:1097\u001b[0m, in \u001b[0;36mcompute_and_set_divisions\u001b[0;34m(df, **kwargs)\u001b[0m\n\u001b[1;32m   1096\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_and_set_divisions\u001b[39m(df: DataFrame, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m DataFrame:\n\u001b[0;32m-> 1097\u001b[0m     mins, maxes, lens \u001b[38;5;241m=\u001b[39m \u001b[43m_compute_partition_stats\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mindex\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mallow_overlap\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1098\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(mins) \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mlen\u001b[39m(df\u001b[38;5;241m.\u001b[39mdivisions) \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m   1099\u001b[0m         df\u001b[38;5;241m.\u001b[39m_divisions \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(mins) \u001b[38;5;241m+\u001b[39m (maxes[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m],)\n",
      "File \u001b[0;32m~/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/shuffle.py:1067\u001b[0m, in \u001b[0;36m_compute_partition_stats\u001b[0;34m(column, allow_overlap, **kwargs)\u001b[0m\n\u001b[1;32m   1062\u001b[0m non_empty_maxes \u001b[38;5;241m=\u001b[39m [m \u001b[38;5;28;01mfor\u001b[39;00m m, length \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(maxes, lens) \u001b[38;5;28;01mif\u001b[39;00m length \u001b[38;5;241m!=\u001b[39m \u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m   1063\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m   1064\u001b[0m     \u001b[38;5;28msorted\u001b[39m(non_empty_mins) \u001b[38;5;241m!=\u001b[39m non_empty_mins\n\u001b[1;32m   1065\u001b[0m     \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28msorted\u001b[39m(non_empty_maxes) \u001b[38;5;241m!=\u001b[39m non_empty_maxes\n\u001b[1;32m   1066\u001b[0m ):\n\u001b[0;32m-> 1067\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m   1068\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPartitions are not sorted ascending by \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcolumn\u001b[38;5;241m.\u001b[39mname\u001b[38;5;250m \u001b[39m\u001b[38;5;129;01mor\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mthe index\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   1069\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mIn your dataset the (min, max, len) values of \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcolumn\u001b[38;5;241m.\u001b[39mname\u001b[38;5;250m \u001b[39m\u001b[38;5;129;01mor\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mthe index\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1070\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfor each partition are : \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlist\u001b[39m(\u001b[38;5;28mzip\u001b[39m(mins,\u001b[38;5;250m \u001b[39mmaxes,\u001b[38;5;250m \u001b[39mlens))\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   1071\u001b[0m     )\n\u001b[1;32m   1072\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m allow_overlap \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28many\u001b[39m(\n\u001b[1;32m   1073\u001b[0m     a \u001b[38;5;241m<\u001b[39m\u001b[38;5;241m=\u001b[39m b \u001b[38;5;28;01mfor\u001b[39;00m a, b \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(non_empty_mins[\u001b[38;5;241m1\u001b[39m:], non_empty_maxes[:\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n\u001b[1;32m   1074\u001b[0m ):\n\u001b[1;32m   1075\u001b[0m     warnings\u001b[38;5;241m.\u001b[39mwarn(\n\u001b[1;32m   1076\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPartitions have overlapping values, so divisions are non-unique.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1077\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUse `set_index(sorted=True)` with no `divisions` to allow dask to fix the overlap. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1080\u001b[0m         \u001b[38;5;167;01mUserWarning\u001b[39;00m,\n\u001b[1;32m   1081\u001b[0m     )\n",
      "\u001b[0;31mValueError\u001b[0m: ('Partitions are not sorted ascending by id', \"In your dataset the (min, max, len) values of id for each partition are : [('a', 'b1', 3), ('b1', 'b2', 3), ('a', 'g', 3)]\")"
     ]
    }
   ],
   "source": [
    "# Correctly produces the sorted on unsorted error\n",
    "rows = {\n",
    "        \"id\": [\"a\", \"a\", \"b1\", \"b1\", \"b2\", \"b2\", \"a\", \"a\", \"g\"],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df3_2 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=True, sort=False)\n",
    "print(df3_2.divisions)\n",
    "for i in range(3):\n",
    "    print(df3_2.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4644275e-f08f-4ac9-b46a-2ac68e5de566",
   "metadata": {},
   "source": [
    "## Case 4: Non-unique indices interaction with parquet row-groups\n",
    "\n",
    "Dask reads parquet partitions more directly, do divisions maintain the same flexibility when being populated by parquet files?\n",
    "\n",
    "Findings:\n",
    "* Parquet does strongarm partitions, to the point where you can have spillage for a non-unique index. Using `calculate_divisions` within the `read_parquet` call seems unwise given this.\n",
    "* This case is the first instance of actual spillage, and Dask doesn't seem to handle it well, given that loc was missing some occurences of a given ID.\n",
    "* By having dask not retrieve the index from parquet, and instead set it afterwards it seems to be get the flexible behavior found above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "df6e2fdf-1939-486b-8c5d-1ab8b264824c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(None, None, None, None)\n"
     ]
    }
   ],
   "source": [
    "# Dataframe setup\n",
    "rows = {\n",
    "        \"id\": [1, 1, 1, 2, 2, 3, 3, 4, 4],\n",
    "        \"time\": [10.1, 10.2, 10.2, 11.1, 11.2, 11.3, 11.4, 15.0, 15.1],\n",
    "        \"flux\": [1.0, 2.0, 5.0, 3.0, 1.0, 2.0, 3.0, 4.0, 5.0],\n",
    "        }\n",
    "\n",
    "df4_1 = dd.from_dict(rows, npartitions=3).set_index(\"id\", sorted=False, sort=False)\n",
    "print(df4_1.divisions)\n",
    "\n",
    "df4_1.to_parquet(\"divdata/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fc996eb-ccc5-4fe5-8c19-1250ead00809",
   "metadata": {},
   "source": [
    "### Case 4.1: Reading with read_parquet's `calculate_divisions`\n",
    "\n",
    "This creates ID spillage. Operations like loc don't look in both partitions, so Dask seems unsuited to natively handle these spills."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "1d95ca95-e7a3-4f68-9844-ca506a69003c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 2, 3, 4)\n",
      "    time  flux\n",
      "id            \n",
      "1   10.1   1.0\n",
      "1   10.2   2.0\n",
      "1   10.2   5.0\n",
      "    time  flux\n",
      "id            \n",
      "2   11.1   3.0\n",
      "2   11.2   1.0\n",
      "3   11.3   2.0\n",
      "    time  flux\n",
      "id            \n",
      "3   11.4   3.0\n",
      "4   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "df4_1 = dd.read_parquet(\"divdata/*.parquet\", index=\"id\", calculate_divisions=True)\n",
    "\n",
    "print(df4_1.divisions)\n",
    "for i in range(3):\n",
    "    print(df4_1.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "045c709a-2e97-4ff9-906b-d9e99599fc1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    time  flux\n",
      "id            \n",
      "3   11.4   3.0\n",
      "    time  flux\n",
      "id            \n",
      "2   11.1   3.0\n",
      "2   11.2   1.0\n",
      "3   11.3   2.0\n",
      "3   11.4   3.0\n",
      "4   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "print(df4_1.loc[3].compute()) # Retrieve a row\n",
    "print(df4_1.loc[2:4].compute()) #Retrieve a range of rows across partition boundaries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d793f934-0184-482d-b4be-91a4cecc01cf",
   "metadata": {},
   "source": [
    "### Case 4.2: Can we get around this?\n",
    "\n",
    "This is possibly just a bug in read_parquet. By setting the index after a read, we can retrieve the expected behavior as when loading from dictionaries. The second cell achieves this by setting index=False, this seems more elegant than a reset_index call, but not sure if it has any less of a performance penalty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ddab1e48-905b-4cc6-b8ed-427df7b002f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(None, None, None, None)\n",
      "    time  flux\n",
      "id            \n",
      "1   10.1   1.0\n",
      "1   10.2   2.0\n",
      "1   10.2   5.0\n",
      "    time  flux\n",
      "id            \n",
      "2   11.1   3.0\n",
      "2   11.2   1.0\n",
      "3   11.3   2.0\n",
      "    time  flux\n",
      "id            \n",
      "3   11.4   3.0\n",
      "4   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/dbranton/miniforge3/envs/ray310/lib/python3.10/site-packages/dask/dataframe/core.py:5243: UserWarning: New index has same name as existing, this is a no-op.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# If the index is present in the parquet files, set index will be a no-op\n",
    "\n",
    "df4_2 = dd.read_parquet(\"divdata/*.parquet\").set_index(\"id\", sort=True)\n",
    "\n",
    "print(df4_2.divisions)\n",
    "for i in range(3):\n",
    "    print(df4_2.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "6582b3a4-b7fc-49bf-b577-cc093edb07a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 2, 3, 4)\n",
      "    time  flux\n",
      "id            \n",
      "1   10.1   1.0\n",
      "1   10.2   2.0\n",
      "1   10.2   5.0\n",
      "    time  flux\n",
      "id            \n",
      "2   11.1   3.0\n",
      "2   11.2   1.0\n",
      "    time  flux\n",
      "id            \n",
      "3   11.3   2.0\n",
      "3   11.4   3.0\n",
      "4   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "# By turning the index off, this allows for set_index to produce the expected divisions \n",
    "\n",
    "df4_2 = dd.read_parquet(\"divdata/*.parquet\", index=False).set_index(\"id\", sorted=True)\n",
    "\n",
    "print(df4_2.divisions)\n",
    "for i in range(3):\n",
    "    print(df4_2.partitions[i].compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "f0def599-bd8b-4fc9-a463-fab05558bda0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    time  flux\n",
      "id            \n",
      "3   11.3   2.0\n",
      "3   11.4   3.0\n",
      "    time  flux\n",
      "id            \n",
      "2   11.1   3.0\n",
      "2   11.2   1.0\n",
      "3   11.3   2.0\n",
      "3   11.4   3.0\n",
      "4   15.0   4.0\n",
      "4   15.1   5.0\n"
     ]
    }
   ],
   "source": [
    "print(df4_2.loc[3].compute()) # Retrieve a row\n",
    "print(df4_2.loc[2:4].compute()) #Retrieve a range of rows across partition boundaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1e3ca52-58a1-4839-ad2b-be2e65151a4c",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}