06-Arrow-DataFusion-Ballista/06.001 - 10+ minutes to pyArrow and DataFusion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "687af001-f355-4232-a4c5-d183f7c13460",
   "metadata": {},
   "source": [
    "# 10+ Minutes to pyArrow and Arrow DataFusion "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2f3a832-1aa8-4d57-a3f5-90a7b536fdc7",
   "metadata": {},
   "source": [
    "### <font color='green'>__Support for Google Colab__  </font>  \n",
    "    \n",
    "open this notebook in Colab using the following button:  \n",
    "  \n",
    "<a href=\"https://colab.research.google.com/github/shauryashaurya/learn-data-munging/blob/main/06-Arrow-DataFusion-Ballista/06.001%20-%2010%2B%20minutes%20to%20pyArrow%20and%20DataFusion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>  \n",
    "\n",
    "  \n",
    "<font color='green'>uncomment and execute the cell below to setup and run this notebook on Google Colab.</font>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d47412f-7828-4c9d-9d79-8d96a7f8facb",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "19d5c0db-f62b-408d-8ab5-9a744ba6a05a",
   "metadata": {},
   "source": [
    "This is WIP, just the basics are in place."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "667de2f7-bbf3-42bf-baee-3ece99d47797",
   "metadata": {},
   "source": [
    "# pyArrow Basics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b6765161-ba33-4474-b3c5-418038b2ebaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow as pa"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19784171-dd18-4fe2-87fd-eefb44e9c835",
   "metadata": {},
   "source": [
    "## Arrays and Tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5b44856d-ff9f-4077-a55d-5bc633e6c99c",
   "metadata": {},
   "outputs": [],
   "source": [
    "days = pa.array([1,12,17,23,28], type=pa.int8())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "37b929a6-cccf-4778-850b-c57021a7c427",
   "metadata": {},
   "outputs": [],
   "source": [
    "months = pa.array([1, 3, 5, 7, 1], type=pa.int8())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ef9b8641-b1dd-49de-804a-474d3842f595",
   "metadata": {},
   "outputs": [],
   "source": [
    "years = pa.array([1990, 2000, 1995, 2000, 1995], type=pa.int16())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "53674302-c2e6-433c-99df-dc2ed635cec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "birthdays_table = pa.table([days, months, years],\n",
    "                           names=[\"days\", \"months\", \"years\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "677010b6-238a-4547-85a1-b1cb9164a48b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pyarrow.Table\n",
       "days: int8\n",
       "months: int8\n",
       "years: int16\n",
       "----\n",
       "days: [[1,12,17,23,28]]\n",
       "months: [[1,3,5,7,1]]\n",
       "years: [[1990,2000,1995,2000,1995]]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "birthdays_table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cce883b7-7d61-4106-8472-cb2b05edf07f",
   "metadata": {},
   "source": [
    "## Saving and Loading Tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cf328825-b7a1-4549-a44d-de44f9963e0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow.parquet as pq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56e29ac9-8acb-4bbc-800b-8875908f4f6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create data folder if it doesn't exit\n",
    "!md ..\\data\\pyArrow-DataFusion-Ballista"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "cd807179-82c8-420b-a470-fad7c0f2c9e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "pq.write_table(birthdays_table, '../data/pyArrow-DataFusion-Ballista/birthdays.parquet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8db97584-fd51-4f65-afa2-326df358d159",
   "metadata": {},
   "outputs": [],
   "source": [
    "reloaded_birthdays = pq.read_table('../data/pyArrow-DataFusion-Ballista/birthdays.parquet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "24bb3b43-8de5-421c-ba3c-d1d7ee8f1d0b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pyarrow.Table\n",
       "days: int8\n",
       "months: int8\n",
       "years: int16\n",
       "----\n",
       "days: [[1,12,17,23,28]]\n",
       "months: [[1,3,5,7,1]]\n",
       "years: [[1990,2000,1995,2000,1995]]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reloaded_birthdays"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "481f4fa6-42a9-4fe3-bcc5-3eb02f44e287",
   "metadata": {},
   "source": [
    "## Performing Computations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17b85521-418b-4241-b887-72dcb2750bdf",
   "metadata": {},
   "source": [
    "Here's a [list of available compute functions](https://arrow.apache.org/docs/python/compute.html#compute) for our reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "77eca07c-b6bb-4264-9409-15c3cb385bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow.compute as pc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1ff2ea42-1d1a-4168-a798-d49ff0b41996",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyarrow.lib.StructArray object at 0x000001C6D39C8FA0>\n",
       "-- is_valid: all not null\n",
       "-- child 0 type: int16\n",
       "  [\n",
       "    1990,\n",
       "    2000,\n",
       "    1995\n",
       "  ]\n",
       "-- child 1 type: int64\n",
       "  [\n",
       "    1,\n",
       "    2,\n",
       "    2\n",
       "  ]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pc.value_counts(birthdays_table[\"years\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01fa9b56-acc3-4b93-b2b6-29e0675e8977",
   "metadata": {},
   "source": [
    "## Working with large data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c4a00a7b-263b-41d0-9c2d-414b71e5361c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow.dataset as ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec4a7ffd-1c97-45c7-aa3b-d2b18f6e00f7",
   "metadata": {},
   "source": [
    "Arrow also provides the ```pyarrow.dataset``` API to work with large data, which will handle for you partitioning of your data in smaller chunks.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4cedb93d-8257-49e6-a8fd-4415ec3f4eef",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.write_dataset(birthdays_table, \"../data/pyArrow-DataFusion-Ballista/pyArrow-large-data/2\", format=\"parquet\",\n",
    "                 partitioning=ds.partitioning(\n",
    "                    pa.schema([birthdays_table.schema.field(\"years\")])\n",
    "                ))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9de1745a-eb7f-4a1f-9c95-b51216838607",
   "metadata": {},
   "source": [
    "Loading back the partitioned dataset will detect the chunks\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "744d0013-1bbd-4a13-9d24-607f5ea26c54",
   "metadata": {},
   "outputs": [],
   "source": [
    "birthdays_dataset = ds.dataset(\"../data/pyArrow-DataFusion-Ballista/pyArrow-large-data\", format=\"parquet\", partitioning=[\"years\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "41a66555-6a16-4ba5-9374-64a06b706000",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyarrow._dataset.FileSystemDataset at 0x1c6d39c9420>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "birthdays_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "69c9fceb-7a42-4a2b-831e-303c1d8e9f1a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['../data/pyArrow-DataFusion-Ballista/pyArrow-large-data/2/1990/part-0.parquet',\n",
       " '../data/pyArrow-DataFusion-Ballista/pyArrow-large-data/2/1995/part-0.parquet',\n",
       " '../data/pyArrow-DataFusion-Ballista/pyArrow-large-data/2/2000/part-0.parquet']"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "birthdays_dataset.files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38656b0e-c4ef-4a89-b9e8-d31b4eb07380",
   "metadata": {},
   "source": [
    "Arrow will lazily load chunks of data only when iterating over them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "5d799074-aa8c-4d98-947d-9053923f20e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "\n",
    "current_year = datetime.datetime.utcnow().year"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4ac08939-f0a1-4f11-9d78-943afd1a654c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AGES [\n",
      "  2021\n",
      "]\n",
      "AGES [\n",
      "  2021,\n",
      "  2021\n",
      "]\n",
      "AGES [\n",
      "  2021,\n",
      "  2021\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "for table_chunk in birthdays_dataset.to_batches():\n",
    "    print(\"AGES\", pc.subtract(current_year, table_chunk[\"years\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efb03430-8a43-477d-8224-92e1fdf941d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# more to come, including DataFusion basics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f30b1f97-5e97-4d8c-9044-43a84ac16c28",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}