arc_to_parquet/arc_to_parquet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/User/repos/functions/arc_to_parquet'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nuclio: ignore\n",
    "import nuclio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%nuclio config kind = \"job\"\n",
    "%nuclio config spec.image = \"mlrun/ml-base\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow.parquet as pq\n",
    "import pyarrow as pa\n",
    "\n",
    "from mlrun.execution import MLClientCtx\n",
    "from mlrun.datastore import DataItem\n",
    "\n",
    "from typing import List, Optional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _chunk_readwrite(\n",
    "    archive_url,\n",
    "    dest_path,\n",
    "    chunksize, \n",
    "    header, \n",
    "    encoding, \n",
    "    dtype\n",
    "):\n",
    "    \"\"\"stream read and write archives\n",
    "    \n",
    "    pandas reads and parquet writes\n",
    "    \n",
    "    notes\n",
    "    -----\n",
    "    * dest_path can be either a file.parquet, or in hte case of partitioned parquet\n",
    "      it will be only the destination folder of the parquet partition files\n",
    "    \"\"\"\n",
    "    pqwriter = None\n",
    "    header = []\n",
    "    for i, df in enumerate(pd.read_csv(archive_url, chunksize=chunksize, \n",
    "                                       names=header, encoding=encoding, \n",
    "                                       dtype=dtype)):\n",
    "        table = pa.Table.from_pandas(df)\n",
    "        if i == 0:\n",
    "            if dataset:\n",
    "                header = copy(table.schema)\n",
    "            else:\n",
    "                pqwriter = pq.ParquetWriter(dest_path, table.schema)\n",
    "        if dataset:\n",
    "            pq.write_to_dataset(table, root_path=dest_path, partition_cols=partition_cols)\n",
    "        else:\n",
    "            pqwriter.write_table(table)\n",
    "    if pqwriter:\n",
    "        pqwriter.close()\n",
    "    \n",
    "    return header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def arc_to_parquet(\n",
    "    context: MLClientCtx,\n",
    "    archive_url: DataItem,\n",
    "    header: List[str] = [None],\n",
    "    chunksize: int = 0,\n",
    "    dtype=None,\n",
    "    encoding: str = \"latin-1\",\n",
    "    key: str = \"data\",\n",
    "    dataset: str = \"None\",\n",
    "    part_cols = [],\n",
    "    file_ext: str = \"parquet\",\n",
    "    index: bool= False,\n",
    "    refresh_data: bool = False,\n",
    "    stats: bool = False\n",
    ") -> None:\n",
    "    \"\"\"Open a file/object archive and save as a parquet file or dataset\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    * this function is typically for large files, please be sure to check all settings\n",
    "    * partitioning requires precise specification of column types.\n",
    "    * the archive_url can be any file readable by pandas read_csv, which includes tar files\n",
    "    * if the `dataset` parameter is not empty, then a partitioned dataset will be created\n",
    "    instead of a single file in the folder `dataset`\n",
    "    * if a key exists already then it will not be re-acquired unless the `refresh_data` param\n",
    "    is set to `True`.  This is in case the original file is corrupt, or a refresh is \n",
    "    required.\n",
    "\n",
    "    :param context:        the function context\n",
    "    :param archive_url:    MLRun data input (DataItem object)\n",
    "    :param chunksize:      (0) when > 0, row size (chunk) to retrieve\n",
    "                           per iteration\n",
    "    :param dtype           destination data type of specified columns\n",
    "    :param encoding        (\"latin-8\") file encoding\n",
    "    :param key:            key in artifact store (when log_data=True)\n",
    "    :param dataset:        (None) if not None then \"target_path/dataset\"\n",
    "                           is folder for partitioned files\n",
    "    :param part_cols:      ([]) list of partitioning columns\n",
    "    :param file_ext:       (parquet) csv/parquet file extension\n",
    "    :param index:          (False) pandas save index option\n",
    "    :param refresh_data:   (False) overwrite existing data at that location\n",
    "    :param stats:          (None) calculate table stats when logging artifact\n",
    "    \"\"\"\n",
    "    base_path = context.artifact_path\n",
    "    os.makedirs(base_path, exist_ok=True)\n",
    "    \n",
    "    archive_url = archive_url.local()\n",
    "    \n",
    "    if dataset is not None:\n",
    "        dest_path = os.path.join(base_path, dataset)\n",
    "        exists = os.path.isdir(dest_path)\n",
    "    else:\n",
    "        dest_path = os.path.join(base_path, key+f\".{file_ext}\")\n",
    "        exists = os.path.isfile(dest_path)\n",
    "        \n",
    "    if not exists:\n",
    "        context.logger.info(\"destination file does not exist, downloading\")\n",
    "        if chunksize > 0:\n",
    "            header = _chunk_readwrite(archive_url, dest_path, chunksize,\n",
    "                                      encoding, dtype)\n",
    "            context.log_dataset(key=key, stats=stats, format='parquet', \n",
    "                                target_path=dest_path)\n",
    "        else:\n",
    "            df = pd.read_csv(archive_url)\n",
    "            context.log_dataset(key, df=df, format=file_ext, index=index)\n",
    "    else:\n",
    "        context.logger.info(\"destination file already exists, nothing done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nuclio: end-code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### mlconfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlrun import mlconf\n",
    "import os\n",
    "\n",
    "mlconf.dbpath = mlconf.dbpath or 'http://mlrun-api:8080'\n",
    "mlconf.artifact_path = mlconf.artifact_path or f'{os.environ[\"HOME\"]}/artifacts'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlrun import code_to_function \n",
    "# create job function object from notebook code\n",
    "fn = code_to_function('arc_to_parquet')\n",
    "\n",
    "# add metadata (for templates and reuse)\n",
    "fn.spec.default_handler = 'arc_to_parquet'\n",
    "fn.spec.description = \"retrieve remote archive, open and save as parquet\"\n",
    "fn.metadata.categories = ['data-movement', 'utils']\n",
    "fn.metadata.labels = {\"author\": \"yjb\"}\n",
    "fn.export('function.yaml')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"V3IO_HOME\" in os.environ:\n",
    "    from mlrun import mount_v3io\n",
    "    fn.apply(mount_v3io())\n",
    "else:\n",
    "    # is you set up mlrun using the instructions at https://github.com/mlrun/mlrun/blob/master/hack/local/README.md\n",
    "    from mlrun.platforms import mount_pvc\n",
    "    fn.apply(mount_pvc('nfsvol', 'nfsvol', '/home/joyan/data'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlrun import NewTask    \n",
    "\n",
    "task_params = {\"name\": \"tasks archive to parquet\", \n",
    "               \"params\":{\"key\": \"higgs-sample\"}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_URL = \"https://raw.githubusercontent.com/yjb-ds/testdata/master/data/higgs-sample.csv.gz\"\n",
    "# original large file \"https://archive.ics.uci.edu/ml/machine-learning-databases/00280/HIGGS.csv.gz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlrun import run_local\n",
    "\n",
    "run = run_local(NewTask(**task_params),\n",
    "          handler=arc_to_parquet,\n",
    "          inputs={\"archive_url\" : DATA_URL})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### the following will run quickly if your artifact path hasn't changed, the large file will be detected and not downloaded a second time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlrun import NewTask\n",
    "run = fn.run(NewTask(**task_params),\n",
    "             inputs={\"archive_url\" : DATA_URL},\n",
    "             artifact_path=mlconf.artifact_path)"
   ]
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}