03.lazyarray-udf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# User Defined Functions\n",
    "\n",
    "In this section, we will see how to do computations with NDArray and/or NumPy arrays using functions defined by ourselves (aka User-Defined-Functions).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-11-27T12:40:17.380740Z",
     "start_time": "2024-11-27T12:40:17.378001Z"
    },
    "is_executing": true
   },
   "outputs": [],
   "source": [
    "import numba as nb\n",
    "import numpy as np\n",
    "\n",
    "import blosc2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A simple example\n",
    "First, let's create a NDArray, a NumPy array and regular scalar, which we will use to create a LazyArray."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-11-27T12:40:17.507724Z",
     "start_time": "2024-11-27T12:40:17.455663Z"
    }
   },
   "outputs": [],
   "source": [
    "shape = (5_000, 2_000)\n",
    "a = np.linspace(0, 1, np.prod(shape), dtype=np.int32).reshape(shape)\n",
    "b = blosc2.arange(np.prod(shape), dtype=np.float32, shape=shape)\n",
    "s = 2.1  # a regular scalar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's define our function. This function can be executed for each chunk and will always receive 3 parameters. The first one is the inputs tuple to which we can pass any operand such as a NDArray, NumPy array or Python scalar. The second is the output buffer to be filled and the third is an offset corresponding to the starting point inside the array of the chunk being filled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-11-27T12:40:17.539888Z",
     "start_time": "2024-11-27T12:40:17.536601Z"
    }
   },
   "outputs": [],
   "source": [
    "def myudf(inputs_tuple, output, offset):\n",
    "    x, y, s = inputs_tuple   # at this point, all are either numpy arrays or scalars\n",
    "    output[:] = x**3 + np.sin(y) + s + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, this function will take the first input, add one and save the result in output.\n",
    "\n",
    "Now, to actually create a `LazyUDF` object (which also follows the [LazyArray interface](https://www.blosc.org/python-blosc2/reference/lazyarray.html)) we will use its constructor `lazyudf`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-11-27T12:40:17.576645Z",
     "start_time": "2024-11-27T12:40:17.561841Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Type: <class 'blosc2.lazyexpr.LazyUDF'>\n"
     ]
    }
   ],
   "source": [
    "larray = blosc2.lazyudf(myudf, (a, b, s), a.dtype)\n",
    "print(f\"Type: {type(larray)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, to execute and get the result of your function you can choose between the `__getitem__` and `compute` methods.\n",
    "The main difference is that the first one will return the computed result as a NumPy array whereas the second one will return a NDArray. Let's see `__getitem__` first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-11-27T12:40:17.730251Z",
     "start_time": "2024-11-27T12:40:17.655380Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Type: <class 'numpy.ndarray'>\n",
      "CPU times: user 213 ms, sys: 51.9 ms, total: 265 ms\n",
      "Wall time: 231 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "npc = larray[:]\n",
    "print(f\"Type: {type(npc)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Now, let's use `compute` for the same purpose. The advantage of using this method is that you can pass some construction parameters for the resulting NDArray like the `urlpath` to store the resulting array on-disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-11-27T12:40:17.838439Z",
     "start_time": "2024-11-27T12:40:17.764610Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Type: <class 'blosc2.ndarray.NDArray'>\n",
      "type    : NDArray\n",
      "shape   : (5000, 2000)\n",
      "chunks  : (500, 2000)\n",
      "blocks  : (20, 2000)\n",
      "dtype   : int32\n",
      "cratio  : 258.64\n",
      "cparams : CParams(codec=<Codec.ZSTD: 5>, codec_meta=0, clevel=1, use_dict=False, typesize=4,\n",
      "        : nthreads=11, blocksize=160000, splitmode=<SplitMode.AUTO_SPLIT: 3>,\n",
      "        : filters=[<Filter.NOFILTER: 0>, <Filter.NOFILTER: 0>, <Filter.NOFILTER: 0>,\n",
      "        : <Filter.NOFILTER: 0>, <Filter.NOFILTER: 0>, <Filter.SHUFFLE: 1>], filters_meta=[0, 0,\n",
      "        : 0, 0, 0, 0], tuner=<Tuner.STUNE: 0>)\n",
      "dparams : DParams(nthreads=11)\n",
      "\n",
      "CPU times: user 284 ms, sys: 71.5 ms, total: 356 ms\n",
      "Wall time: 259 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "c = larray.compute(urlpath=\"array-compute.b2nd\", mode=\"w\")\n",
    "print(f\"Type: {type(c)}\")\n",
    "print(c.info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Numba\n",
    "Let's see how Python-Blosc2 can use Numba as an UDF. For this, let's decorate the same function with Numba."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-11-27T12:40:17.924259Z",
     "start_time": "2024-11-27T12:40:17.920094Z"
    }
   },
   "outputs": [],
   "source": [
    "@nb.jit(nopython=True, parallel=True)\n",
    "def myudf_numba(inputs_tuple, output, offset):\n",
    "    x, y, s = inputs_tuple\n",
    "    output[:] = x**3 + np.sin(y) + s + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-11-27T12:40:18.571717Z",
     "start_time": "2024-11-27T12:40:18.002076Z"
    }
   },
   "outputs": [],
   "source": [
    "larray2 = blosc2.lazyudf(myudf_numba, (a, b, s), a.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cool! We made our first Numba UDF function.  Now, let's evaluate it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.12 s, sys: 180 ms, total: 1.3 s\n",
      "Wall time: 1.26 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "npc2 = larray2[:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Incidentally, the pure Python version was faster than Numba.  This is because Numba has\n",
    "large initialization overheads and the function is quite simple.  For more complex functions, or larger arrays, the difference will be less noticeable or favorable to it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise\n",
    "\n",
    "Check which array size Numba UDF starts to be competitive.  If you master Numba enough, you may also want to unroll loops in UDF and see whether you can make it faster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "We have seen how to build new LazyArray objects coming from other NDArray or NumPy objects and use User Defined Functions (UDFs) to create the desired result.  We have also demonstrated that integrating Numba in UDF is pretty easy."
   ]
  }
 ],
 "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}