Skip to content

Instantly share code, notes, and snippets.

@rjzamora

rjzamora/read_file_list.ipynb

Created October 30, 2020 20:02
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save rjzamora/a8bb8e752882326ede3e86062f8111b6 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save rjzamora/a8bb8e752882326ede3e86062f8111b6 to your computer and use it in GitHub Desktop.
Download ZIP
Simple example of mapping multiple parquet files to each dask_cudf.DataFrame partition.
Display the source blob
Display the rendered blob
Raw
read_file_list.ipynb
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Partitioning by groups of files with `read_parquet`\n",
"\n",
"Below is a simple example of a `dask_cudf.read_parquet`-like API for mapping groups of files to `dask_cudf.DataFrame` partitions. Note that the critical `kwarg` is `chunksize`, which specifies the total **on-disk** file size to map to each output partition. Keep in mind that the on-disk size may be a lot smaller than the in memory size if there is compression.\n",
"\n",
"The experimental API is called `read_parquet_files`:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import cudf\n",
"from dask.base import tokenize\n",
"from dask.dataframe.core import new_dd_object\n",
"from dask.utils import parse_bytes\n",
"from fsspec.core import get_fs_token_paths\n",
"import pyarrow.parquet as pq\n",
"\n",
"def read_parquet_files(paths, chunksize=None, columns=None, storage_options=None):\n",
"\n",
" chunksize = parse_bytes(chunksize) if chunksize else parse_bytes(\"25MB\")\n",
" chunksize = parse_bytes(chunksize)\n",
"\n",
" storage_options = storage_options or {}\n",
" fs, fs_token, _ = get_fs_token_paths(paths, mode=\"rb\", storage_options=storage_options)\n",
" \n",
" dataset = pq.ParquetDataset(paths, filesystem=fs, validate_schema=False)\n",
" pieces = generate_pieces(dataset, fs, chunksize)\n",
"\n",
" name = \"parquet-to-ddf-\" + tokenize(fs_token, pieces, columns)\n",
" dsk = {\n",
" (name, p): (read_part, piece, columns)\n",
" for p, piece in enumerate(pieces)\n",
" }\n",
"\n",
" meta = cudf.read_parquet(dataset.pieces[0].path, num_rows=5, columns=columns).iloc[:0]\n",
" divisions = [None] * (len(pieces) + 1)\n",
"\n",
" return new_dd_object(dsk, name, meta, divisions)\n",
"\n",
"\n",
"def generate_pieces(dataset, fs, chunksize):\n",
" pieces = []\n",
" item = []\n",
" size = 0\n",
" for piece in dataset.pieces:\n",
" path = piece.path\n",
" file_size = fs.du(path)\n",
" if size + file_size <= chunksize:\n",
" # This piece can be added to the item list\n",
" size += file_size\n",
" item.append(path)\n",
" else:\n",
" # This piece is too big to add to the item list.\n",
" # We need to flush whatever we already have in\n",
" # `item` before adding this piece.\n",
" if size:\n",
" pieces.append(item)\n",
" size = file_size\n",
" item = [path]\n",
" if size:\n",
" pieces.append(item)\n",
" return pieces\n",
"\n",
"\n",
"def read_part(paths, columns):\n",
" return cudf.read_parquet(paths, columns=columns)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Simple Use Example"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from dask.datasets import timeseries\n",
"\n",
"# Write test dataset\n",
"path = \"./test.parquet\"\n",
"ddf = timeseries().reset_index(drop=True)\n",
"ddf.to_parquet(path)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Note that the test dataset contains 30 files (because the DataFrame has 30 partitions)."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div><strong>Dask DataFrame Structure:</strong></div>\n",
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>id</th>\n",
" <th>name</th>\n",
" <th>x</th>\n",
" <th>y</th>\n",
" </tr>\n",
" <tr>\n",
" <th>npartitions=30</th>\n",
" <th></th>\n",
" <th></th>\n",
" <th></th>\n",
" <th></th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th></th>\n",
" <td>int64</td>\n",
" <td>object</td>\n",
" <td>float64</td>\n",
" <td>float64</td>\n",
" </tr>\n",
" <tr>\n",
" <th></th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>...</th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th></th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th></th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>\n",
"<div>Dask Name: reset_index, 60 tasks</div>"
],
"text/plain": [
"Dask DataFrame Structure:\n",
" id name x y\n",
"npartitions=30 \n",
" int64 object float64 float64\n",
" ... ... ... ...\n",
"... ... ... ... ...\n",
" ... ... ... ...\n",
" ... ... ... ...\n",
"Dask Name: reset_index, 60 tasks"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ddf"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"9.5K\t./test.parquet/_common_metadata\n",
"19K\t./test.parquet/_metadata\n",
"1.9M\t./test.parquet/part.0.parquet\n",
"1.9M\t./test.parquet/part.10.parquet\n",
"1.9M\t./test.parquet/part.11.parquet\n",
"1.9M\t./test.parquet/part.12.parquet\n",
"1.9M\t./test.parquet/part.13.parquet\n",
"1.9M\t./test.parquet/part.14.parquet\n",
"1.9M\t./test.parquet/part.15.parquet\n",
"1.9M\t./test.parquet/part.16.parquet\n",
"1.9M\t./test.parquet/part.17.parquet\n",
"1.9M\t./test.parquet/part.18.parquet\n",
"1.9M\t./test.parquet/part.19.parquet\n",
"1.9M\t./test.parquet/part.1.parquet\n",
"1.9M\t./test.parquet/part.20.parquet\n",
"1.9M\t./test.parquet/part.21.parquet\n",
"1.9M\t./test.parquet/part.22.parquet\n",
"1.9M\t./test.parquet/part.23.parquet\n",
"1.9M\t./test.parquet/part.24.parquet\n",
"1.9M\t./test.parquet/part.25.parquet\n",
"1.9M\t./test.parquet/part.26.parquet\n",
"1.9M\t./test.parquet/part.27.parquet\n",
"1.9M\t./test.parquet/part.28.parquet\n",
"1.9M\t./test.parquet/part.29.parquet\n",
"1.9M\t./test.parquet/part.2.parquet\n",
"1.9M\t./test.parquet/part.3.parquet\n",
"1.9M\t./test.parquet/part.4.parquet\n",
"1.9M\t./test.parquet/part.5.parquet\n",
"1.9M\t./test.parquet/part.6.parquet\n",
"1.9M\t./test.parquet/part.7.parquet\n",
"1.9M\t./test.parquet/part.8.parquet\n",
"1.9M\t./test.parquet/part.9.parquet\n"
]
}
],
"source": [
"!du -sh ./test.parquet/*"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"ddf_read = read_parquet_files(path, chunksize=\"7MB\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Note that we now have only 10 partitions after reading back, because we set `chunksize` to be a bit larger than 3x the size of each file."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div><strong>Dask DataFrame Structure:</strong></div>\n",
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>id</th>\n",
" <th>name</th>\n",
" <th>x</th>\n",
" <th>y</th>\n",
" </tr>\n",
" <tr>\n",
" <th>npartitions=10</th>\n",
" <th></th>\n",
" <th></th>\n",
" <th></th>\n",
" <th></th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th></th>\n",
" <td>int64</td>\n",
" <td>object</td>\n",
" <td>float64</td>\n",
" <td>float64</td>\n",
" </tr>\n",
" <tr>\n",
" <th></th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>...</th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th></th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th></th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>\n",
"<div>Dask Name: parquet-to-ddf, 10 tasks</div>"
],
"text/plain": [
"<dask_cudf.DataFrame | 10 tasks | 10 npartitions>"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ddf_read"
]
},
{
"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>id</th>\n",
" <th>name</th>\n",
" <th>x</th>\n",
" <th>y</th>\n",
" </tr>\n",
" <tr>\n",
" <th>__null_dask_index__</th>\n",
" <th></th>\n",
" <th></th>\n",
" <th></th>\n",
" <th></th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>1036</td>\n",
" <td>Edith</td>\n",
" <td>-0.735192</td>\n",
" <td>-0.223191</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>989</td>\n",
" <td>Wendy</td>\n",
" <td>-0.619498</td>\n",
" <td>-0.123222</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>1025</td>\n",
" <td>Hannah</td>\n",
" <td>-0.878546</td>\n",
" <td>-0.043274</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>1000</td>\n",
" <td>Edith</td>\n",
" <td>-0.284171</td>\n",
" <td>-0.514583</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <td>1028</td>\n",
" <td>Frank</td>\n",
" <td>-0.336215</td>\n",
" <td>-0.763271</td>\n",
" </tr>\n",
" <tr>\n",
" <th>...</th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>86395</th>\n",
" <td>986</td>\n",
" <td>George</td>\n",
" <td>0.593016</td>\n",
" <td>-0.595666</td>\n",
" </tr>\n",
" <tr>\n",
" <th>86396</th>\n",
" <td>977</td>\n",
" <td>Frank</td>\n",
" <td>0.437606</td>\n",
" <td>-0.565509</td>\n",
" </tr>\n",
" <tr>\n",
" <th>86397</th>\n",
" <td>1020</td>\n",
" <td>Zelda</td>\n",
" <td>0.883307</td>\n",
" <td>0.919918</td>\n",
" </tr>\n",
" <tr>\n",
" <th>86398</th>\n",
" <td>1031</td>\n",
" <td>Michael</td>\n",
" <td>0.684471</td>\n",
" <td>0.485971</td>\n",
" </tr>\n",
" <tr>\n",
" <th>86399</th>\n",
" <td>985</td>\n",
" <td>Quinn</td>\n",
" <td>0.581061</td>\n",
" <td>-0.720021</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"<p>2592000 rows × 4 columns</p>\n",
"</div>"
],
"text/plain": [
" id name x y\n",
"__null_dask_index__ \n",
"0 1036 Edith -0.735192 -0.223191\n",
"1 989 Wendy -0.619498 -0.123222\n",
"2 1025 Hannah -0.878546 -0.043274\n",
"3 1000 Edith -0.284171 -0.514583\n",
"4 1028 Frank -0.336215 -0.763271\n",
"... ... ... ... ...\n",
"86395 986 George 0.593016 -0.595666\n",
"86396 977 Frank 0.437606 -0.565509\n",
"86397 1020 Zelda 0.883307 0.919918\n",
"86398 1031 Michael 0.684471 0.485971\n",
"86399 985 Quinn 0.581061 -0.720021\n",
"\n",
"[2592000 rows x 4 columns]"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ddf_read.compute()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"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.8"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment