pandas_pyarrow_simple.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow.parquet as pq\n",
    "import pyarrow as pa\n",
    "import random\n",
    "import string\n",
    "from multiprocessing import Pool\n",
    "\n",
    "# Define input parameters\n",
    "tmpdir = '.'                    # Parent Directory to write/read data\n",
    "size   = 1000000                # Length of global dataframe\n",
    "npart  = 8                      # Number of dataframe partitions\n",
    "nprocs = 8                      # Number of processes to use\n",
    "path = tmpdir + '/test_dataset' # Parquet dataset directory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Investigating parallel I/O with Pandas, Apache Arrow and Parquet\n",
    "\n",
    "- Rick Zamora (original: 5/14/2019)\n",
    "\n",
    "**Update 6/11/2019**: This document has been modified to reflect recent progress in metadata handling within `arrow.parquet` (e.g. [PR#4405](https://github.com/apache/arrow/pull/4405)).\n",
    "\n",
    "\n",
    "This notebook includes a simple exploration of the pyarrow-parquet writer/reader interface. The general goal is to discover the requirements for efficient parallel processing of partitioned pandas datasets.\n",
    "\n",
    "\n",
    "To begin this exploration, we generate a *large* global pandas dataframe, and partition it into `npart` pieces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<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>a</th>\n",
       "      <th>b</th>\n",
       "      <th>c</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>875000</th>\n",
       "      <td>C</td>\n",
       "      <td>875000</td>\n",
       "      <td>0.900198</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>875001</th>\n",
       "      <td>C</td>\n",
       "      <td>875001</td>\n",
       "      <td>0.352782</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>875002</th>\n",
       "      <td>D</td>\n",
       "      <td>875002</td>\n",
       "      <td>0.716656</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>875003</th>\n",
       "      <td>C</td>\n",
       "      <td>875003</td>\n",
       "      <td>0.161382</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>875004</th>\n",
       "      <td>A</td>\n",
       "      <td>875004</td>\n",
       "      <td>0.197359</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        a       b         c\n",
       "875000  C  875000  0.900198\n",
       "875001  C  875001  0.352782\n",
       "875002  D  875002  0.716656\n",
       "875003  C  875003  0.161382\n",
       "875004  A  875004  0.197359"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Define the global dataframe\n",
    "df = pd.DataFrame(\n",
    "    {\n",
    "    'a': [random.choice(['A', 'B', 'C', 'D']) for c in range(size)], \n",
    "    'b': np.arange(size),\n",
    "    'c': [random.random() for c in range(size)], \n",
    "    }\n",
    ")\n",
    "\n",
    "# Break the global dataframe into partitions\n",
    "lsize = size//npart\n",
    "df_part = [ df.iloc[ x * (size//npart) : size if (x == npart-1) else (x * (size//npart) + size//npart) ] for x in range(npart) ]\n",
    "df_part[npart-1].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Above, we are showing the `head()` of the last partition. In order to write each of these partions, we define the `write_partition` function, which returns a *metadata* object for each partition. Note that the metadata object is actually a dictionary, with `'meta'`, and `'schema'` keys.\n",
    "\n",
    "In order to write each *partition* of the dataframe, we define a specialized `write_partition` function that takes in the partition (a dataframe), as well as the file name (file-system path). We will use the `write_table` function from `arrow.parquet` to perform the actual I/O opperation.  This funtion allows the user to pass in an empty list object using the `metadata_collector` keyword argument.  In the cell below, `metadata_list` will be populated with a parquet-metadata (`pyarrow._parquet.FileMetaData`) object for each dataset partition. After collecting this metadata object, we need to use the `set_file_path` method to correctly define the location of the **data** corresponding to this partition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define `write_partition` funciton\n",
    "def write_partition(df, i, ipath):\n",
    "    filepath = ipath + \"/part.%i.parquet\" % i\n",
    "    t = pa.Table.from_pandas(df)\n",
    "    metadata_list = []\n",
    "    with open(filepath, \"wb\") as fil:\n",
    "        pq.write_table(\n",
    "            t, \n",
    "            fil, \n",
    "            metadata_collector=metadata_list\n",
    "        )\n",
    "        # Set file path attribute in metadata:\n",
    "        metadata_list[0].set_file_path(\"part.%i.parquet\" % i)\n",
    "    # Return metadata & schema\n",
    "    return {'meta': metadata_list[0], 'schema': t.schema}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to write the pandas dataframe in parallel, we can use a simple `multiprocessing.Pool` approach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 44 ms, sys: 16 ms, total: 60 ms\n",
      "Wall time: 122 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['part.2.parquet',\n",
       " 'part.6.parquet',\n",
       " 'part.7.parquet',\n",
       " 'part.4.parquet',\n",
       " 'part.5.parquet',\n",
       " 'part.3.parquet',\n",
       " 'part.0.parquet',\n",
       " 'part.1.parquet']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Write the datafram in parallel (to multiple parquet files)\n",
    "def fwrite(args):\n",
    "    i = args[0]\n",
    "    part = args[1]\n",
    "    ipath = args[2]\n",
    "    filename = ipath + \"/part.%i.parquet\" % i\n",
    "    meta = write_partition(part, i, ipath)\n",
    "    return meta\n",
    "\n",
    "fargs = [ [i, part, path] for i, part in enumerate(df_part) ]\n",
    "os.system('rm -rf ' + path)\n",
    "os.system('mkdir ' + path)\n",
    "p = Pool(nprocs)\n",
    "%time meta = p.map(fwrite, fargs)\n",
    "\n",
    "# Check that there are multiple files\n",
    "files = os.listdir(path)\n",
    "files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we can see that each of the `npart` partitions is written as a separate file called `'part.<part>.parquet'` (where `part` is just an integer in this case). After collecting the file metadata (and schemas) from all partitions within the `meta` list, we can check that the `file_path` attribute is properly set for one of the partitions..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'created_by': 'parquet-cpp version 1.5.1-SNAPSHOT',\n",
       " 'num_columns': 3,\n",
       " 'num_rows': 125000,\n",
       " 'num_row_groups': 1,\n",
       " 'row_groups': [{'num_columns': 3,\n",
       "   'num_rows': 125000,\n",
       "   'total_byte_size': 2063898,\n",
       "   'columns': [{'file_offset': 31587,\n",
       "     'file_path': 'part.0.parquet',\n",
       "     'physical_type': 'BYTE_ARRAY',\n",
       "     'num_values': 125000,\n",
       "     'path_in_schema': 'a',\n",
       "     'is_stats_set': True,\n",
       "     'statistics': {'has_min_max': True,\n",
       "      'min': b'A',\n",
       "      'max': b'D',\n",
       "      'null_count': 0,\n",
       "      'distinct_count': 0,\n",
       "      'num_values': 125000,\n",
       "      'physical_type': 'BYTE_ARRAY'},\n",
       "     'compression': 'SNAPPY',\n",
       "     'encodings': ('PLAIN_DICTIONARY', 'PLAIN', 'RLE'),\n",
       "     'has_dictionary_page': True,\n",
       "     'dictionary_page_offset': 4,\n",
       "     'data_page_offset': 40,\n",
       "     'total_compressed_size': 31583,\n",
       "     'total_uncompressed_size': 31575},\n",
       "    {'file_offset': 797909,\n",
       "     'file_path': 'part.0.parquet',\n",
       "     'physical_type': 'INT64',\n",
       "     'num_values': 125000,\n",
       "     'path_in_schema': 'b',\n",
       "     'is_stats_set': True,\n",
       "     'statistics': {'has_min_max': True,\n",
       "      'min': 0,\n",
       "      'max': 124999,\n",
       "      'null_count': 0,\n",
       "      'distinct_count': 0,\n",
       "      'num_values': 125000,\n",
       "      'physical_type': 'INT64'},\n",
       "     'compression': 'SNAPPY',\n",
       "     'encodings': ('PLAIN_DICTIONARY', 'PLAIN', 'RLE'),\n",
       "     'has_dictionary_page': True,\n",
       "     'dictionary_page_offset': 31633,\n",
       "     'data_page_offset': 531941,\n",
       "     'total_compressed_size': 766276,\n",
       "     'total_uncompressed_size': 1265970},\n",
       "    {'file_offset': 2064033,\n",
       "     'file_path': 'part.0.parquet',\n",
       "     'physical_type': 'DOUBLE',\n",
       "     'num_values': 125000,\n",
       "     'path_in_schema': 'c',\n",
       "     'is_stats_set': True,\n",
       "     'statistics': {'has_min_max': True,\n",
       "      'min': 5.44345e-07,\n",
       "      'max': 0.999996,\n",
       "      'null_count': 0,\n",
       "      'distinct_count': 0,\n",
       "      'num_values': 125000,\n",
       "      'physical_type': 'DOUBLE'},\n",
       "     'compression': 'SNAPPY',\n",
       "     'encodings': ('PLAIN_DICTIONARY', 'PLAIN', 'RLE'),\n",
       "     'has_dictionary_page': True,\n",
       "     'dictionary_page_offset': 797994,\n",
       "     'data_page_offset': 1798065,\n",
       "     'total_compressed_size': 1266039,\n",
       "     'total_uncompressed_size': 1265970}]}],\n",
       " 'format_version': '1.0',\n",
       " 'serialized_size': 1043}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "meta[0]['meta'].to_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we are collecting the metadata/schema information for each of the partitions that is written, we can write the metadata and schema information to specialized files called `_metadata` and `_common_metadata`, respectively.\n",
    "\n",
    "First, we can write `_common_metadata` using the `write_metadata()` function available in `pyarrow.parquet`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "a: string\n",
       "b: int64\n",
       "c: double\n",
       "metadata\n",
       "--------\n",
       "{b'pandas': b'{\"index_columns\": [{\"kind\": \"range\", \"name\": null, \"start\": 0, \"'\n",
       "            b'stop\": 125000, \"step\": 1}], \"column_indexes\": [{\"name\": null, \"f'\n",
       "            b'ield_name\": null, \"pandas_type\": \"unicode\", \"numpy_type\": \"objec'\n",
       "            b't\", \"metadata\": {\"encoding\": \"UTF-8\"}}], \"columns\": [{\"name\": \"a'\n",
       "            b'\", \"field_name\": \"a\", \"pandas_type\": \"unicode\", \"numpy_type\": \"o'\n",
       "            b'bject\", \"metadata\": null}, {\"name\": \"b\", \"field_name\": \"b\", \"pan'\n",
       "            b'das_type\": \"int64\", \"numpy_type\": \"int64\", \"metadata\": null}, {\"'\n",
       "            b'name\": \"c\", \"field_name\": \"c\", \"pandas_type\": \"float64\", \"numpy_'\n",
       "            b'type\": \"float64\", \"metadata\": null}], \"creator\": {\"library\": \"py'\n",
       "            b'arrow\", \"version\": \"0.13.1.dev281+g5999609\"}, \"pandas_version\": '\n",
       "            b'\"0.24.2\"}'}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Write a `_common_metadata` (schema) file using the 0th partition\n",
    "# (Note that this is not actually needed...)\n",
    "schema = meta[0]['schema']\n",
    "common_metadata_path = path + '/_common_metadata' \n",
    "with open(common_metadata_path, \"wb\") as f:\n",
    "    pq.write_metadata(schema, f)\n",
    "schema"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the `pq.write_metadata()` call above uses the `schema` object for the 0th partition. This actually means that not **all** information in `_common_metadata` will be true for the entire dataset. More specifically, the `\"index_columns\"` item will have something like this:\n",
    "\n",
    "```[{\"kind\": \"range\", \"name\": null, \"start\": 0, \"stop\": 125000, \"step\": 1}]```\n",
    "            \n",
    "Therefore, the `\"start\"` and `\"stop\"` values will correspond to the 0th partition, rather than the global dataset.\n",
    "\n",
    "Although we just took the time to write out a `_common_metadata` file, it does not currently make much sense to use the file to read back the dataset.  I say this, because that file does not seem include the type of metadata we really need for typical data-processing tasks (such as the number of partitions, the schema, and the row statistics).\n",
    "\n",
    "Instead, what we really want is the so-called `_metadata` file, which (by convention) contains the aggregated metadata for all partitions within a single file (including row-group statistics). To write this file, we must aggregate our list of metadata objects into a single metadata object (using `append_row_groups`). Then we use `` to write a dedicated metadata file..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "_meta = meta[0]['meta']\n",
    "for i in range(1,len(meta)):\n",
    "    _meta.append_row_groups(meta[i]['meta'])\n",
    "metadata_path = path + '/_metadata'\n",
    "with open(metadata_path, \"wb\") as f:\n",
    "    _meta.write_metadata_file(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have an aggregated metadata file, we can read that file (and only that file) to learn a lot about the dataset without using a file-path glob and parsing the metadata of every single file in the dataset. Since `_metadata` specifies the file-system path for each row-group column in the dataset, the files can also have a complex/nested/remote organization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of dataset rows: 1000000\n",
      "Number of dataset columns: 3\n",
      "Number of dataset row_groups: 8\n"
     ]
    }
   ],
   "source": [
    "# Read back the metadata\n",
    "meta = pq.read_metadata(metadata_path)\n",
    "md = meta.to_dict()\n",
    "\n",
    "print(\"Number of dataset rows: \"+str(md['num_rows']))\n",
    "print(\"Number of dataset columns: \"+str(md['num_columns']))\n",
    "print(\"Number of dataset row_groups: \"+str(md['num_row_groups']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an example/demonstration, we can use `multiprocessing.Pool` again to generate a *metadata* dataframe (note: there are probably much more efficient ways to parse the `_metadata` file)..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 0 ns, sys: 4 ms, total: 4 ms\n",
      "Wall time: 9.5 ms\n"
     ]
    }
   ],
   "source": [
    "# Get files and global statistics\n",
    "def fparse(args):\n",
    "    i = args[0]\n",
    "    rg = args[1]\n",
    "    rg_dict = {'rgID': [], 'colID': [], 'cmin': [], 'cmax': [], 'file_path': []}\n",
    "\n",
    "    for j, col in enumerate(rg['columns']):\n",
    "        rg_dict['rgID'].append(i)\n",
    "        rg_dict['colID'].append(col['path_in_schema'])\n",
    "        if isinstance(col['statistics']['min'], bytes):\n",
    "            rg_dict['cmin'].append(col['statistics']['min'].decode(\"utf-8\"))\n",
    "            rg_dict['cmax'].append(col['statistics']['max'].decode(\"utf-8\"))\n",
    "        else:\n",
    "            rg_dict['cmin'].append(col['statistics']['min'])\n",
    "            rg_dict['cmax'].append(col['statistics']['max'])\n",
    "        rg_dict['file_path'].append(col['file_path'])\n",
    "    return pd.DataFrame(rg_dict)\n",
    "\n",
    "fargs = [ [i, rg] for i, rg in enumerate(md['row_groups']) ]\n",
    "p = Pool(nprocs)\n",
    "%time frames = p.map(fparse, fargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<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>rgID</th>\n",
       "      <th>colID</th>\n",
       "      <th>cmin</th>\n",
       "      <th>cmax</th>\n",
       "      <th>file_path</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>a</td>\n",
       "      <td>A</td>\n",
       "      <td>D</td>\n",
       "      <td>part.0.parquet</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>b</td>\n",
       "      <td>0</td>\n",
       "      <td>124999</td>\n",
       "      <td>part.0.parquet</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>c</td>\n",
       "      <td>5.44345e-07</td>\n",
       "      <td>0.999996</td>\n",
       "      <td>part.0.parquet</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>a</td>\n",
       "      <td>A</td>\n",
       "      <td>D</td>\n",
       "      <td>part.1.parquet</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>b</td>\n",
       "      <td>125000</td>\n",
       "      <td>249999</td>\n",
       "      <td>part.1.parquet</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   rgID colID         cmin      cmax       file_path\n",
       "0     0     a            A         D  part.0.parquet\n",
       "1     0     b            0    124999  part.0.parquet\n",
       "2     0     c  5.44345e-07  0.999996  part.0.parquet\n",
       "0     1     a            A         D  part.1.parquet\n",
       "1     1     b       125000    249999  part.1.parquet"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "meta_df = pd.concat(frames)\n",
    "meta_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On systems where each file access is slow, it can be very useful to work with an in-memory dataframe (like this one) to avoid openeing files when it isn't necessary. For example, if you are just looking for the maximum value in column `'c'`, you can use the following approach to determine the file that needs to be parsed (in this case, `part.5.parquet`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<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>rgID</th>\n",
       "      <th>colID</th>\n",
       "      <th>cmin</th>\n",
       "      <th>cmax</th>\n",
       "      <th>file_path</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>5</td>\n",
       "      <td>c</td>\n",
       "      <td>3.81598e-06</td>\n",
       "      <td>0.999997</td>\n",
       "      <td>part.5.parquet</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   rgID colID         cmin      cmax       file_path\n",
       "2     5     c  3.81598e-06  0.999997  part.5.parquet"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Find maximum value for column == col_id\n",
    "col_id = 'c'\n",
    "c_max = meta_df.groupby(['colID']).max()['cmax'][col_id]\n",
    "mdf = meta_df[ meta_df['cmax'] == c_max ]\n",
    "mdf[ mdf['colID'] == col_id ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use the `_metadata` file to get a list of all files containing the partitions for out dataset.  This is useful when the file-system glob operations are slow, but is also critical when the data-files are not organized in a simple/predictible way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 0 ns, sys: 0 ns, total: 0 ns\n",
      "Wall time: 162 µs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['part.2.parquet',\n",
       " 'part.0.parquet',\n",
       " 'part.1.parquet',\n",
       " 'part.4.parquet',\n",
       " 'part.6.parquet',\n",
       " 'part.5.parquet',\n",
       " 'part.3.parquet',\n",
       " 'part.7.parquet']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time rg_files = list(set(meta_df['file_path']))\n",
    "rg_files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using this list of metadata files, we can do the reverse of the parallel *write* operation above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 52 ms, sys: 44 ms, total: 96 ms\n",
      "Wall time: 149 ms\n"
     ]
    }
   ],
   "source": [
    "def read_partition(part_path):\n",
    "    with open(part_path, mode=\"rb\") as f:\n",
    "        table = pq.read_table(\n",
    "            f,\n",
    "            use_pandas_metadata=True,\n",
    "            use_threads=False\n",
    "        )\n",
    "    df = table.to_pandas(use_threads=False)\n",
    "    return df\n",
    "\n",
    "def fread(args):\n",
    "    ipath = args[0]\n",
    "    fname = args[1]\n",
    "    dfrd = read_partition(ipath + '/' + fname)\n",
    "    return dfrd\n",
    "\n",
    "fargs = [ [path, ifile] for ifile in rg_files ]\n",
    "p = Pool(nprocs)\n",
    "%time df_part_rd = p.map(fread, fargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 68 ms, sys: 40 ms, total: 108 ms\n",
      "Wall time: 107 ms\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>a</th>\n",
       "      <th>b</th>\n",
       "      <th>c</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>D</td>\n",
       "      <td>0</td>\n",
       "      <td>0.626840</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>A</td>\n",
       "      <td>1</td>\n",
       "      <td>0.952662</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>B</td>\n",
       "      <td>2</td>\n",
       "      <td>0.765985</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>B</td>\n",
       "      <td>3</td>\n",
       "      <td>0.205139</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>C</td>\n",
       "      <td>4</td>\n",
       "      <td>0.116506</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   a  b         c\n",
       "0  D  0  0.626840\n",
       "1  A  1  0.952662\n",
       "2  B  2  0.765985\n",
       "3  B  3  0.205139\n",
       "4  C  4  0.116506"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time df_rd = pd.concat(df_part_rd).sort_index()\n",
    "df_rd.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that we wrote and read the same global dataset\n",
    "for col in df:\n",
    "    assert list(df[col]) == list(df_rd[col])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we can also use `pq.ParquetDataset` functionality to read/analyze the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8 ms, sys: 0 ns, total: 8 ms\n",
      "Wall time: 4.94 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'fs': <pyarrow.filesystem.LocalFileSystem at 0x7fc6caf86240>,\n",
       " 'paths': ['./test_dataset/part.2.parquet',\n",
       "  './test_dataset/part.0.parquet',\n",
       "  './test_dataset/part.1.parquet',\n",
       "  './test_dataset/part.4.parquet',\n",
       "  './test_dataset/part.6.parquet',\n",
       "  './test_dataset/part.5.parquet',\n",
       "  './test_dataset/part.3.parquet',\n",
       "  './test_dataset/part.7.parquet'],\n",
       " 'memory_map': True,\n",
       " 'pieces': [ParquetDatasetPiece('./test_dataset/part.2.parquet', row_group=None, partition_keys=[]),\n",
       "  ParquetDatasetPiece('./test_dataset/part.0.parquet', row_group=None, partition_keys=[]),\n",
       "  ParquetDatasetPiece('./test_dataset/part.1.parquet', row_group=None, partition_keys=[]),\n",
       "  ParquetDatasetPiece('./test_dataset/part.4.parquet', row_group=None, partition_keys=[]),\n",
       "  ParquetDatasetPiece('./test_dataset/part.6.parquet', row_group=None, partition_keys=[]),\n",
       "  ParquetDatasetPiece('./test_dataset/part.5.parquet', row_group=None, partition_keys=[]),\n",
       "  ParquetDatasetPiece('./test_dataset/part.3.parquet', row_group=None, partition_keys=[]),\n",
       "  ParquetDatasetPiece('./test_dataset/part.7.parquet', row_group=None, partition_keys=[])],\n",
       " 'partitions': None,\n",
       " 'common_metadata_path': None,\n",
       " 'metadata_path': None,\n",
       " 'common_metadata': None,\n",
       " 'metadata': <pyarrow._parquet.FileMetaData object at 0x7fc7963fb2c8>\n",
       "   created_by: parquet-cpp version 1.5.1-SNAPSHOT\n",
       "   num_columns: 3\n",
       "   num_rows: 1000000\n",
       "   num_row_groups: 8\n",
       "   format_version: 1.0\n",
       "   serialized_size: 2996,\n",
       " 'schema': <pyarrow._parquet.ParquetSchema object at 0x7fc6c7bae2b0>\n",
       " a: BYTE_ARRAY UTF8\n",
       " b: INT64\n",
       " c: DOUBLE\n",
       "  ,\n",
       " 'split_row_groups': False}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "myfiles = [path + '/' + f for f in rg_files]\n",
    "%time dataset = pq.ParquetDataset(myfiles, metadata = meta)\n",
    "# Note that 'common_metadata_path' and 'metadata_path' are only found\n",
    "# when the root directory of the dataset is passed in to the ParquetDataset constructor\n",
    "dataset.__dict__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `ParquetDataset` function takes in a file, directory or list of files, and parses the footer (metadata) in each of the files.  Although a `metadata` key-word argument is allowed, my understanding is that this only results in the schema information being used for validation purposes (not to avoid parsing the footer of every file). **In the future, the user should be allowed to skip the redundant parsing in the case that a metadata object is passed in.  However, this might require some refactoring in `arrow.parquet`.**\n",
    "\n",
    "The main outcome of calling `ParquetDataset` is that the user (or Dask) will then have access of `dataset.pieces`, which is a list of `pq.ParquetDatasetPiece` objects...."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ParquetDatasetPiece('./test_dataset/part.2.parquet', row_group=None, partition_keys=[]),\n",
       " ParquetDatasetPiece('./test_dataset/part.0.parquet', row_group=None, partition_keys=[]),\n",
       " ParquetDatasetPiece('./test_dataset/part.1.parquet', row_group=None, partition_keys=[]),\n",
       " ParquetDatasetPiece('./test_dataset/part.4.parquet', row_group=None, partition_keys=[]),\n",
       " ParquetDatasetPiece('./test_dataset/part.6.parquet', row_group=None, partition_keys=[]),\n",
       " ParquetDatasetPiece('./test_dataset/part.5.parquet', row_group=None, partition_keys=[]),\n",
       " ParquetDatasetPiece('./test_dataset/part.3.parquet', row_group=None, partition_keys=[]),\n",
       " ParquetDatasetPiece('./test_dataset/part.7.parquet', row_group=None, partition_keys=[])]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.pieces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, the only information needed to generate a `ParquetDatasetPiece` object (for each partition of the dataset) is the *path* of the partition. Therefore, we can also use the `_metadata` file information directly to generate each piece (independently) using the `pq.ParquetDatasetPiece` constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4 ms, sys: 0 ns, total: 4 ms\n",
      "Wall time: 1.21 ms\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>a</th>\n",
       "      <th>b</th>\n",
       "      <th>c</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>250000</th>\n",
       "      <td>C</td>\n",
       "      <td>250000</td>\n",
       "      <td>0.074626</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>250001</th>\n",
       "      <td>C</td>\n",
       "      <td>250001</td>\n",
       "      <td>0.678798</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>250002</th>\n",
       "      <td>C</td>\n",
       "      <td>250002</td>\n",
       "      <td>0.200794</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>250003</th>\n",
       "      <td>C</td>\n",
       "      <td>250003</td>\n",
       "      <td>0.507092</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>250004</th>\n",
       "      <td>B</td>\n",
       "      <td>250004</td>\n",
       "      <td>0.596354</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        a       b         c\n",
       "250000  C  250000  0.074626\n",
       "250001  C  250001  0.678798\n",
       "250002  C  250002  0.200794\n",
       "250003  C  250003  0.507092\n",
       "250004  B  250004  0.596354"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pieces = [pq.ParquetDatasetPiece(path + '/' + f) for f in rg_files] \n",
    "piece = pieces[0]\n",
    "%time md = piece.get_metadata(lambda fn: pq.ParquetFile(open(fn, mode=\"rb\")))\n",
    "with open(piece.path, mode=\"rb\") as f:\n",
    "    table = piece.read(\n",
    "        use_pandas_metadata=True,\n",
    "        file=f,\n",
    "        use_threads=False\n",
    "    )\n",
    "table.to_pandas().head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given the current state of `arrow.parquet`, it seems that `ParquetDataset` should only be used to open the entire dataset by `Dask` when the `_metadata` file is **not** present. When the metadata file **is** available, Dask can just allow each dataset piece to be created/read indepenently.\n",
    "\n",
    "However, the most ideal solution is likely to modify the arrow.parquet code to skip unnecessary I/O when the metadata is known. Therefore, a 1st pass at `dask.parquet` refactoring should probably use `ParquetDataset`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pre-`_metadata` version of the read phase..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8 ms, sys: 0 ns, total: 8 ms\n",
      "Wall time: 6.29 ms\n"
     ]
    }
   ],
   "source": [
    "%time dataset = pq.ParquetDataset(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 0 ns, sys: 0 ns, total: 0 ns\n",
      "Wall time: 967 µs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<pyarrow._parquet.FileMetaData object at 0x7fc7953d64f8>\n",
       "  created_by: parquet-cpp version 1.5.1-SNAPSHOT\n",
       "  num_columns: 3\n",
       "  num_rows: 125000\n",
       "  num_row_groups: 1\n",
       "  format_version: 1.0\n",
       "  serialized_size: 995"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pieces = dataset.pieces\n",
    "piece = pieces[0]\n",
    "%time md = piece.get_metadata(lambda fn: pq.ParquetFile(open(fn, mode=\"rb\")))\n",
    "columns = md.schema.names\n",
    "md"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown above, rather than reading the metadata from a dedicated metadata file, we can use the `dataset` obect to read the metadata from each piece (partition).  If we are interested in the row statistics of each partition, we can iterate trhough each file and generate our own `stats` structure. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'num-rows': 125000,\n",
       " 'columns': [{'name': 'a', 'min': b'A', 'max': b'D', 'null_count': 0},\n",
       "  {'name': 'b', 'min': 0, 'max': 124999, 'null_count': 0},\n",
       "  {'name': 'c', 'min': 5.44345e-07, 'max': 0.999996, 'null_count': 0}]}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "row_groups = [\n",
    "    piece.get_metadata(\n",
    "        lambda fn: pq.ParquetFile(open(fn, mode=\"rb\"))\n",
    "    ).row_group(0)\n",
    "    for piece in pieces\n",
    "]\n",
    "stats = []\n",
    "for row_group in row_groups:\n",
    "    s = {\"num-rows\": row_group.num_rows, \"columns\": []}\n",
    "    for i, name in enumerate(columns):\n",
    "        column = row_group.column(i)\n",
    "        d = {\"name\": name}\n",
    "        if column.statistics:\n",
    "            d.update(\n",
    "                {\n",
    "                    \"min\": column.statistics.min,\n",
    "                    \"max\": column.statistics.max,\n",
    "                    \"null_count\": column.statistics.null_count,\n",
    "                }\n",
    "            )\n",
    "        s[\"columns\"].append(d)\n",
    "    stats.append(s)\n",
    "stats[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `pieces` member of `ParquetDataset` class, it is straightforward to read back each partition of the dataset in parallel. Since we know the names of the columns and the *pieces* to iterate through, we can use the same `multiprocessing.Pool` approach that we used to write the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_partition(piece, columns):\n",
    "    with open(piece.path, mode=\"rb\") as f:\n",
    "        table = piece.read(\n",
    "            columns=columns,\n",
    "            use_pandas_metadata=True,\n",
    "            file=f,\n",
    "            use_threads=False\n",
    "        )\n",
    "    return df[list(columns)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 156 ms, sys: 136 ms, total: 292 ms\n",
      "Wall time: 452 ms\n"
     ]
    }
   ],
   "source": [
    "def fread(args):\n",
    "    piece = args[0]\n",
    "    columns = args[1]\n",
    "    dfrd = read_partition(piece, columns=columns)\n",
    "    return dfrd\n",
    "\n",
    "fargs = [ [piece, columns] for piece in pieces ]\n",
    "p = Pool(nprocs)\n",
    "%time df_part_rd = p.map(fread, fargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we can also confirm that each partition of `df_part_rd` is the same as `df_part`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert(df_part_rd[i] == df_part[i] for i in range(npart))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Warning: This document should be viewed as my own personal notes.  There may be serious mistakes/errors."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}