{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e3c59492",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-09T01:58:42.167774Z",
     "start_time": "2022-11-09T01:58:41.631144Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running typecheck\n",
      "Success: no issues found in 1 source file\n",
      "\n",
      "running flake8\n",
      "flake8 success\n",
      "file size: 13 KB\n",
      "[2022-11-08 19:58:42.159 hdf5_repack] copying key1 0 of 1\n",
      "file size: 11 KB\n"
     ]
    }
   ],
   "source": [
    "%%checkall\n",
    "from __future__ import annotations\n",
    "import h5py\n",
    "import string\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import datetime\n",
    "from typing import Any\n",
    "from pyqstrat.pq_utils import get_temp_dir, get_child_logger, assert_\n",
    "\n",
    "_logger = get_child_logger(__name__)\n",
    "\n",
    "\n",
    "def np_arrays_to_hdf5(data: dict[str, np.ndarray], \n",
    "                      filename: str, \n",
    "                      key: str, \n",
    "                      dtypes: dict[str, str] | None = None, \n",
    "                      as_utf8: list[str] | None = None,\n",
    "                      compression_args: dict[Any, Any] | None = None) -> None:\n",
    "    '''\n",
    "    Write a list of numpy arrays to hdf5\n",
    "    Args:\n",
    "        data: list of numpy one dimensional arrays along with the name of the array\n",
    "        filename: filename of the hdf5 file\n",
    "        key: group and or / subgroups to write to.  For example, \"g1/g2\" will write to the subgrp g2 within the grp g1\n",
    "        dtypes: dict used to override datatype for a column.  For example, {\"col1\": \"f4\"} will write a 4 byte float array for col1\n",
    "        as_utf_8: each column listed here will be saved with utf8 encoding. For all other strings, we will compute the max length\n",
    "        and store as a fixed length byte array with ascii encoding, i.e. a S[max length] datatype. This is much faster to read and process\n",
    "        compression_args: if you want to compress the hdf5 file. You can use the hdf5plugin module and arguments such as hdf5plugin.Blosc()\n",
    "    '''\n",
    "    if not len(data): return\n",
    "    tmp_key = key + '_tmp'\n",
    "    \n",
    "    if compression_args is None:\n",
    "        compression_args = {}\n",
    "        \n",
    "    if as_utf8 is None:\n",
    "        as_utf8 = []\n",
    "    \n",
    "    with h5py.File(filename, 'a') as f:\n",
    "        if tmp_key in f: del f[tmp_key]\n",
    "        grp = f.create_group(tmp_key)\n",
    "        for colname, array in data.items():\n",
    "            if dtypes is not None and colname in dtypes:\n",
    "                dtype = np.dtype(dtypes[colname])\n",
    "                if dtype.kind == 'M':  # datetime\n",
    "                    dtype = h5py.opaque_dtype(dtype)\n",
    "                    array = array.astype(dtype)\n",
    "                else:\n",
    "                    array = array.astype(dtype)\n",
    "            else:  # we need to figure out datatype\n",
    "                dtype = array.dtype\n",
    "                if colname in as_utf8:\n",
    "                    array = np.char.encode(array.astype('U'), 'utf-8')\n",
    "                elif dtype.kind == 'O':\n",
    "                    array = np.where(array == None, '', array)  # noqa: E711 comparison to None should be 'if cond is None:'\n",
    "                    array = array.astype('S')\n",
    "                elif dtype.kind == 'M':  # datetime\n",
    "                    dtype = h5py.opaque_dtype(dtype)\n",
    "                    array = array.astype(dtype)\n",
    "            if colname in grp:\n",
    "                del grp[colname]\n",
    "            grp.create_dataset(name=colname, data=array, shape=[len(array)], dtype=array.dtype, **compression_args)\n",
    "            \n",
    "        grp.attrs['type'] = 'dataframe'\n",
    "        grp.attrs['timestamp'] = str(datetime.datetime.now())\n",
    "        grp.attrs['rows'] = len(array)\n",
    "        grp.attrs['columns'] = ','.join([tup[0] for tup in data])\n",
    "        grp.attrs['utf8_cols'] = ','.join(as_utf8)\n",
    "\n",
    "        if key in f: \n",
    "            del f[key]\n",
    "        f.move(tmp_key, key)\n",
    "        f.flush()\n",
    "        \n",
    "\n",
    "def hdf5_to_np_arrays(filename: str, key: str) -> dict[str, np.ndarray]:\n",
    "    '''\n",
    "    Read a list of numpy arrays previously written out by np_arrays_to_hdf5\n",
    "    Args:\n",
    "        filename: path of the hdf5 file to read\n",
    "        key: group and or / subgroups to read from.  For example, \"g1/g2\" will read from the subgrp g2 within the grp g1\n",
    "    Return:\n",
    "        a list of numpy arrays along with their names\n",
    "        '''\n",
    "    ret: dict[str, np.ndarray] = {}\n",
    "    with h5py.File(filename, 'r') as f:\n",
    "        assert_(key in f, f'{key} not found in {filename}')\n",
    "        grp = f[key]\n",
    "        assert_('type' in grp.attrs and grp.attrs['type'] == 'dataframe', f'{key} not a dataframe')\n",
    "        columns = grp.attrs['columns'].split(',')\n",
    "        utf8_cols: list[str] = []\n",
    "        if 'utf8_cols' in grp.attrs:\n",
    "            utf8_cols = grp.attrs['utf8_cols'].split(',')\n",
    "        for col in columns:\n",
    "            array = grp[col][:]\n",
    "            if col in utf8_cols:\n",
    "                array = np.char.decode(array, 'utf-8')\n",
    "                dtype = f'U{array.dtype.itemsize}'\n",
    "            if array.dtype.kind == 'S':\n",
    "                # decode bytes to numpy unicode\n",
    "                dtype = f'U{array.dtype.itemsize}'\n",
    "                array = array.astype(dtype)\n",
    "            elif array.dtype == 'O':\n",
    "                array = array.astype('S')\n",
    "                array = np.char.decode(array, encoding='utf-8')\n",
    "            ret[col] = array\n",
    "    return ret\n",
    "        \n",
    "        \n",
    "def df_to_hdf5(df: pd.DataFrame, \n",
    "               filename: str, \n",
    "               key: str, \n",
    "               dtypes: dict[str, str] | None = None,\n",
    "               as_utf8: list[str] | None = None) -> None:\n",
    "    '''\n",
    "    Write out a pandas dataframe to hdf5 using the np_arrays_to_hdf5 function\n",
    "    '''\n",
    "    arrays: dict[str, np.ndarray] = {}\n",
    "    for column in df.columns:\n",
    "        arrays[column] = df[column].values\n",
    "    np_arrays_to_hdf5(arrays, filename, key, dtypes, as_utf8)\n",
    "    \n",
    "\n",
    "def hdf5_to_df(filename: str, key: str) -> pd.DataFrame:\n",
    "    '''\n",
    "    Read a pandas dataframe previously written out using df_to_hdf5 or np_arrays_to_hdf5\n",
    "    '''\n",
    "    arrays = hdf5_to_np_arrays(filename, key)\n",
    "    array_dict = {name: array for name, array in arrays.items()}\n",
    "    return pd.DataFrame(array_dict)\n",
    "\n",
    "\n",
    "def hdf5_repack(inp_filename: str, out_filename: str) -> None:\n",
    "    '''\n",
    "    Copy groups from input filename to output filename.\n",
    "    Serves the same purpose as the h5repack command line tool, i.e. \n",
    "    discards empty space in the input file so the output file may be smaller\n",
    "    '''\n",
    "    with h5py.File(inp_filename, 'r') as inpf:\n",
    "        num_items = len(list(inpf.keys()))\n",
    "        with h5py.File(out_filename + '.tmp', 'w') as outf:\n",
    "            for i, name in enumerate(inpf):\n",
    "                _logger.info(f'copying {name} {i} of {num_items}')\n",
    "                inpf.copy(inpf[name], outf)\n",
    "    os.rename(out_filename + '.tmp', out_filename)\n",
    "\n",
    "\n",
    "def test_hdf5_to_df():\n",
    "    size = int(100)\n",
    "    a = np.random.randint(0, 10000, size)\n",
    "    b = a * 1.1\n",
    "    letters = np.random.choice(list(string.ascii_letters), (size, 5))\n",
    "    c = np.empty(size, dtype='O')\n",
    "    for i, row in enumerate(letters):\n",
    "        c[i] = ''.join(row)\n",
    "    c[1] = None\n",
    "    d = (a * 1000).astype('M8[m]')\n",
    "    temp_dir = get_temp_dir()\n",
    "    # os.remove(f'{temp_dir}/test.hdf5')\n",
    "    if os.path.isfile(f'{temp_dir}/test.hdf5'): os.remove(f'{temp_dir}/test.hdf5')\n",
    "    np_arrays_to_hdf5({\"b\": b, \"a\": a, \"c\": c, \"d\": d}, f'{temp_dir}/test.hdf5', 'key1/key2')\n",
    "    file_size = os.path.getsize(f'{temp_dir}/test.hdf5')\n",
    "    print(f\"file size: {file_size / 1e3:.0f} KB\")\n",
    "    \n",
    "    hdf5_repack(f'{temp_dir}/test.hdf5', f'{temp_dir}/test.hdf5.tmp')\n",
    "    if os.path.isfile(f'{temp_dir}/test.hdf5'): os.remove(f'{temp_dir}/test.hdf5')\n",
    "    os.rename(f'{temp_dir}/test.hdf5.tmp', f'{temp_dir}/test.hdf5')\n",
    "    file_size = os.path.getsize(f'{temp_dir}/test.hdf5')\n",
    "    print(f\"file size: {file_size / 1e3:.0f} KB\")\n",
    "    assert_(file_size > 10000 and file_size < 14000, f'invalid file size: {file_size}')\n",
    "\n",
    "    if os.path.isfile(f'{temp_dir}/test.hdf5'): os.remove(f'{temp_dir}/test.hdf5')\n",
    "    df_in = pd.DataFrame(dict(a=a, b=b, c=c, d=d))\n",
    "    df_to_hdf5(df_in, f'{temp_dir}/test.hdf5', 'key1/key2', dtypes={'d': 'M8[m]'})\n",
    "    df_out = hdf5_to_df(f'{temp_dir}/test.hdf5', 'key1/key2')\n",
    "    df_out.c = np.where(df_out.c == '', None, df_out.c)\n",
    "    from pandas.testing import assert_frame_equal\n",
    "    assert_frame_equal(df_in, df_out)\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    test_hdf5_to_df()"
   ]
  }
 ],
 "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.9.13"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
