{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8bf9f48f-cbf9-42fe-8d86-ef5389971dee",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Buckaroo has been enabled as the default DataFrame viewer.  To return to default dataframe visualization use `from buckaroo import disable; disable()`\n"
     ]
    }
   ],
   "source": [
    "from df_cereal.df_widget import DFWidget, BaseWidget\n",
    "import pandas as pd\n",
    "import pyarrow as pa\n",
    "import numpy as np\n",
    "import pyarrow.feather as feather\n",
    "import base64\n",
    "import json\n",
    "import io\n",
    "import time\n",
    "from buckaroo import BuckarooWidget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "663b81ed-27bc-4278-9275-93a91c0056f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from traitlets import Unicode, List, Dict, observe, Any, Bytes, observe, Bool\n",
    "\n",
    "class Base64Widget(BaseWidget):\n",
    "    _view_name = Unicode('Base64WidgetView').tag(sync=True)\n",
    "    df_base64 = Any(\"\").tag(sync=True)\n",
    "\n",
    "class BytesWidget(BaseWidget):\n",
    "    _view_name = Unicode('BytesWidgetView').tag(sync=True)\n",
    "    df_arrow_bytes = Bytes().tag(sync=True)\n",
    "\n",
    "class SimpleBytesBenchmarkWidget(BaseWidget):\n",
    "    _view_name = Unicode('BytesBenchmarkWidgetView').tag(sync=True)\n",
    "    df_arrow_bytes = Bytes().tag(sync=True)\n",
    "    timing_info = Dict({}).tag(sync=True)\n",
    "    do_calc = Bool(True).tag(sync=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec91023e-4ef0-4090-a8d8-63ecd11e1a02",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def df_to_arrow_bytes(df):\n",
    "    table = pa.Table.from_pandas(df)\n",
    "    fake_file = io.BytesIO()\n",
    "    feather.write_feather(table, fake_file, compression='uncompressed')\n",
    "    fake_file.seek(0)\n",
    "    return fake_file.read()\n",
    "\n",
    "def df_to_base64(df):\n",
    "    return base64.b64encode(df_to_arrow_bytes(df)).decode('utf8')\n",
    "\n",
    "def df_gen_func(exp):\n",
    "    N = 4 * (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",
    "    return df\n",
    "\n",
    "def second_convert(ser):\n",
    "    return ser.dt.seconds + ser.dt.microseconds / (10**6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbb1e102-8e10-48f0-912a-30257211bafa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from functools import wraps\n",
    "def yield_for_change(widget, attribute):\n",
    "    \"\"\"Pause a generator to wait for a widget change event.\n",
    "        \n",
    "    This is a decorator for a generator function which pauses the generator on yield\n",
    "    until the given widget attribute changes. The new value of the attribute is\n",
    "    sent to the generator and is the value of the yield.\n",
    "    \"\"\"\n",
    "    def f(iterator):\n",
    "        @wraps(iterator)\n",
    "        def inner():\n",
    "            i = iterator()\n",
    "            def next_i(change):\n",
    "                try:\n",
    "                    i.send(change.new)\n",
    "                except StopIteration as e:\n",
    "                    widget.unobserve(next_i, attribute)\n",
    "            widget.observe(next_i, attribute)\n",
    "            # start the generator\n",
    "            next(i)\n",
    "        return inner\n",
    "    return f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93e3dc39-ee90-4cbb-8e86-8b682970477d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ipywidgets import IntSlider, VBox, HTML\n",
    "slider2=IntSlider()\n",
    "\n",
    "@yield_for_change(slider2, 'value')\n",
    "def f():\n",
    "    for i in range(10):\n",
    "        print('did work %s'%i)\n",
    "        x = yield\n",
    "        print('generator function continued with value %s'%x)\n",
    "f()\n",
    "\n",
    "slider2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0d62f4c-fe17-49b7-878d-ac86865f8362",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%gui asyncio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4626e26-4b0d-432d-9822-63289b2da2fb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import asyncio\n",
    "def wait_for_change(widget, value):\n",
    "    future = asyncio.Future()\n",
    "    def getvalue(change):\n",
    "        # make the new value available\n",
    "        future.set_result(change.new)\n",
    "        widget.unobserve(getvalue, value)\n",
    "    widget.observe(getvalue, value)\n",
    "    return future"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "878646b0-647f-47dc-9d10-67e80b02ab65",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "w = [SimpleBytesBenchmarkWidget(df_arrow_bytes = df_to_arrow_bytes(df_gen_func(6)))]\n",
    "out = Output()\n",
    "results = []\n",
    "async def f():\n",
    "    x = await wait_for_change(w[0], 'timing_info')\n",
    "    out.append_stdout('async function continued with value ' + str(x) + '\\n')\n",
    "    results.append(x)\n",
    "    w[0].close()\n",
    "    del w[0]\n",
    "asyncio.ensure_future(f())\n",
    "display(out)\n",
    "display(w[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c01ffc61-4e36-45e8-a48d-c9e91d2dce2d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def run_benchmark():\n",
    "    results = []\n",
    "\n",
    "    for i in range(0,9):\n",
    "        w = [SimpleBytesBenchmarkWidget(df_arrow_bytes = df_to_arrow_bytes(df_gen_func(6)))]\n",
    "        out = Output()\n",
    "        async def f():\n",
    "            x = await wait_for_change(w[0], 'timing_info')\n",
    "            out.append_stdout('async function continued with value ' + str(x) + '\\n')\n",
    "            results.append(x)\n",
    "            w[0].close()\n",
    "            del w[0]\n",
    "        asyncio.ensure_future(f())\n",
    "        display(out)\n",
    "        display(w[0])\n",
    "        while i > len(results):\n",
    "            time.sleep(.5)\n",
    "run_benchmark()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7b33c1c-d3f0-4d46-aaca-6a160b7cf76a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def run_benchmark():\n",
    "    results = []\n",
    "\n",
    "    for i in range(9):\n",
    "        w = [SimpleBytesBenchmarkWidget(df_arrow_bytes = df_to_arrow_bytes(df_gen_func(6)))]\n",
    "        out = Output()\n",
    "        async def f():\n",
    "            x = await wait_for_change(w[0], 'timing_info')\n",
    "            out.append_stdout('async function continued with value ' + str(x) + '\\n')\n",
    "            results.append(x)\n",
    "            w[0].close()\n",
    "            del w[0]\n",
    "        asyncio.ensure_future(f())\n",
    "        display(out)\n",
    "        display(w[0])\n",
    "run_benchmark()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11d2caa3-c95e-42dd-83c7-446562f27b5f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ipywidgets import IntSlider, Output\n",
    "slider = IntSlider()\n",
    "out = Output()\n",
    "\n",
    "async def f():\n",
    "    for i in range(10):\n",
    "        out.append_stdout('did work ' + str(i) + '\\n')\n",
    "        x = await wait_for_change(slider, 'value')\n",
    "        out.append_stdout('async function continued with value ' + str(x) + '\\n')\n",
    "asyncio.ensure_future(f())\n",
    "\n",
    "slider"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d61b3006-6d2f-4bd7-8f65-e51df09acead",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b67829be-2d85-4e4e-9cc6-395040a95309",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#sbw = \n",
    "out = Output()\n",
    "w = [SimpleBytesBenchmarkWidget(df_arrow_bytes = df_to_arrow_bytes(df_gen_func(4)))]\n",
    "#results = []\n",
    "async def f():\n",
    "#    for i in range(10):\n",
    "    #out.append_stdout('did work ' + str(i) + '\\n')\n",
    "    x = await wait_for_change(w[0], 'timing_info')\n",
    "    out.append_stdout('async function continued with value ' + str(x) + '\\n')\n",
    "    results.append(x)\n",
    "    sbw.close()\n",
    "    del w[0]\n",
    "asyncio.ensure_future(f())\n",
    "display(out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "773cca8a-3129-4012-bed6-57e3fa4d8ce8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# I can't get this threading method to work consistently\n",
    "import threading\n",
    "from IPython.display import display\n",
    "import ipywidgets as widgets\n",
    "import time\n",
    "res = []\n",
    "def work(w2):\n",
    "    total = 100\n",
    "    for i in range(100):\n",
    "        print(\"testing\")\n",
    "        if w2.timing_info:\n",
    "            print(w2.timing_info)\n",
    "            res.append(w2.timing_info)\n",
    "            w2.close()\n",
    "            del w2\n",
    "            return\n",
    "        time.sleep(0.2)\n",
    "        #print(\"waiting\")\n",
    "w = SimpleBytesBenchmarkWidget(df_arrow_bytes = df_to_arrow_bytes(df_gen_func(6)))\n",
    "thread = threading.Thread(target=work, args=(w,))\n",
    "display(w)\n",
    "thread.start()\n",
    "#thread.join()\n",
    "while len(res) == 0:\n",
    "    time.sleep(0.2)\n",
    "print(\"got a result\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b806962-300e-4b54-837f-28be92551c97",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleBytesBenchmarkWidget(BaseWidget):\n",
    "    _view_name = Unicode('BytesBenchmarkWidgetView').tag(sync=True)\n",
    "    df_arrow_bytes = Bytes().tag(sync=True)\n",
    "    timing_info = Dict({}).tag(sync=True)\n",
    "    do_calc = Bool(True).tag(sync=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0b25468-fba4-4995-b639-81520a066791",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class BaseBenchmarkWidget(BaseWidget):\n",
    "    _view_name = Unicode('BytesBenchmarkWidgetView').tag(sync=True)\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__\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\n",
    "        self.timing_info = {}\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()\n",
    "        self.current_timing['rows'] = len(gen_df)\n",
    "        self.current_timing['elements'] = len(gen_df) * len(gen_df.columns)\n",
    "        byts = self.convert_df(gen_df)\n",
    "        self.current_timing['bytes'] = len(byts)\n",
    "\n",
    "        self.current_timing['serial_t2'] = pd.Timestamp.utcnow()\n",
    "        self.do_calc = True\n",
    "        #print(\"setting arrow bytes\", len(byts) / 1_000_000)\n",
    "        setattr(self, self.target_prop, byts)\n",
    "        self.current_timing['serial_t3'] = pd.Timestamp.utcnow()\n",
    "        #print(\"after set arrow bytes\")\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",
    "            #time.sleep(.1)\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",
    "        return df_to_base64(df)\n",
    "class BytesBenchmarkWidget(BaseBenchmarkWidget):\n",
    "    _view_name = Unicode('BytesBenchmarkWidgetView').tag(sync=True)\n",
    "\n",
    "    def convert_df(self, df):\n",
    "        return df_to_arrow_bytes(df)\n",
    "    df_arrow_bytes = Bytes().tag(sync=True)\n",
    "    target_prop = 'df_arrow_bytes'\n",
    "\n",
    "class DFDataBenchmarkWidget(BaseBenchmarkWidget):\n",
    "    _view_name = Unicode('DFDataBenchmarkWidgetView').tag(sync=True)\n",
    "\n",
    "    def convert_df(self, df):\n",
    "        return json.loads(df.to_json(orient='records'))\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": "ed20f3a6-69b2-4988-8a63-042c00e895ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_obs_dfs = []\n",
    "dfd_bn = DFDataBenchmarkWidget(df_gen_func, min_n=1, max_n=4)\n",
    "dfd_bn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6d32ae2-581b-45db-816e-53fa2b4e683b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "all_obs_dfs.append(dfd_bn.bw.raw_df)\n",
    "dfd_bn.close()\n",
    "dfd_bn.bw.close()\n",
    "del dfd_bn.bw\n",
    "del dfd_bn\n",
    "b64_bn = Base64BenchmarkWidget(df_gen_func, min_n=1, max_n=6)\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",
    "\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.close()\n",
    "del byt_bn"
   ]
  },
  {
   "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['rows_per_second'] = full_df['roundtrip'] / full_df['rows']\n",
    "full_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.8.18"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "e7fbbc3c2be94a3c8de73940b8d678b3": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "fb7f994535094a82b2cae38e21445309": {
      "model_module": "df_cereal",
      "model_module_version": "^0.6.5",
      "model_name": "ReactWidgetModel",
      "state": {
       "_model_module_version": "^0.6.5",
       "_model_name": "ReactWidgetModel",
       "_view_module_version": "^0.6.5",
       "_view_name": "BytesBenchmarkWidgetView",
       "df_arrow_bytes": {},
       "do_calc": true,
       "layout": "IPY_MODEL_e7fbbc3c2be94a3c8de73940b8d678b3",
       "timing_info": {
        "t1": "2024-03-19T12:06:43.704Z",
        "t2": "2024-03-19T12:06:43.708Z",
        "t3": "2024-03-19T12:06:43.783Z"
       }
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
