ex_07_ray_data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fe990923",
   "metadata": {},
   "source": [
    "# A Gentle introduction to Ray datasets APIs\n",
    "\n",
    "© 2019-2022, Anyscale. All Rights Reserved\n",
    "\n",
    "📖 [Back to Table of Contents](./ex_00_tutorial_overview.ipynb)<br>\n",
    "⬅️ [Previous notebook](./ex_06_ray_api_calls.ipynb) <br>\n",
    "\n",
    "### Overview\n",
    "\n",
    "Ray data is built primarly for Ray AIR for easy data ingestion and\n",
    "preprocessing for machine learning training, tuning, and scoring. As a library built atop Ray, Ray data allows you to: \n",
    " * exchange data among Ray tasks, actors, libraries, and applications\n",
    " * read/write training data from different file sources and storage\n",
    " * supports myriad [file formats and data sources](https://docs.ray.io/en/latest/data/dataset.html#datasource-compatibility).\n",
    " * provide basic transformations like `map`, `filter`, and `partition` \n",
    "\n",
    "**Note**: Ray datasets is not a replacement for a full-fledged data processing library for doing exploratory data analysis (EDA), extract, transform and load (ETL) or a subsitute for Apache Spark or Dask or Pandas DataFrames. Its primary objective is the last-mile rudimentary and efficient distributed data preprocessing and data ingestion for ML training.\n",
    "\n",
    "<img src=\"images/dataset.png\" width=\"75%\" height=\"45%\">\n",
    "\n",
    "### Learning objectives\n",
    "\n",
    "In this introductory tutorial you will learn:\n",
    " * Ray datasets concepts\n",
    " * how to create, transform, read, and save Ray datasets\n",
    " * use shards for parallel processing of large datasets\n",
    " * use datasets for ML ingestion for distributed training\n",
    " \n",
    "**Note**: Even though you will get an introduction into Ray data as a library here, most of its functionality is subsumed within [Ray AIR](), covered in the next tutorial, for model training, batch inference, and preprocessing workloads. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e6abd1f",
   "metadata": {},
   "source": [
    "### Key concepts\n",
    "\n",
    "To work with Ray Datasets, you need to understand how `Datasets` and `Dataset Pipelines` work. A quick few code examples will shed light into overall merits of Ray Datasets.\n",
    "\n",
    "Let's start with the internal format. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caa7d84d",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Ray Datasets\n",
    "\n",
    "A Ray dataset is backed by [Apache Arrow](https://arrow.apache.org/). As such, a Dataset consists of a *list of Ray object references* to blocks in the *shared memory object store*. Each block holds a set of items in either an [Arrow table](https://arrow.apache.org/docs/python/data.html#tables) (when created from or transformed to tabular or tensor data), a Pandas DataFrame (when created from or transformed to Pandas data), or a Python list (otherwise).\n",
    "\n",
    "<img src=\"images/dataset-arch.png\" width=\"60%\" height=\"30%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9456d3a7",
   "metadata": {},
   "source": [
    "#### Dataset Pipelines\n",
    "Datasets execute their transformations synchronously or eagerly in blocking calls. However, it can be useful to overlap dataset computations with output. This can be done with a `DatasetPipeline`.\n",
    "\n",
    "A `DatasetPipeline` is a unified iterator over a sequence of Ray Datasets, each of which represents a window over the original data. Conceptually, it is similar to a `Spark DStream`, but manages execution over a bounded source of data instead of an unbounded stream. Ray computes each dataset window and stitches their output together into a single logical data iterator. \n",
    "\n",
    "`DatasetPipeline` offers the same transformation and output methods as Datasets (e.g., `map`, `filter`, `split`, `iter_rows`, `to_torch`, etc.).\n",
    "\n",
    "<img src=\"images/ray_dataset_pipeline.jpg\" width=\"60%\" height=\"30%\">\n",
    "\n",
    "#### Common use case for Dataset Pipeline: Saturate your GPU resources\n",
    "\n",
    "The best use case for pipelines is efficient use of CPU/GPU resources\n",
    "\n",
    "<img src=\"images/gpu_saturation.png\" width=\"70%\" height=\"35%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f6c6074",
   "metadata": {},
   "source": [
    "### Datasets Execution Model\n",
    "\n",
    "Let's get the feel of dataset execution model.\n",
    "\n",
    "#### Reading Data\n",
    "Datasets uses Ray tasks, for parallelism, to read data from remote storage or source. When reading from a file-based datasource (e.g., S3, GCS), it creates a number of parallel\n",
    "read tasks equal to the specified read parallelism. You can further [fine tune parallelism](https://docs.ray.io/en/master/data/performance-tips.html#tuning-read-parallelism).\n",
    "One or more files will be assigned to each read task. Each read task reads its assigned files and produces one or more output blocks (Ray objects):\n",
    "\n",
    "<img src=\"https://docs.ray.io/en/master/_images/dataset-read.svg\" height=\"25%\" width=\"50%\">\n",
    "\n",
    "Normally, each read task produces a single output block. Read tasks may split the output into multiple blocks if the data exceeds the target max block size (2GiB by default). This automatic block splitting avoids out-of-memory errors when reading very large single files (e.g., a 100-gigabyte CSV file). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d01de3f6",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Deferred Read Task Execution\n",
    "\n",
    "When a Dataset is created using `ray.data.read_*`, only the first read is executed initially. This avoids blocking Dataset creation on the reading of all data files, enabling inspection functions like `ds.schema()` without incurring high read costs. `<ray.data.Dataset.schema>`() and `ds.show()` can be used right away. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45d75d13",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Dataset Transforms\n",
    "\n",
    "Datasets use either Ray tasks or Ray actors to transform datasets (i.e., for `ds.map_batches()`, `ds.map()`, or `ds.flat_map()`). By default, tasks are used `(compute=\"tasks\")`. Actors can be specified with `compute=\"actors\"`, along with an autoscaling pool of Ray actors used for transformations. \n",
    "\n",
    "Use `ActorStrategyPool` for expensive state initialization (e.g., for GPU-based tasks) that can be re-used. For examples, when workers take time to initialize GPU memory or certain values within an actor that could be resused again.\n",
    "\n",
    "<img src=\"https://docs.ray.io/en/master/_images/dataset-map.svg\" height=\"25%\" width=\"50%\">\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "442e0e08",
   "metadata": {},
   "source": [
    "#### Shuffling Data\n",
    "\n",
    "Certain operations like `ds.sort()` and `ds.groupby()` require data blocks to be partitioned by value.\n",
    "\n",
    "You can also change the partitioning of a Dataset using `ds.random_shuffle()` or `ds.repartition()`. Use the former if you want to randomize the order of elements in the dataset; use the second if you only want to equalize the size of the Dataset blocks (e.g., after a read or transformation that may skew the distribution of block sizes). \n",
    "\n",
    "Note that repartition has two modes, `shuffle=False`, which performs the minimal data movement needed to equalize block sizes, and `shuffle=True`, which performs a full (non-random) distributed shuffle:\n",
    "\n",
    "<img src=\"https://docs.ray.io/en/master/_images/dataset-shuffle.svg\" height=\"25%\" width=\"50%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a03fce60",
   "metadata": {},
   "source": [
    "#### Fault tolerance\n",
    "\n",
    "Datasets relies on task-based 🤹‍♀️ [fault tolerance](https://docs.ray.io/en/latest/ray-core/tasks/fault-tolerance.html) in Ray core. Specifically, a `Dataset` will be automatically recovered by Ray in case of failures. This works through **lineage reconstruction**: a Dataset is a collection of Ray objects stored in shared memory, and if any of these objects are lost, then Ray will recreate them by re-executing the task(s) that created them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc15e040-58d9-4c88-99b4-6d9f9f872188",
   "metadata": {},
   "source": [
    "### Any Questions?\n",
    "\n",
    "Let's move on to some code ...."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c48005fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging, os, random, warnings\n",
    "import ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ab95cf05",
   "metadata": {},
   "outputs": [],
   "source": [
    "warnings.filterwarnings(\"ignore\")\n",
    "os.environ[\"PYTHONWARNINGS\"] = \"ignore\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "92d52759",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"margin-left: 50px;display: flex;flex-direction: row;align-items: center\">\n",
       "        <h3 style=\"color: var(--jp-ui-font-color0)\">Ray</h3>\n",
       "        <svg version=\"1.1\" id=\"ray\" width=\"3em\" viewBox=\"0 0 144.5 144.6\" style=\"margin-left: 3em;margin-right: 3em\">\n",
       "            <g id=\"layer-1\">\n",
       "                <path fill=\"#00a2e9\" class=\"st0\" d=\"M97.3,77.2c-3.8-1.1-6.2,0.9-8.3,5.1c-3.5,6.8-9.9,9.9-17.4,9.6S58,88.1,54.8,81.2c-1.4-3-3-4-6.3-4.1\n",
       "                    c-5.6-0.1-9.9,0.1-13.1,6.4c-3.8,7.6-13.6,10.2-21.8,7.6C5.2,88.4-0.4,80.5,0,71.7c0.1-8.4,5.7-15.8,13.8-18.2\n",
       "                    c8.4-2.6,17.5,0.7,22.3,8c1.3,1.9,1.3,5.2,3.6,5.6c3.9,0.6,8,0.2,12,0.2c1.8,0,1.9-1.6,2.4-2.8c3.5-7.8,9.7-11.8,18-11.9\n",
       "                    c8.2-0.1,14.4,3.9,17.8,11.4c1.3,2.8,2.9,3.6,5.7,3.3c1-0.1,2,0.1,3,0c2.8-0.5,6.4,1.7,8.1-2.7s-2.3-5.5-4.1-7.5\n",
       "                    c-5.1-5.7-10.9-10.8-16.1-16.3C84,38,81.9,37.1,78,38.3C66.7,42,56.2,35.7,53,24.1C50.3,14,57.3,2.8,67.7,0.5\n",
       "                    C78.4-2,89,4.7,91.5,15.3c0.1,0.3,0.1,0.5,0.2,0.8c0.7,3.4,0.7,6.9-0.8,9.8c-1.7,3.2-0.8,5,1.5,7.2c6.7,6.5,13.3,13,19.8,19.7\n",
       "                    c1.8,1.8,3,2.1,5.5,1.2c9.1-3.4,17.9-0.6,23.4,7c4.8,6.9,4.6,16.1-0.4,22.9c-5.4,7.2-14.2,9.9-23.1,6.5c-2.3-0.9-3.5-0.6-5.1,1.1\n",
       "                    c-6.7,6.9-13.6,13.7-20.5,20.4c-1.8,1.8-2.5,3.2-1.4,5.9c3.5,8.7,0.3,18.6-7.7,23.6c-7.9,5-18.2,3.8-24.8-2.9\n",
       "                    c-6.4-6.4-7.4-16.2-2.5-24.3c4.9-7.8,14.5-11,23.1-7.8c3,1.1,4.7,0.5,6.9-1.7C91.7,98.4,98,92.3,104.2,86c1.6-1.6,4.1-2.7,2.6-6.2\n",
       "                    c-1.4-3.3-3.8-2.5-6.2-2.6C99.8,77.2,98.9,77.2,97.3,77.2z M72.1,29.7c5.5,0.1,9.9-4.3,10-9.8c0-0.1,0-0.2,0-0.3\n",
       "                    C81.8,14,77,9.8,71.5,10.2c-5,0.3-9,4.2-9.3,9.2c-0.2,5.5,4,10.1,9.5,10.3C71.8,29.7,72,29.7,72.1,29.7z M72.3,62.3\n",
       "                    c-5.4-0.1-9.9,4.2-10.1,9.7c0,0.2,0,0.3,0,0.5c0.2,5.4,4.5,9.7,9.9,10c5.1,0.1,9.9-4.7,10.1-9.8c0.2-5.5-4-10-9.5-10.3\n",
       "                    C72.6,62.3,72.4,62.3,72.3,62.3z M115,72.5c0.1,5.4,4.5,9.7,9.8,9.9c5.6-0.2,10-4.8,10-10.4c-0.2-5.4-4.6-9.7-10-9.7\n",
       "                    c-5.3-0.1-9.8,4.2-9.9,9.5C115,72.1,115,72.3,115,72.5z M19.5,62.3c-5.4,0.1-9.8,4.4-10,9.8c-0.1,5.1,5.2,10.4,10.2,10.3\n",
       "                    c5.6-0.2,10-4.9,9.8-10.5c-0.1-5.4-4.5-9.7-9.9-9.6C19.6,62.3,19.5,62.3,19.5,62.3z M71.8,134.6c5.9,0.2,10.3-3.9,10.4-9.6\n",
       "                    c0.5-5.5-3.6-10.4-9.1-10.8c-5.5-0.5-10.4,3.6-10.8,9.1c0,0.5,0,0.9,0,1.4c-0.2,5.3,4,9.8,9.3,10\n",
       "                    C71.6,134.6,71.7,134.6,71.8,134.6z\"/>\n",
       "            </g>\n",
       "        </svg>\n",
       "        <table>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left\"><b>Python version:</b></td>\n",
       "                <td style=\"text-align: left\"><b>3.8.13</b></td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left\"><b>Ray version:</b></td>\n",
       "                <td style=\"text-align: left\"><b> 2.0.0</b></td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "    <td style=\"text-align: left\"><b>Dashboard:</b></td>\n",
       "    <td style=\"text-align: left\"><b><a href=\"http://127.0.0.1:8269\" target=\"_blank\">http://127.0.0.1:8269</a></b></td>\n",
       "</tr>\n",
       "\n",
       "        </table>\n",
       "    </div>\n",
       "</div>\n"
      ],
      "text/plain": [
       "RayContext(dashboard_url='127.0.0.1:8269', python_version='3.8.13', ray_version='2.0.0', ray_commit='cba26cc83f6b5b8a2ff166594a65cb74c0ec8740', address_info={'node_ip_address': '127.0.0.1', 'raylet_ip_address': '127.0.0.1', 'redis_address': None, 'object_store_address': '/tmp/ray/session_2022-10-17_14-50-03_628568_56165/sockets/plasma_store', 'raylet_socket_name': '/tmp/ray/session_2022-10-17_14-50-03_628568_56165/sockets/raylet', 'webui_url': '127.0.0.1:8269', 'session_dir': '/tmp/ray/session_2022-10-17_14-50-03_628568_56165', 'metrics_export_port': 59952, 'gcs_address': '127.0.0.1:56022', 'address': '127.0.0.1:56022', 'dashboard_agent_listen_port': 52365, 'node_id': '1e8b39be5f1f84f45b6de93b3b936fd108737f866cce495ed2d6dc70'})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if ray.is_initialized:\n",
    "    ray.shutdown()\n",
    "ray.init(logging_level=logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd7a40fc",
   "metadata": {},
   "source": [
    "### Creating a simple Ray Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c86adcdd",
   "metadata": {},
   "source": [
    "For quick illustration, let's create a generic dataset of 1K integers and look at the schema and underlying datatype. [Ray Data API Documentation](https://docs.ray.io/en/latest/data/package-ref.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "74585645",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = ray.data.range(1000)\n",
    "ds.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "790b11f3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "int"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.schema()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa13aa79-802f-4af1-9d4a-82c8eec49abd",
   "metadata": {},
   "source": [
    "The difference between `show` and `take` is that the former takes one item at time and prints it, while the latter iterates over row items from the dataset, appends to a list and returns it. Underneath, `ds.show()` calls `ds.take()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7fa1bf3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n"
     ]
    }
   ],
   "source": [
    "ds.show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a52f3e2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 1, 2, 3, 4]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.take(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36d6b809",
   "metadata": {},
   "source": [
    "### Creating a large Ray Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9604bc8",
   "metadata": {},
   "source": [
    "Let's create a synthetic dataset, *Homeowners*, of Arrow records (800K) with several columns. To illustrate some simple transformational functions, we'll use this *Homeowners* generated data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f8b5a7c0-75d9-45fa-8915-83e3531dad0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_anonomized_name(num_letters=6):\n",
    "    import string\n",
    "    return (\"anonomized-\" + \"\".join(random.choices(string.ascii_uppercase + string.ascii_lowercase, k=num_letters)))\n",
    "\n",
    "def get_anonomized_ssn():\n",
    "    return (\"anonomized-\" + \"\".join([\"{}\".format(random.randint(0, 9)) for num in range(1, 10)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "de1fad4c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'id': 1,\n",
       "  'ssn': 'anonomized-518861753',\n",
       "  'name': 'anonomized-BICjkW',\n",
       "  'amount': 15.0,\n",
       "  'interest': 0.30000000000000004,\n",
       "  'state': 'NM',\n",
       "  'marital_status': 'undeclared',\n",
       "  'property': 'rental',\n",
       "  'dependents': 5,\n",
       "  'defaulted': 0,\n",
       "  'gender': 'M'}]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NUM_ROWS = 800001\n",
    "STATES = [\"CA\", \"AZ\", \"OR\", \"WA\", \"TX\", \"UT\", \"NV\", \"NM\"]\n",
    "M_STATUS = [\"married\", \"single\", \"domestic\", \"divorced\", \"undeclared\"]\n",
    "GENDER = [\"F\", \"M\", \"U\"]\n",
    "HOME_OWNER = [\"condo\", \"house\", \"rental\", \"cottage\"]\n",
    "\n",
    "items = [{\"id\": i,\n",
    "          \"ssn\": get_anonomized_ssn(),\n",
    "          \"name\": get_anonomized_name(),\n",
    "          \"amount\": i * 1.5 * 10, \n",
    "          \"interest\": random.randint(1,5) * .1,\n",
    "          \"state\": random.choice(STATES),\n",
    "          \"marital_status\": random.choice(M_STATUS),\n",
    "          \"property\": random.choice(HOME_OWNER),\n",
    "          \"dependents\": random.randint(1, 5),\n",
    "          \"defaulted\": random.randint(0,1),\n",
    "          \"gender\":random.choice(GENDER) } for i in range(1,NUM_ROWS)]\n",
    "items[:1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38232ec7",
   "metadata": {},
   "source": [
    "#### Creating a dataset from list of dictionary items\n",
    "\n",
    "Ray data can be created from a dictionary of items. The number of blocks can be [tuned as need](https://docs.ray.io/en/master/data/performance-tips.html#tuning-read-parallelism). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5968c4bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset(num_blocks=200, num_rows=800000, schema={id: int64, ssn: string, name: string, amount: double, interest: double, state: string, marital_status: string, property: string, dependents: int64, defaulted: int64, gender: string})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arrow_ds = ray.data.from_items(items)\n",
    "arrow_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "09a87ef2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "800000"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arrow_ds.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "845b3a6f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ArrowRow({'id': 1,\n",
       "           'ssn': 'anonomized-518861753',\n",
       "           'name': 'anonomized-BICjkW',\n",
       "           'amount': 15.0,\n",
       "           'interest': 0.30000000000000004,\n",
       "           'state': 'NM',\n",
       "           'marital_status': 'undeclared',\n",
       "           'property': 'rental',\n",
       "           'dependents': 5,\n",
       "           'defaulted': 0,\n",
       "           'gender': 'M'})]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arrow_ds.take(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8f84ceb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "id: int64\n",
       "ssn: string\n",
       "name: string\n",
       "amount: double\n",
       "interest: double\n",
       "state: string\n",
       "marital_status: string\n",
       "property: string\n",
       "dependents: int64\n",
       "defaulted: int64\n",
       "gender: string"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arrow_ds.schema()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b622e02",
   "metadata": {},
   "source": [
    "### Saving datasets and reading as a parquet files 🗃\n",
    "Ray datasets support myriad data formats. Let's save this dataset as a parquet file and create five partitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c934bf84",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Repartition: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 24.71it/s]\n",
      "Write Progress: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 21.41it/s]\n"
     ]
    }
   ],
   "source": [
    "path = os.path.abspath(\"data_homeowners/interest.parquet\")\n",
    "arrow_ds.repartition(5).write_parquet(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "65ef04e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 53824\n",
      "-rw-r--r--  1 jules  staff  5418197 Oct 17 14:50 2c358ebfdc1845a09b629141d62a3819_000000.parquet\n",
      "-rw-r--r--  1 jules  staff  5533296 Oct 17 14:50 2c358ebfdc1845a09b629141d62a3819_000001.parquet\n",
      "-rw-r--r--  1 jules  staff  5535217 Oct 17 14:50 2c358ebfdc1845a09b629141d62a3819_000002.parquet\n",
      "-rw-r--r--  1 jules  staff  5533656 Oct 17 14:50 2c358ebfdc1845a09b629141d62a3819_000003.parquet\n",
      "-rw-r--r--  1 jules  staff  5533257 Oct 17 14:50 2c358ebfdc1845a09b629141d62a3819_000004.parquet\n"
     ]
    }
   ],
   "source": [
    "!ls -l data_homeowners/interest.parquet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cab42bd-a2df-462b-b2c8-febfc688b859",
   "metadata": {},
   "source": [
    "### Reading data from parquet files\n",
    "\n",
    "Next, we read a few of the files from the dataset. This read is semi-lazy, where reading of the first file is eagerly executed, but reading of all other files is delayed until the underlying data is needed by downstream operations (e.g., consuming the data with `ds.take()`, or transforming the data with `ds.map_batches())`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "95d585be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.83 ms, sys: 4.27 ms, total: 7.1 ms\n",
      "Wall time: 14.4 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "arrow_ds = ray.data.read_parquet(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa926d11-fd83-4e36-974e-d437c91b0761",
   "metadata": {},
   "source": [
    "Wow that was fast!\n",
    "\n",
    "We can easily inspect the schema of this dataset. For Parquet files, we don’t even have to read the actual data to get the schema or take first rows; we can read it from the lightweight Parquet metadata!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1e9cc56b-92e3-4514-b2b6-93ce05eb5c04",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 12 µs, sys: 6 µs, total: 18 µs\n",
      "Wall time: 18.8 µs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "id: int64\n",
       "ssn: string\n",
       "name: string\n",
       "amount: double\n",
       "interest: double\n",
       "state: string\n",
       "marital_status: string\n",
       "property: string\n",
       "dependents: int64\n",
       "defaulted: int64\n",
       "gender: string"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "arrow_ds.schema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0028f9bc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.65 ms, sys: 2.63 ms, total: 4.28 ms\n",
      "Wall time: 74.1 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[ArrowRow({'id': 1,\n",
       "           'ssn': 'anonomized-518861753',\n",
       "           'name': 'anonomized-BICjkW',\n",
       "           'amount': 15.0,\n",
       "           'interest': 0.30000000000000004,\n",
       "           'state': 'NM',\n",
       "           'marital_status': 'undeclared',\n",
       "           'property': 'rental',\n",
       "           'dependents': 5,\n",
       "           'defaulted': 0,\n",
       "           'gender': 'M'})]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "arrow_ds.take(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34ef1509",
   "metadata": {},
   "source": [
    "### Transforming data with simple methods\n",
    "\n",
    "Ray datasets support transformation in parallel using `map`. It uses Ray tasks to execute eagerly or synchronously. Among others [transformations](https://docs.ray.io/en/latest/data/package-ref.html#dataset-api), it supports`filter`, `flat_map`, `groupBy`etc.\n",
    "\n",
    "Let's try using `.map()`, `.filter()` and `.groupBy` on our dataset. \n",
    "\n",
    "The `map()` and `filter()` are row-based operations. This can be expensive for large datasets. However, you can use `map_batches(...)` with `batch_size=4096` as default. This will create a task per block, and each batch is executed in parallel. Ray tasks are created per block for a map operation. \n",
    "\n",
    "Let's try first the row-based transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a7297bc0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Read->Filter: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:05<00:00,  1.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 56.5 ms, sys: 174 ms, total: 230 ms\n",
      "Wall time: 5.23 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[ArrowRow({'id': 667,\n",
       "           'ssn': 'anonomized-840675366',\n",
       "           'name': 'anonomized-ieiNZs',\n",
       "           'amount': 10005.0,\n",
       "           'interest': 0.30000000000000004,\n",
       "           'state': 'NV',\n",
       "           'marital_status': 'divorced',\n",
       "           'property': 'house',\n",
       "           'dependents': 2,\n",
       "           'defaulted': 0,\n",
       "           'gender': 'F'})]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "arrow_ds.filter(lambda x: x['amount'] > 10000).take(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "788b0897",
   "metadata": {},
   "source": [
    "Let's try a `.map_batches()`. We should expect faster execution. \n",
    "\n",
    "*Question: Why the `.map()` returned `ArrowRow` and `.map_batches` returned `PandasRow`*?\n",
    "\n",
    "Because `map_batch(..., batch_format='native')` promotes it to Pandas DataFrame. You can\n",
    "promote it to [other formats](https://docs.ray.io/en/master/data/package-ref.html#ray.data.Dataset.map_batches) such as `pyarrow`. The default is `native`, which is pandas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "c8116c45",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Read->Map_Batches: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 14.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 86 ms, sys: 39.7 ms, total: 126 ms\n",
      "Wall time: 441 ms\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[PandasRow({'id': 667,\n",
       "            'ssn': 'anonomized-840675366',\n",
       "            'name': 'anonomized-ieiNZs',\n",
       "            'amount': 10005.0,\n",
       "            'interest': 0.30000000000000004,\n",
       "            'state': 'NV',\n",
       "            'marital_status': 'divorced',\n",
       "            'property': 'house',\n",
       "            'dependents': 2,\n",
       "            'defaulted': 0,\n",
       "            'gender': 'F'})]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "arrow_ds.map_batches(lambda df: df[df[\"amount\"] > 10000]).take(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95926ece",
   "metadata": {},
   "source": [
    "You can see that `.map_batches()` is a lot faster than row based. So for large datasets use \n",
    "`.map_batches()`.\n",
    "\n",
    "Let's try a filter operation: both per row operation and per batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9164836e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Read->Filter: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:01<00:00,  2.55it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 28.6 ms, sys: 67.1 ms, total: 95.7 ms\n",
      "Wall time: 1.97 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[ArrowRow({'id': 669,\n",
       "           'ssn': 'anonomized-564839131',\n",
       "           'name': 'anonomized-PxhAue',\n",
       "           'amount': 10035.0,\n",
       "           'interest': 0.1,\n",
       "           'state': 'CA',\n",
       "           'marital_status': 'divorced',\n",
       "           'property': 'condo',\n",
       "           'dependents': 5,\n",
       "           'defaulted': 0,\n",
       "           'gender': 'M'})]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "arrow_ds.filter(lambda x: x[\"amount\"] > 10000.00 and x[\"state\"] == 'CA').take(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "3f98958f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Read->Map_Batches: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 30.06it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 21.2 ms, sys: 13.7 ms, total: 34.9 ms\n",
      "Wall time: 175 ms\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[PandasRow({'id': 669,\n",
       "            'ssn': 'anonomized-564839131',\n",
       "            'name': 'anonomized-PxhAue',\n",
       "            'amount': 10035.0,\n",
       "            'interest': 0.1,\n",
       "            'state': 'CA',\n",
       "            'marital_status': 'divorced',\n",
       "            'property': 'condo',\n",
       "            'dependents': 5,\n",
       "            'defaulted': 0,\n",
       "            'gender': 'M'})]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "arrow_ds.map_batches(lambda df: df[(df[\"amount\"] > 10000) & (df[\"state\"] == \"CA\")]).take(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e754a08",
   "metadata": {},
   "source": [
    "Use `groupBy` state and compute the count, which will will include a shuffle \n",
    "of data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "7588807b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Read: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 100.63it/s]\n",
      "Sort Sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 2705.65it/s]\n",
      "Shuffle Map: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:01<00:00,  4.30it/s]\n",
      "Shuffle Reduce: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 804.25it/s]\n"
     ]
    }
   ],
   "source": [
    "results = arrow_ds.groupby(\"state\").count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "ec43ecf4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'state': 'AZ', 'count()': 99806}\n",
      "{'state': 'CA', 'count()': 99973}\n",
      "{'state': 'NM', 'count()': 99848}\n",
      "{'state': 'NV', 'count()': 99923}\n",
      "{'state': 'OR', 'count()': 100066}\n",
      "{'state': 'TX', 'count()': 100295}\n",
      "{'state': 'UT', 'count()': 100182}\n",
      "{'state': 'WA', 'count()': 99907}\n"
     ]
    }
   ],
   "source": [
    "results.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10edfea7",
   "metadata": {},
   "source": [
    "Get the max of certain columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "7491c77a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Read: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 107.42it/s]\n",
      "Shuffle Map: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 711.53it/s]\n",
      "Shuffle Reduce: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 326.18it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ArrowRow({'max(amount)': 12000000.0,\n",
       "          'max(interest)': 0.5,\n",
       "          'max(dependents)': 5})"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results = arrow_ds.max([\"amount\", \"interest\", \"dependents\"])\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38bf33e9-be3d-4fae-92ee-ddf246116f96",
   "metadata": {},
   "source": [
    "### Any questions?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86f9ae36",
   "metadata": {},
   "source": [
    "### Accessing datasets using batches or iterating by rows\n",
    "\n",
    "Datasets can be passed to Ray tasks or actors and iterated over with `.iter_batches()` or `.iter_rows()`. (This does not incur a copy, since the blocks of the Dataset are passed by reference as Ray objects.) Splitting data as shards and passing to individual Ray Actors to process a shard is a common Ray pattern used in distributed training with Ray actors.\n",
    "\n",
    "Let's examine how we can process a list of shards with a `BatchWorker` Actor  in a distributed fashion.\n",
    "\n",
    "<img src=\"images/batch_worker.jpg\" width=\"70%\" height=\"35%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "678f8e93-f64a-4e2f-9261-74888e43418d",
   "metadata": {},
   "source": [
    "A Ray actor `BatchWorker` working through shards in a batch size of 1024."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3ccb43ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class BatchWorker:\n",
    "    def __init__(self, rank):\n",
    "        self.rank = rank         # this could be rank of CPU/GPU or worker id\n",
    "        self.processed = 0       # how much was processed keeps state of each actor\n",
    "    \n",
    "    @ray.method(num_returns=2)   # we want to return a tuple\n",
    "    def process_shard_list(self, shard: ray.data.Dataset) -> tuple:\n",
    "        for batch in shard.iter_batches(batch_size=1024):\n",
    "            # here you could do something with the batch such as feature\n",
    "            # preprocessing, image processing, minor transformation and then\n",
    "            # saving as a parquet file, etc\n",
    "            self.processed = self.processed + len(batch)\n",
    "            \n",
    "        # return items processed, worker id\n",
    "        return (self.processed, self.rank)     "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcbe723a",
   "metadata": {},
   "source": [
    "#### Create batch workers as Ray actors\n",
    "Each actor will get a shard (list of rows) to work on. We split\n",
    "our dataset `arrow_ds` into five shards. Each `BatchWorker` gets a shard.\n",
    "\n",
    "`.split`() splits shards across these batch of workers by using the `locality_hints`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "ecefae2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Read progress: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 101.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shard row: Dataset(num_blocks=1, num_rows=160000, schema={id: int64, ssn: string, name: string, amount: double, interest: double, state: string, marital_status: string, property: string, dependents: int64, defaulted: int64, gender: string})\n",
      "Number of shards:5\n",
      "Number of shard workers:5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Comprehension list to create five actors, each a BatchWorker\n",
    "batch_workers = [BatchWorker.remote(i) for i in range(1, 6)]\n",
    "\n",
    "# Split into five shards, each one for an actor and co-locate \n",
    "# the blocks of the dataset with each actor to maximize data locality.\n",
    "shards = arrow_ds.split(len(batch_workers), locality_hints=batch_workers)\n",
    "\n",
    "print(f\"Shard row: {shards[0]}\")\n",
    "print(f\"Number of shards:{len(shards)}\")\n",
    "print(f\"Number of shard workers:{len(batch_workers)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7305ab3b",
   "metadata": {},
   "source": [
    "### Launch `BatchWorker` actors\n",
    "\n",
    "Each batch worker processes a shard. Each `BatchWorker.process_shard_list()` returns a object RefID with a tuple as its value. What we get from this comprehension is a list objectRefs as tuples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "b3e8db4b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([[ObjectRef(e0513aa905843ccf04e6e446d797e395d8a313740100000001000000),\n",
       "   ObjectRef(e0513aa905843ccf04e6e446d797e395d8a313740100000002000000)],\n",
       "  [ObjectRef(5f70e045687d2f9ad99d47e53cb935d22040e41b0100000001000000),\n",
       "   ObjectRef(5f70e045687d2f9ad99d47e53cb935d22040e41b0100000002000000)],\n",
       "  [ObjectRef(a4dc031465f905f8940cb4d05c0455a64b8da7f10100000001000000),\n",
       "   ObjectRef(a4dc031465f905f8940cb4d05c0455a64b8da7f10100000002000000)],\n",
       "  [ObjectRef(9e7872a82e7456d9a9e44d80c0a6e3326fba84460100000001000000),\n",
       "   ObjectRef(9e7872a82e7456d9a9e44d80c0a6e3326fba84460100000002000000)],\n",
       "  [ObjectRef(cd25e647a728676b105824817a18aace132ccfb00100000001000000),\n",
       "   ObjectRef(cd25e647a728676b105824817a18aace132ccfb00100000002000000)]],\n",
       " 5)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "object_refs = [w.process_shard_list.remote(s) for w, s in zip(batch_workers, shards)]\n",
    "object_refs, len(object_refs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6afaddf5",
   "metadata": {},
   "source": [
    "Fetch the values from the returned list of ObjectRefs, which is a tuple of (batch_size, worker_rank)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "7aa6d80f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[160000, 1], [160000, 2], [160000, 3], [160000, 4], [160000, 5]]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "values = [ray.get(ref) for ref in object_refs]\n",
    "values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "b48dea6d-87bf-486e-b95e-1b271ed276e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[None, None, None, None, None]"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Kill the BatchWorkers\n",
    "[ray.kill(w) for w in batch_workers]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e864e774-8b14-44a9-b9ee-d0ca9209c14b",
   "metadata": {},
   "source": [
    "### Any questions?\n",
    "\n",
    "Let's examine `DataPipelines`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea260db0",
   "metadata": {},
   "source": [
    "### Creating and using Ray dataset pipelines\n",
    "\n",
    "What are dataset pipelines and how they are different from Ray datasets? \n",
    "\n",
    "Datasets perform transformations or operations eagerly or synchronously, whereas [DataPipelines](https://docs.ray.io/en/latest/data/package-ref.html#datasetpipeline-api) can execute in an overlapped pipeline executions. For example, if you had operations that require reading from a file, transforming data, and then doing some minor feature engineering, these operations can be executed in a normal pipeline fashion. This allows for the overlapped execution of data input (e.g., reading files), computation (e.g., feature preprocessing), and training (e.g., distributed ML training). \n",
    "\n",
    "<img src=\"images/pipeline_window.jpg\" width=\"70%\" height=\"35%\">\n",
    "\n",
    "#### Two ways to create `DatasetPipeline`\n",
    "\n",
    "A `DatasetPipeline` can be constructed in two ways: either by pipelining the execution of an existing Dataset (via `Dataset.window`) or generating repeats of an existing Dataset (via `Dataset.repeat`). \n",
    "\n",
    "Let's have a go at it with the first method and see what we can do with our simple and synthetic data from above using `Dataset.window`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c8e4aec",
   "metadata": {},
   "source": [
    "### Using Dataset.window"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4270d8d",
   "metadata": {},
   "source": [
    "Create simple functions or operations to be executed in a overlapped manner in the pipeline. These functions are simple to illustrate a point. But normally they can be complex for a particular use case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "fde9105f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def divide_row_value(row, n) -> int:\n",
    "    return round(row / n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "d544e0d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def double_row_value(row, n) -> int:\n",
    "    return row * n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "25ae7078",
   "metadata": {},
   "outputs": [],
   "source": [
    "def modulo_row_value(row , n) -> int:\n",
    "    return row % random.randint(1, n)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8e447c5",
   "metadata": {},
   "source": [
    "#### Create a window based pipeline\n",
    "With each window of 50 blocks. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "c0575078",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetPipeline(num_windows=1, num_stages=2)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Use our original simple dataset from above with 1K rows in integer\n",
    "ds_pipe = ds.window(blocks_per_window=50)\n",
    "ds_pipe"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df1cc58b",
   "metadata": {},
   "source": [
    "### Applying overlapping transforms to pipelines\n",
    "\n",
    "#### Example 1:\n",
    "This adds various overlapping functions with our simple dataset from above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "5b722702",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetPipeline(num_windows=1, num_stages=5)\n"
     ]
    }
   ],
   "source": [
    "ds_pipe = ds_pipe.map(lambda row: divide_row_value(row, 2))\n",
    "ds_pipe = ds_pipe.map(lambda row: double_row_value(row, 3))\n",
    "ds_pipe = ds_pipe.map(lambda row: modulo_row_value(row, 4))\n",
    "print(ds_pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be3f9992",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Iterate our pipeline\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "3c0c2594",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stage 0:   0%|                                                                                                                             | 0/1 [00:00<?, ?it/s]\n",
      "  0%|                                                                                                                                      | 0/1 [00:00<?, ?it/s]\u001b[A\n",
      "Stage 1:   0%|                                                                                                                             | 0/1 [00:00<?, ?it/s]\u001b[A\n",
      "Stage 1: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.74it/s]\u001b[A\n",
      "Stage 0: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total value of the results: 393\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "results=[]\n",
    "for row in ds_pipe.iter_rows():\n",
    "    results.append(row)\n",
    "# print(f\"Results from each pipeline map function:{results}\")\n",
    "print(f\"Total value of the results: {sum(results)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4e28847",
   "metadata": {},
   "source": [
    "#### Example 2\n",
    "\n",
    "Let's try a `Datapipeline` with our synthetic data *Homewowners* generated from above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "184f2504",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to count or return based on the condition\n",
    "def count_state(row, state) -> int:\n",
    "    return 1 if row['state'] == state and row[\"defaulted\"] else 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "f5ce3775",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetPipeline(num_windows=1, num_stages=2)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a window-based pipeline\n",
    "arrow_ds_pipe = arrow_ds.window(blocks_per_window=50)\n",
    "arrow_ds_pipe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "b5db9b91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetPipeline(num_windows=1, num_stages=3)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Add the first stage to the pipeline\n",
    "arrow_ds_pipe = arrow_ds_pipe.map(lambda row: count_state(row, \"CA\"))\n",
    "arrow_ds_pipe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "bc6376fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Stage 0:   0%|                                                                                                                                | 0/1 [00:00<?, ?it/s]\n",
      "  0%|                                                                                                                                         | 0/1 [00:00<?, ?it/s]\u001b[A\n",
      "Stage 1:   0%|                                                                                                                                | 0/1 [00:00<?, ?it/s]\u001b[A\n",
      "Stage 1: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.88s/it]\u001b[A\n",
      "Stage 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total rows for CA state and defaulted loans rows: 50029\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Execute the pipeline \n",
    "results=[]\n",
    "for row in arrow_ds_pipe.iter_rows():\n",
    "    results.append(row)\n",
    "print(f\"Total rows for CA state and defaulted loans rows: {sum(results)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12f8f3df-cb13-48ae-8cda-e53c8c890fc4",
   "metadata": {},
   "source": [
    "### Any questions?\n",
    "\n",
    "Let's take the pattern we used above for `BatchWorker` a step further and use it\n",
    "for distributed model training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51b7071a",
   "metadata": {},
   "source": [
    "## Ingesting data into Model Trainers\n",
    "Let's define a toy `Trainer` actor that takes our synthetic data and trains the model and returns loss value for that trainer. This is a common pattern in Ray libraries for distributing data to Trainers in a Ray cluster.\n",
    "\n",
    "<img src=\"images/trainer_worker.jpg\" width=\"60%\" height=\"30%\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "e66729d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Our dummy model computing loss or accuracy\n",
    "import time\n",
    "\n",
    "def model(input):\n",
    "    # do some training and compute the loss\n",
    "    # ....\n",
    "    # time.sleep(0.005)\n",
    "    return random.uniform(0, 1)\n",
    "\n",
    "@ray.remote\n",
    "class Trainer:\n",
    "    def __init__(self, rank, model):\n",
    "        self.rank = rank\n",
    "        self.model = model # copy of the model (or ray.put(model))\n",
    "        self.loss = 0.0\n",
    "        \n",
    "    def train(self, shard:ray.data.Dataset) -> float:\n",
    "        for epoch in range(1,21):\n",
    "            for batch in shard.iter_batches(batch_size=1024):\n",
    "                output = self.model(batch)\n",
    "                self.loss = output \n",
    "            if epoch % 5 == 0:\n",
    "                print(f'rank: {self.rank} epoch: {epoch}, loss: {self.loss:.3f}')\n",
    "        return self.loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "873a4fb8",
   "metadata": {},
   "source": [
    "#### Create five trainers, each with a copy of the model and each training on its respective shard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "632e4cd3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Actor(Trainer, f5248240a89ca6c89f93df9201000000),\n",
       " Actor(Trainer, f2e51c72c5b5eaab8cf882db01000000),\n",
       " Actor(Trainer, e0947e0b2c3d3f6328b44fb401000000),\n",
       " Actor(Trainer, 6c444dfefdf6bd2dbe4af95901000000),\n",
       " Actor(Trainer, 29b4147167f5992c157e18b401000000)]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainers = [Trainer.remote(i, model) for i in range(1, 6)]\n",
    "trainers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcba3400",
   "metadata": {},
   "source": [
    "#### Split the shards across all trainers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "78fccfbe",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Read progress:   0%|                                                                                                                          | 0/5 [00:00<?, ?it/s]\u001b[2m\u001b[36m(_map_block_nosplit pid=15728)\u001b[0m E0818 20:56:43.668332000 6212366336 chttp2_transport.cc:1111]          Received a GOAWAY with error code ENHANCE_YOUR_CALM and debug data equal to \"too_many_pings\"\n",
      "Read progress: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:01<00:00,  4.57it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[Dataset(num_blocks=1, num_rows=160000, schema={id: int64, ssn: string, name: string, amount: double, interest: double, state: string, marital_status: string, property: string, dependents: int64, defaulted: int64, gender: string}),\n",
       " Dataset(num_blocks=1, num_rows=160000, schema={id: int64, ssn: string, name: string, amount: double, interest: double, state: string, marital_status: string, property: string, dependents: int64, defaulted: int64, gender: string}),\n",
       " Dataset(num_blocks=1, num_rows=160000, schema={id: int64, ssn: string, name: string, amount: double, interest: double, state: string, marital_status: string, property: string, dependents: int64, defaulted: int64, gender: string}),\n",
       " Dataset(num_blocks=1, num_rows=160000, schema={id: int64, ssn: string, name: string, amount: double, interest: double, state: string, marital_status: string, property: string, dependents: int64, defaulted: int64, gender: string}),\n",
       " Dataset(num_blocks=1, num_rows=160000, schema={id: int64, ssn: string, name: string, amount: double, interest: double, state: string, marital_status: string, property: string, dependents: int64, defaulted: int64, gender: string})]"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "shards = arrow_ds.split(n=len(trainers), locality_hints=trainers)\n",
    "shards"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5e27a15",
   "metadata": {},
   "source": [
    "#### Launch our trainers in a distributed fashion\n",
    "\n",
    "This will run across the cluster. Check the dashbard to see five actors launched. On a cluster, they will be on five different nodes, whereas on a single node on five different cores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "276306c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(Trainer pid=16152)\u001b[0m rank: 4 epoch: 5, loss: 0.895\n",
      "\u001b[2m\u001b[36m(Trainer pid=16151)\u001b[0m rank: 3 epoch: 5, loss: 0.448\n",
      "\u001b[2m\u001b[36m(Trainer pid=16149)\u001b[0m rank: 1 epoch: 5, loss: 0.518\n",
      "\u001b[2m\u001b[36m(Trainer pid=16153)\u001b[0m rank: 5 epoch: 5, loss: 0.651\n",
      "\u001b[2m\u001b[36m(Trainer pid=16150)\u001b[0m rank: 2 epoch: 5, loss: 0.297\n",
      "\u001b[2m\u001b[36m(Trainer pid=16149)\u001b[0m rank: 1 epoch: 10, loss: 0.860\n",
      "\u001b[2m\u001b[36m(Trainer pid=16153)\u001b[0m rank: 5 epoch: 10, loss: 0.158\n",
      "\u001b[2m\u001b[36m(Trainer pid=16152)\u001b[0m rank: 4 epoch: 10, loss: 0.750\n",
      "\u001b[2m\u001b[36m(Trainer pid=16150)\u001b[0m rank: 2 epoch: 10, loss: 0.049\n",
      "\u001b[2m\u001b[36m(Trainer pid=16151)\u001b[0m rank: 3 epoch: 10, loss: 0.604\n",
      "\u001b[2m\u001b[36m(Trainer pid=16149)\u001b[0m rank: 1 epoch: 15, loss: 0.569\n",
      "\u001b[2m\u001b[36m(Trainer pid=16153)\u001b[0m rank: 5 epoch: 15, loss: 0.347\n",
      "\u001b[2m\u001b[36m(Trainer pid=16152)\u001b[0m rank: 4 epoch: 15, loss: 0.444\n",
      "\u001b[2m\u001b[36m(Trainer pid=16150)\u001b[0m rank: 2 epoch: 15, loss: 0.297\n",
      "\u001b[2m\u001b[36m(Trainer pid=16151)\u001b[0m rank: 3 epoch: 15, loss: 0.116\n",
      "\u001b[2m\u001b[36m(Trainer pid=16149)\u001b[0m rank: 1 epoch: 20, loss: 0.865\n",
      "\u001b[2m\u001b[36m(Trainer pid=16153)\u001b[0m rank: 5 epoch: 20, loss: 0.992\n",
      "\u001b[2m\u001b[36m(Trainer pid=16152)\u001b[0m rank: 4 epoch: 20, loss: 0.929\n",
      "\u001b[2m\u001b[36m(Trainer pid=16150)\u001b[0m rank: 2 epoch: 20, loss: 0.436\n",
      "\u001b[2m\u001b[36m(Trainer pid=16151)\u001b[0m rank: 3 epoch: 20, loss: 0.006\n"
     ]
    }
   ],
   "source": [
    "object_refs = [t.train.remote(s) for t, s in zip(trainers, shards)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "d113e8e0-ca87-4cd4-9c8e-0d08f56cd51d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.8645845130210484,\n",
       " 0.4362154258789336,\n",
       " 0.005786463921397145,\n",
       " 0.9289078395198194,\n",
       " 0.9924620574567161]"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ray.get(object_refs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "cd4776d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f2afff0",
   "metadata": {},
   "source": [
    "### Exercises\n",
    "\n",
    "Use [Ray data API](https://docs.ray.io/en/latest/data/package-ref.html) for reference:\n",
    "\n",
    " 1. Write some simple transformers, filters, and aggregators with our synthetic data. For example:\n",
    "  * use [`.add_column()`](https://docs.ray.io/en/latest/data/package-ref.html) to add an `age` column\n",
    "  * filter by gender == 'U'\n",
    "  * aggregate (or groupby `property`) and count each. \n",
    " 2. Add additional pipleline stages function `def count_tx(...)` with our synthetic data. For example, count all people in state of `TX` who are `married` and `defaulted`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12f2d47a",
   "metadata": {},
   "source": [
    "### Homework\n",
    "\n",
    "So far we have covered the basics of Ray Datasets. There are advanced topics that you can now explore since you know the basics. Below is a list of tasks you will want to work through at home.\n",
    "\n",
    "1. Re-write all `map_batches()` tasks transformations using [`ActorPoolStrategy` pool](https://docs.ray.io/en/latest/data/transforming-datasets.html#compute-strategy)\n",
    "2. Work through the [NYC example tutorial](extra/ray_data_nyc.ipynb). This explores how you use `.map_batches()` for filtering and map operations and Parquet push-down predicates for projection.\n",
    "3. Peruse the user guides for advanced examples in [data transformation](https://docs.ray.io/en/master/data/transforming-datasets.html#transforming-datasets) and [ML preprocessing](https://docs.ray.io/en/master/data/dataset-ml-preprocessing.html#datasets-ml-preprocessing).\n",
    "4. Read how to do large scale [ML ingest](https://docs.ray.io/en/master/data/examples/big_data_ingestion.html).\n",
    "5. Check out the advanced [pipeline usage](https://docs.ray.io/en/latest/data/advanced-pipelines.html#)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "734cfca1-e601-4e9c-b2c2-f05e4af86b4d",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "1. [Ray Data Documentation](https://docs.ray.io/en/latest/data/dataset.html)\n",
    "2. [Ray Data Webinar Talk](https://www.anyscale.com/events/2022/02/23/ray-datasets-scalable-data-preprocessing-for-distributed-ml)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86a252c3-c756-4f2e-976b-f3056f6722cb",
   "metadata": {},
   "source": [
    "📖 [Back to Table of Contents](./ex_00_tutorial_overview.ipynb)<br>\n",
    "⬅️ [Previous notebook](./ex_06_ray_api_calls.ipynb) <br>\n",
    "\n",
    "Done! 🍻"
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}