Add an elementwise operation¶
This notebook will show you how to add a new elementwise operation to an existing backend.
We are going to add julianday
, a function supported by the SQLite database, to the SQLite Ibis backend.
The Julian day of a date, is the number of days since January 1st, 4713 BC. For more information check the Julian day wikipedia page.
Step 1: Define the Operation¶
Let's define the julianday
operation as a function that takes one string input argument and returns a float.
def julianday(date: str) -> float:
"""Julian date"""
import ibis.expr.datatypes as dt
import ibis.expr.rules as rlz
from ibis.expr.operations import ValueOp
class JulianDay(ValueOp):
arg = rlz.string
output_dtype = dt.float32
output_shape = rlz.shape_like('arg')
We just defined a JulianDay
class that takes one argument of type string or binary, and returns a float.
Step 2: Define the API¶
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.
We still need to add a method to StringValue
and BinaryValue
(this needs to work on both scalars and columns).
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.
We can do this by defining a function and assigning it to the appropriate class of expressions.
from ibis.expr.types import BinaryValue, StringValue
def julianday(string_value):
return JulianDay(string_value).to_expr()
StringValue.julianday = julianday
Interlude: Create some expressions with sha1
¶
import ibis
t = ibis.table([('string_col', 'string')], name='t')
t.string_col.julianday()
r0 := UnboundTable: t string_col string Selection[r0] selections: JulianDay(string_col): JulianDay(r0.string_col)
Step 3: Turn the Expression into SQL¶
import sqlalchemy as sa
@ibis.sqlite.add_operation(JulianDay)
def _julianday(translator, expr):
# pull out the arguments to the expression
(arg,) = expr.args
# compile the argument
compiled_arg = translator.translate(arg)
# return a SQLAlchemy expression that calls into the SQLite julianday function
return sa.func.julianday(compiled_arg)
Step 4: Putting it all Together¶
!curl -LsS -o geography.db 'https://storage.googleapis.com/ibis-tutorial-data/geography.db'
import os
import tempfile
import ibis
db_fname = 'geography.db'
con = ibis.sqlite.connect(db_fname)
Create and execute a julianday
expression¶
independence = con.table('independence')
independence
DatabaseTable: independence country_code string independence_date date independence_from string
day = independence.independence_date.cast('string')
day
r0 := DatabaseTable: independence country_code string independence_date date independence_from string Selection[r0] selections: Cast(independence_date, string): Cast(r0.independence_date, to=string)
julianday_expr = day.julianday().name("jday")
julianday_expr
r0 := DatabaseTable: independence country_code string independence_date date independence_from string Selection[r0] selections: jday: JulianDay(Cast(r0.independence_date, to=string))
sql_expr = julianday_expr.compile()
print(sql_expr)
SELECT julianday(CAST(t0.independence_date AS TEXT)) AS jday FROM independence AS t0
result = julianday_expr.to_pandas()
result.head()
0 2422189.5 1 2419734.5 2 2437850.5 3 2442727.5 4 2444909.5 Name: jday, dtype: float32
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
scalar = ibis.literal('2010-03-14')
scalar
'2010-03-14'
julianday_scalar = scalar.julianday()
con.execute(julianday_scalar)
2455269.5