{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extending Ibis Part 1: Adding a New Elementwise Expression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook will show you how to add a new elementwise operation to an existing backend.\n",
    "\n",
    "We are going to add `julianday`, a function supported by the SQLite database, to the SQLite Ibis backend.\n",
    "\n",
    "The Julian day of a date, is the number of days since January 1st, 4713 BC. For more information check the [Julian day](https://en.wikipedia.org/wiki/Julian_day) wikipedia page."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Define the Operation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define the `julianday` operation as a function that takes one string input argument and returns a float.\n",
    "\n",
    "```python\n",
    "def julianday(date: str) -> float:\n",
    "    \"\"\"Julian date\"\"\"\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ibis.expr.datatypes as dt\n",
    "import ibis.expr.rules as rlz\n",
    "\n",
    "from ibis.expr.operations import ValueOp, Arg\n",
    "\n",
    "\n",
    "class JulianDay(ValueOp):\n",
    "    arg = Arg(rlz.string)\n",
    "    output_type = rlz.shape_like('arg', 'float')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We just defined a `JulianDay` class that takes one argument of type string or binary, and returns a float."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Define the API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we know the output type of the operation, to make an expression out of ``JulianDay`` we simply need to construct it and call its `ibis.expr.types.Node.to_expr` method.\n",
    "\n",
    "We still need to add a method to `StringValue` and `BinaryValue` (this needs to work on both scalars and columns).\n",
    "\n",
    "When you add a method to any of the expression classes whose name matches `*Value` both the scalar and column child classes will pick it up, making it easy to define operations for both scalars and columns in one place.\n",
    "\n",
    "We can do this by defining a function and assigning it to the appropriate class\n",
    "of expressions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ibis.expr.types import StringValue, BinaryValue\n",
    "\n",
    "\n",
    "def julianday(string_value):\n",
    "    return JulianDay(string_value).to_expr()\n",
    "\n",
    "\n",
    "StringValue.julianday = julianday"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interlude: Create some expressions with `sha1`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ibis\n",
    "\n",
    "t = ibis.table([('string_col', 'string')], name='t')\n",
    "\n",
    "t.string_col.julianday()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Turn the Expression into SQL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlalchemy as sa\n",
    "\n",
    "\n",
    "@ibis.sqlite.add_operation(JulianDay)\n",
    "def _julianday(translator, expr):\n",
    "    # pull out the arguments to the expression\n",
    "    arg, = expr.op().args\n",
    "    \n",
    "    # compile the argument\n",
    "    compiled_arg = translator.translate(arg)\n",
    "    \n",
    "    # return a SQLAlchemy expression that calls into the SQLite julianday function\n",
    "    return sa.func.julianday(compiled_arg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Putting it all Together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "import ibis\n",
    "\n",
    "db_fname = str(pathlib.Path().resolve().parent.parent / 'tutorial' / 'data' / 'geography.db')\n",
    "\n",
    "con = ibis.sqlite.connect(db_fname)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and execute a `julianday` expression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "independence = con.table('independence')\n",
    "independence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "day = independence.independence_date.cast('string')\n",
    "day"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "julianday_expr = day.julianday()\n",
    "julianday_expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_expr = julianday_expr.compile()\n",
    "print(sql_expr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = julianday_expr.execute()\n",
    "result.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we've defined our operation on `StringValue`, and not just on `StringColumn` we get operations on both string scalars *and* string columns for free"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scalar = ibis.literal('2010-03-14')\n",
    "scalar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "julianday_scalar = scalar.julianday()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "con.execute(julianday_scalar)"
   ]
  }
 ],
 "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
