Skip to content

Instantly share code, notes, and snippets.

@rjzamora
Created September 7, 2021 19:14
Show Gist options
  • Select an option

  • Save rjzamora/a27dc1795eefc3c69951746ea71a4e62 to your computer and use it in GitHub Desktop.

Select an option

Save rjzamora/a27dc1795eefc3c69951746ea71a4e62 to your computer and use it in GitHub Desktop.
Parquet and FSSpec Experiments
Display the source blob
Display the rendered blob
Raw
{
"cells": [
{
"cell_type": "markdown",
"id": "c09578ad-cd37-46bb-98f9-3941f7cf13b0",
"metadata": {},
"source": [
"## Parquet Remote-Storage Experiments"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "92c5af6d-2040-4441-bed0-53d0f0d2114e",
"metadata": {},
"outputs": [],
"source": [
"# Test file (large file with 100 columns)\n",
"path = \"s3://my-bucket/large-file.parquet\"\n",
"\n",
"# Column and row-group selection\n",
"columns = [\"timestamp\", \"id10\"]\n",
"rgs = None"
]
},
{
"cell_type": "markdown",
"id": "d4b34ad3-bb53-448b-b0a6-39c560106a4b",
"metadata": {},
"source": [
"### Experiment 1 - \"Dummy\" Buffer Construction"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "b6e7a5c2-3559-4973-9650-7f4c94f1183b",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import io\n",
"import numpy as np\n",
"from fsspec.core import get_fs_token_paths\n",
"import cudf"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "bf3f774b-5546-4a31-a20e-ebf760056bf9",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 265 ms, sys: 44.8 ms, total: 309 ms\n",
"Wall time: 825 ms\n"
]
}
],
"source": [
"%%time\n",
"\n",
"# Step 0 - Get fs and size of file\n",
"fs, _, _ = get_fs_token_paths(path, mode=\"rb\")\n",
"file_size = fs.size(path)\n",
"\n",
"# Step 1 - Get 32 KB from tail of file.\n",
"#\n",
"# This \"sample size\" can be tunable, but should\n",
"# always be >= 8 bytes (so we can read the footer size)\n",
"tail_size = 32_000\n",
"footer_sample = fs.tail(path, tail_size)\n",
"\n",
"# Step 2 - Read the footer size and re-read a larger\n",
"# tail if necessary\n",
"footer_size = int.from_bytes(footer_sample[-8:-4], \"little\")\n",
"if tail_size < (footer_size + 8):\n",
" footer_sample = fs.tail(path, footer_size + 8)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "14c792fe-bcd2-4e1c-95d1-b0f50179a918",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 25 ms, sys: 3.46 ms, total: 28.5 ms\n",
"Wall time: 26 ms\n"
]
}
],
"source": [
"%%time\n",
"\n",
"# Step 3 - Collect required byte ranges\n",
"#\n",
"# Parsing allowed with both fastparquet and pyarrow\n",
"#\n",
"\n",
"def get_byte_ranges_pyarrow(footer_sample, columns=None, row_groups=None):\n",
" import pyarrow.parquet as pq\n",
"\n",
" byte_ranges = []\n",
" md = pq.ParquetFile(io.BytesIO(footer_sample)).metadata\n",
" for r in range(md.num_row_groups):\n",
" # Skip this row-group if we are targetting\n",
" # specific row-groups\n",
" if row_groups is None or r in row_groups:\n",
" row_group = md.row_group(r)\n",
" for c in range(row_group.num_columns):\n",
" column = row_group.column(c)\n",
" name = column.path_in_schema\n",
" # Skip this column if we are targetting a\n",
" # specific columns\n",
" if columns is None or name in columns:\n",
" file_offset0 = column.dictionary_page_offset\n",
" if file_offset0 is None:\n",
" file_offset0 = column.data_page_offset\n",
" num_bytes = column.total_uncompressed_size\n",
" byte_ranges.append((file_offset0, num_bytes))\n",
" return byte_ranges\n",
"\n",
"def get_byte_ranges_fastparquet(footer_sample, columns=None, row_groups=None):\n",
" import fastparquet as fp\n",
"\n",
" byte_ranges = []\n",
" pf = fp.ParquetFile(io.BytesIO(footer_sample))\n",
" for r, row_group in enumerate(pf.row_groups):\n",
" # Skip this row-group if we are targetting\n",
" # specific row-groups\n",
" if row_groups is None or r in row_groups:\n",
" for column in row_group.columns:\n",
" name = column.meta_data.path_in_schema[0]\n",
" # Skip this column if we are targetting a\n",
" # specific columns\n",
" if columns is None or name in columns:\n",
" file_offset0 = column.meta_data.dictionary_page_offset\n",
" if file_offset0 is None:\n",
" file_offset0 = column.meta_data.data_page_offset\n",
" num_bytes = column.meta_data.total_uncompressed_size\n",
" byte_ranges.append((file_offset0, num_bytes))\n",
" return byte_ranges\n",
"\n",
"byte_ranges = get_byte_ranges_pyarrow(footer_sample, columns=columns, row_groups=rgs)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "c5393467-f656-44b5-a7d5-21a7592b0bd1",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 5.35 s, sys: 1.04 s, total: 6.39 s\n",
"Wall time: 21.4 s\n"
]
}
],
"source": [
"%%time\n",
"\n",
"# Step 4 - Transfer the data for each byte range in parallel\n",
"#\n",
"# Use this data to construct a \"dummy\" buffer\n",
"#\n",
"\n",
"from queue import Queue\n",
"from threading import Thread\n",
"\n",
"def _assign_block(fs, path, local_buffer, offset, nbytes):\n",
" local_buffer[offset:offset+nbytes] = np.frombuffer(\n",
" fs.read_block(path, offset, nbytes),\n",
" dtype=\"b\"\n",
" )\n",
"\n",
"class ReadBlockWorker(Thread):\n",
"\n",
" def __init__(self, queue, fs, path, local_buffer):\n",
" Thread.__init__(self)\n",
" self.queue = queue\n",
" self.fs = fs\n",
" self.path = path\n",
" self.local_buffer = local_buffer\n",
"\n",
" def run(self):\n",
" while True:\n",
" # Get the work from the queue and expand the tuple\n",
" offset, nbytes = self.queue.get()\n",
" try:\n",
" _assign_block(self.fs, self.path, self.local_buffer, offset, nbytes)\n",
" finally:\n",
" self.queue.task_done()\n",
"\n",
" \n",
"# Start with empty buffer\n",
"buf = np.zeros(file_size, dtype=\"b\")\n",
"\n",
"# Column Chunks\n",
"def read_byte_ranges(ranges, local_buffer, num_threads=1):\n",
" \n",
" # No reason to generate more threads than byte-ranges\n",
" num_threads = min(num_threads, len(ranges))\n",
" \n",
" if num_threads > 1:\n",
" queue = Queue()\n",
" for x in range(num_threads):\n",
" worker = ReadBlockWorker(queue, fs, path, local_buffer)\n",
" worker.daemon = True\n",
" worker.start()\n",
" \n",
" for (offset, nbytes) in ranges:\n",
" if num_threads > 1:\n",
" queue.put((offset, nbytes))\n",
" else:\n",
" _assign_block(fs, path, local_buffer, offset, nbytes)\n",
"\n",
" if num_threads > 1:\n",
" queue.join()\n",
" \n",
"# Call multi-threaded data transfer of\n",
"# remote byte-ranges to local buffer\n",
"read_byte_ranges(byte_ranges, buf, num_threads=len(byte_ranges))\n",
"\n",
"# Add Header & Footer bytes\n",
"buf[:4] = np.frombuffer(b\"PAR1\", dtype=\"b\")\n",
"buf[-(footer_size+8):] = np.frombuffer(footer_sample[-(footer_size+8):], dtype=\"b\")"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "a0b79a04-e473-4a29-9c2f-e63384229c95",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 684 ms, sys: 564 ms, total: 1.25 s\n",
"Wall time: 1.24 s\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>id10</th>\n",
" </tr>\n",
" <tr>\n",
" <th>timestamp</th>\n",
" <th></th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>2000-01-01 00:00:00</th>\n",
" <td>974</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:01</th>\n",
" <td>990</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:02</th>\n",
" <td>1017</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:03</th>\n",
" <td>1035</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:04</th>\n",
" <td>956</td>\n",
" </tr>\n",
" <tr>\n",
" <th>...</th>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:56</th>\n",
" <td>980</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:57</th>\n",
" <td>981</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:58</th>\n",
" <td>1033</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:59</th>\n",
" <td>1042</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-31 00:00:00</th>\n",
" <td>970</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"<p>2592001 rows × 1 columns</p>\n",
"</div>"
],
"text/plain": [
" id10\n",
"timestamp \n",
"2000-01-01 00:00:00 974\n",
"2000-01-01 00:00:01 990\n",
"2000-01-01 00:00:02 1017\n",
"2000-01-01 00:00:03 1035\n",
"2000-01-01 00:00:04 956\n",
"... ...\n",
"2000-01-30 23:59:56 980\n",
"2000-01-30 23:59:57 981\n",
"2000-01-30 23:59:58 1033\n",
"2000-01-30 23:59:59 1042\n",
"2000-01-31 00:00:00 970\n",
"\n",
"[2592001 rows x 1 columns]"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"\n",
"# Step 5 - Read from the local \"dummy\" buffer\n",
"\n",
"cudf.read_parquet(buf.tobytes(), columns=columns, row_groups=rgs)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "13fcbdf2-ee1e-49a8-a4e4-7d5411fe6e7a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 10 s, sys: 2.48 s, total: 12.5 s\n",
"Wall time: 2min 39s\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>id10</th>\n",
" </tr>\n",
" <tr>\n",
" <th>timestamp</th>\n",
" <th></th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>2000-01-01 00:00:00</th>\n",
" <td>974</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:01</th>\n",
" <td>990</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:02</th>\n",
" <td>1017</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:03</th>\n",
" <td>1035</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:04</th>\n",
" <td>956</td>\n",
" </tr>\n",
" <tr>\n",
" <th>...</th>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:56</th>\n",
" <td>980</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:57</th>\n",
" <td>981</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:58</th>\n",
" <td>1033</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:59</th>\n",
" <td>1042</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-31 00:00:00</th>\n",
" <td>970</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"<p>2592001 rows × 1 columns</p>\n",
"</div>"
],
"text/plain": [
" id10\n",
"timestamp \n",
"2000-01-01 00:00:00 974\n",
"2000-01-01 00:00:01 990\n",
"2000-01-01 00:00:02 1017\n",
"2000-01-01 00:00:03 1035\n",
"2000-01-01 00:00:04 956\n",
"... ...\n",
"2000-01-30 23:59:56 980\n",
"2000-01-30 23:59:57 981\n",
"2000-01-30 23:59:58 1033\n",
"2000-01-30 23:59:59 1042\n",
"2000-01-31 00:00:00 970\n",
"\n",
"[2592001 rows x 1 columns]"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"\n",
"# Note that direct read is MUCH slower\n",
"\n",
"cudf.read_parquet(path, columns=columns, row_groups=rgs)"
]
},
{
"cell_type": "markdown",
"id": "f08264a7-db04-474c-b3e8-e086093df817",
"metadata": {},
"source": [
"**NOTES**:\n",
"\n",
"- The dummy-buffer approach is **much** faster when a small number of row-groups is selected from a file with many row-groups. However, the specific algorithm used here scales poorly when a large number of column-chunks is selected."
]
},
{
"cell_type": "markdown",
"id": "b00fd0c6-43fd-49a9-be0c-098ac8a172d2",
"metadata": {},
"source": [
"### Experiment 2 - Using a PyArrow-Native FileSystem\n",
"\n",
"**NOTE**: This experiment uses PyArrow to perform Parquet IO. It is also possible to pass an Arrow-based file object directly to libcudf, but this option is not yet plumbed into the python-cudf API."
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "ce8f89f5-6ed0-42e4-b254-cdb0e70d6e23",
"metadata": {},
"outputs": [],
"source": [
"from pyarrow import fs as pafs\n",
"import pyarrow.dataset as pads\n",
"\n",
"s3, pa_path = pafs.FileSystem.from_uri(path)\n",
"frag = next(pads.dataset(pa_path, filesystem=s3).get_fragments())"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "693442c2-14c4-440d-947b-30e1c1a42c7b",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 539 ms, sys: 219 ms, total: 758 ms\n",
"Wall time: 2.79 s\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>id10</th>\n",
" </tr>\n",
" <tr>\n",
" <th>timestamp</th>\n",
" <th></th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>2000-01-01 00:00:00</th>\n",
" <td>974</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:01</th>\n",
" <td>990</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:02</th>\n",
" <td>1017</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:03</th>\n",
" <td>1035</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:04</th>\n",
" <td>956</td>\n",
" </tr>\n",
" <tr>\n",
" <th>...</th>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:56</th>\n",
" <td>980</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:57</th>\n",
" <td>981</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:58</th>\n",
" <td>1033</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:59</th>\n",
" <td>1042</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-31 00:00:00</th>\n",
" <td>970</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"<p>2592001 rows × 1 columns</p>\n",
"</div>"
],
"text/plain": [
" id10\n",
"timestamp \n",
"2000-01-01 00:00:00 974\n",
"2000-01-01 00:00:01 990\n",
"2000-01-01 00:00:02 1017\n",
"2000-01-01 00:00:03 1035\n",
"2000-01-01 00:00:04 956\n",
"... ...\n",
"2000-01-30 23:59:56 980\n",
"2000-01-30 23:59:57 981\n",
"2000-01-30 23:59:58 1033\n",
"2000-01-30 23:59:59 1042\n",
"2000-01-31 00:00:00 970\n",
"\n",
"[2592001 rows x 1 columns]"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"\n",
"cudf.DataFrame.from_arrow(frag.to_table(columns=columns))"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "a406d1a0-a097-4809-b71d-1ba30ad9bb34",
"metadata": {},
"outputs": [],
"source": [
"frag = next(pads.dataset(path, filesystem=fs).get_fragments())"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "60249dbd-1879-4522-a488-ebf50661a0ea",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 4.2 s, sys: 727 ms, total: 4.93 s\n",
"Wall time: 56 s\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>id10</th>\n",
" </tr>\n",
" <tr>\n",
" <th>timestamp</th>\n",
" <th></th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>2000-01-01 00:00:00</th>\n",
" <td>974</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:01</th>\n",
" <td>990</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:02</th>\n",
" <td>1017</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:03</th>\n",
" <td>1035</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-01 00:00:04</th>\n",
" <td>956</td>\n",
" </tr>\n",
" <tr>\n",
" <th>...</th>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:56</th>\n",
" <td>980</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:57</th>\n",
" <td>981</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:58</th>\n",
" <td>1033</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-30 23:59:59</th>\n",
" <td>1042</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2000-01-31 00:00:00</th>\n",
" <td>970</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"<p>2592001 rows × 1 columns</p>\n",
"</div>"
],
"text/plain": [
" id10\n",
"timestamp \n",
"2000-01-01 00:00:00 974\n",
"2000-01-01 00:00:01 990\n",
"2000-01-01 00:00:02 1017\n",
"2000-01-01 00:00:03 1035\n",
"2000-01-01 00:00:04 956\n",
"... ...\n",
"2000-01-30 23:59:56 980\n",
"2000-01-30 23:59:57 981\n",
"2000-01-30 23:59:58 1033\n",
"2000-01-30 23:59:59 1042\n",
"2000-01-31 00:00:00 970\n",
"\n",
"[2592001 rows x 1 columns]"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"\n",
"cudf.DataFrame.from_arrow(frag.to_table(columns=columns))"
]
},
{
"cell_type": "markdown",
"id": "9d3c2499-0d12-4f85-ac8c-f29b59909174",
"metadata": {},
"source": [
"**NOTES**:\n",
"\n",
"- Using the native PyArrow s3 file object is much faster than using fsspec"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "247805db-19a1-47bd-b357-b386fd06080a",
"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.8.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment