Custom Functions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Apart from the SQL functions that are already implemented in `dask-sql`, it is possible to add custom functions and aggregations.\n",
    "Have a look into [the documentation](https://dask-sql.readthedocs.io/en/latest/pages/custom.html) for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import dask.dataframe as dd\n",
    "import dask.datasets\n",
    "from dask_sql.context import Context"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use some generated test data for the notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = Context()\n",
    "\n",
    "df = dask.datasets.timeseries().reset_index().persist()\n",
    "c.create_table(\"timeseries\", df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a first step, we will create a scalar function to calculate the absolute value of a column.\n",
    "(Please note that this can also be done via the `ABS` function in SQL):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The input to the function will be a dask series\n",
    "def my_abs(x):\n",
    "    return x.abs()\n",
    "\n",
    "# As SQL is a typed language, we need to specify all types \n",
    "c.register_function(my_abs, \"MY_ABS\", parameters=[(\"x\", np.float64)], return_type=np.float64)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now able to use our new function in all queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c.sql(\"\"\"\n",
    "    SELECT\n",
    "        x, y, MY_ABS(x) AS \"abs_x\", MY_ABS(y) AS \"abs_y\"\n",
    "    FROM\n",
    "        \"timeseries\"\n",
    "    WHERE\n",
    "        MY_ABS(x * y) > 0.5\n",
    "\"\"\").compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will register an aggregation, which gets a column as input and returns a single value.\n",
    "An aggregation needs to be an instance of `dask.Aggregation` (see the [dask docu](https://docs.dask.org/en/latest/dataframe-groupby.html#aggregate))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_sum = dd.Aggregation(\"MY_SUM\", lambda x: x.sum(), lambda x: x.sum())\n",
    "\n",
    "c.register_aggregation(my_sum, \"MY_SUM\", [(\"x\", np.float64)], np.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c.sql(\"\"\"\n",
    "    SELECT\n",
    "        name, MY_SUM(x) AS \"my_sum\"\n",
    "    FROM\n",
    "        \"timeseries\"\n",
    "    GROUP BY\n",
    "        name\n",
    "    LIMIT 10\n",
    "\"\"\").compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}