{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bf9f48f-cbf9-42fe-8d86-ef5389971dee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from df_cereal.df_widget import DFWidget, BaseWidget\n",
    "from df_cereal.arrow_utils import df_to_arrow_bytes, df_to_base64\n",
    "from traitlets import Unicode, Dict, observe, Any, Bytes, observe, Bool\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "import time\n",
    "from buckaroo import BuckarooWidget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0b25468-fba4-4995-b639-81520a066791",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def second_convert(ser):\n",
    "    return ser.dt.seconds + ser.dt.microseconds / (10**6)\n",
    "\n",
    "class BaseBenchmarkWidget(BaseWidget):\n",
    "    _view_name = Unicode('BytesBenchmarkWidgetView').tag(sync=True)\n",
    "    \n",
    "    #timing_info must be empty and do_calc must be True to trigger the benchmark, \n",
    "    #we use this to prevent loops, and to allow only calculating the calc time not the widget instantiation time\n",
    "    timing_info = Dict({}).tag(sync=True)  \n",
    "    do_calc = Bool(True).tag(sync=True)\n",
    "    def convert_df(self, df):\n",
    "        pass\n",
    "    target_prop = None\n",
    "\n",
    "    def __init__(self, df_gen_func, min_n, max_n, **kwargs):\n",
    "        self.results_arr = []\n",
    "        self.max_n = max_n\n",
    "        self.current_n = min_n\n",
    "        self.df_gen_func = df_gen_func\n",
    "        self.bw = BuckarooWidget(pd.DataFrame(), pinned_rows=[])\n",
    "        display(self.bw)\n",
    "        super().__init__()\n",
    "        self.set_from_gen()\n",
    "        \n",
    "    def __del__(self):\n",
    "        print(\"in __del__\")\n",
    "        self.bw.close()\n",
    "        del self.bw\n",
    "        return super().__del__()\n",
    "    \n",
    "    def format_results(self):\n",
    "        tdf = pd.DataFrame(self.results_arr)\n",
    "        tdf['klass'] = self.__class__.__name__.replace(\"BenchmarkWidget\",\"\")\n",
    "        tdf['py_serialize'] = second_convert(tdf['serial_t2'] - tdf['serial_t1'])\n",
    "        tdf['widget_xfer'] =  second_convert(tdf['t1'] - tdf['serial_t2'] )\n",
    "        tdf['js_deserialize'] = second_convert(tdf['t2'] - tdf['t1'])\n",
    "        tdf['js_iterate'] = second_convert(tdf['t3'] - tdf['t2'])\n",
    "        tdf['py_notify'] = second_convert(tdf['t4'] - tdf['t3'])\n",
    "        tdf['roundtrip'] = second_convert(tdf['t4'] - tdf['serial_t3'])\n",
    "        self.bw.raw_df = tdf[['rows', 'bytes', 'elements', 'klass', 'py_serialize', 'widget_xfer', 'js_deserialize', 'js_iterate', 'py_notify', 'roundtrip']]\n",
    "    \n",
    "    def set_from_gen(self):\n",
    "        self.do_calc = False # this will trigger view re-evaluate, while doing nothing\n",
    "        self.timing_info = {} # this will trigger view re-evaluate while doing nothing\n",
    "        self.current_timing = {'gen_t1': pd.Timestamp.utcnow()}\n",
    "        gen_df = self.df_gen_func(self.current_n)\n",
    "        self.current_timing['serial_t1'] = pd.Timestamp.utcnow() # START PY_SERIALIZE\n",
    "        self.current_timing['rows'] = len(gen_df)\n",
    "        self.current_timing['elements'] = len(gen_df) * len(gen_df.columns)\n",
    "        byts, self.current_timing['bytes'] = self.convert_df(gen_df) \n",
    "        self.current_timing['serial_t2'] = pd.Timestamp.utcnow()  #END PY_SERIALIZE\n",
    "        setattr(self, self.target_prop, byts) # trigger transfer\n",
    "        self.do_calc = True # trigger js_timing loop (t1, t2, t3)\n",
    "\n",
    "        self.current_timing['serial_t3'] = pd.Timestamp.utcnow()\n",
    "    \n",
    "    @observe('timing_info')\n",
    "    def _timing_info_change(self, change):\n",
    "        if len(self.timing_info) == 0:\n",
    "            print(\"exit early because empty timing_info\")\n",
    "            return\n",
    "        new_res = { k: pd.Timestamp(v) for k,v in self.timing_info.items()}\n",
    "        new_res['t4'] = pd.Timestamp.utcnow()\n",
    "        new_res.update(self.current_timing)\n",
    "        self.results_arr.append(new_res)\n",
    "        if self.current_n > self.max_n:\n",
    "            print(\"done\")\n",
    "            self.format_results()\n",
    "            return\n",
    "        else:\n",
    "            self.do_calc = False \n",
    "            self.timing_info = {}\n",
    "            self.set_from_gen()\n",
    "            self.current_n += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c15e221-0212-4858-a289-2850d06c67f3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Base64BenchmarkWidget(BaseBenchmarkWidget):\n",
    "    _view_name = Unicode('Base64BenchmarkWidgetView').tag(sync=True)\n",
    "    df_base64 = Any(\"\").tag(sync=True)\n",
    "    target_prop = 'df_base64'\n",
    "    \n",
    "    def convert_df(self, df):\n",
    "        txt = df_to_base64(df)\n",
    "        return txt, len(txt)\n",
    "class BytesBenchmarkWidget(BaseBenchmarkWidget):\n",
    "    _view_name = Unicode('BytesBenchmarkWidgetView').tag(sync=True)\n",
    "\n",
    "    def convert_df(self, df):\n",
    "        byts = df_to_arrow_bytes(df)\n",
    "        return byts, len(byts)\n",
    "    df_arrow_bytes = Bytes().tag(sync=True)\n",
    "    target_prop = 'df_arrow_bytes'\n",
    "\n",
    "class JSONDFBenchmarkWidget(BaseBenchmarkWidget):\n",
    "    _view_name = Unicode('DFDataBenchmarkWidgetView').tag(sync=True)\n",
    "\n",
    "    def convert_df(self, df):\n",
    "        txt = df.to_json(orient='records')\n",
    "        return json.loads(txt), len(txt)\n",
    "    df_data = Any([]).tag(sync=True)\n",
    "    target_prop = 'df_data'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1f7de35-635e-4931-b163-e8a2b8fb13e7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def df_gen_func(exp):\n",
    "    N = 20 * (10**exp)//2\n",
    "    big_df = df = pd.DataFrame({\n",
    "        'str': [\"foo\", \"barr\", \"bazz\", \"bofffff\"] * (N//4),\n",
    "        'log_normal': np.random.lognormal(25, .3, N),\n",
    "        'bool': [True, False] * (N//2),\n",
    "        'randint': np.random.randint(0,100,N)\n",
    "    })\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6d32ae2-581b-45db-816e-53fa2b4e683b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "all_obs_dfs = []\n",
    "b64_bn = Base64BenchmarkWidget(df_gen_func, min_n=1, max_n=5)\n",
    "b64_bn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a443f732-3b99-4d52-9f28-fe50e56f000b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "all_obs_dfs.append(b64_bn.bw.raw_df)\n",
    "b64_bn.bw.close()\n",
    "del b64_bn.bw\n",
    "b64_bn.close()\n",
    "del b64_bn\n",
    "byt_bn = BytesBenchmarkWidget(df_gen_func, min_n=1, max_n=6)\n",
    "byt_bn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a1540f0-75a8-4f9a-87c9-6cbc3efb4382",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "all_obs_dfs.append(byt_bn.bw.raw_df)\n",
    "byt_bn.bw.close()\n",
    "del byt_bn.bw\n",
    "byt_bn.close()\n",
    "del byt_bn\n",
    "jsdf_bn = JSONDFBenchmarkWidget(df_gen_func, min_n=1, max_n=4)\n",
    "jsdf_bn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0b9f6aa-8489-4c82-9f97-ed22857e38eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_obs_dfs.append(jsdf_bn.bw.raw_df)\n",
    "jsdf_bn.bw.close()\n",
    "del jsdf_bn.bw\n",
    "jsdf_bn.close()\n",
    "del jsdf_bn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c69302-6a98-4d74-ba30-2e913dd9c25d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def without(arr, removs):\n",
    "    new_arr = []\n",
    "    for k in arr:\n",
    "        if k in removs:\n",
    "            continue\n",
    "        new_arr.append(k)\n",
    "    return new_arr\n",
    "def reorder(df, first_cols):\n",
    "    all_cols = df.columns\n",
    "    rest_columns = without(all_cols, first_cols)\n",
    "    ord_cols = first_cols.copy()\n",
    "    ord_cols.extend(rest_columns)\n",
    "    return df[ord_cols]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5050d68-5346-4db5-aa9c-410a379d8dd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def df_gen_func(exp):  # redefining here for nice screenshots\n",
    "    N = 20 * (10**exp)//2\n",
    "    big_df = df = pd.DataFrame({\n",
    "        'str': [\"foo\", \"barr\", \"bazz\", \"bofffff\"] * (N//4),\n",
    "        'log_normal': np.random.lognormal(25, .3, N),\n",
    "        'bool': [True, False] * (N//2),\n",
    "        'randint': np.random.randint(0,100,N)\n",
    "    })\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a21ef01-262b-444b-8eac-ffc484a2cf3f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "full_df = pd.concat(all_obs_dfs)\n",
    "full_df['els_per_second'] = full_df['elements'] / full_df['roundtrip']\n",
    "full_df['bytes_per_second_full'] = full_df['bytes'] / full_df['roundtrip']\n",
    "full_df['bytes_per_second_xfer'] = full_df['bytes'] / full_df['widget_xfer']\n",
    "\n",
    "reorder(full_df, ['bytes', 'elements', 'klass', 'els_per_second', 'bytes_per_second_xfer'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78b94cf9-ae1a-499a-a64c-330ad5171a7d",
   "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.18"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
