SQL translators

The purpose of this vignette is to walk through how expressions like _.id.mean() are converted into SQL.

This process involves 3 parts

  1. SQL translation functions, e.g. taking column “id” and producing the SQL “ROUND(id)”.

  2. SQL translation from a symbolic call

  • Converting method calls like _.id.round(2) to round(_.id, 2)

  • Looking up SQL translators (e.g. for “mean” function call)

  1. Handling SQL partitions, like in OVER clauses

Using sqlalchemy select statment for convenience

Throughout this vignette, we’ll use a select statement object from sqlalchemy, so we can conveniently access its columns as needed.

[1]:
from sqlalchemy import sql
col_names = ['id', 'x', 'y']
sel = sql.select([sql.column(x) for x in col_names])

print(sel)
print(type(sel.columns))
print(sel.columns)
SELECT id, x, y
<class 'sqlalchemy.sql.base.ImmutableColumnCollection'>
['id', 'x', 'y']

Translator functions

A SQL translator function takes…

  • a first argument that is a sqlalchemy Column

  • (optional) additional arguments for the translation

A simple translator

[2]:
f_simple_round = lambda col, n: sql.func.round(col, n)

sql_expr = f_simple_round(sel.columns.x, 2)

print(sql_expr)
round(x, :round_1)

The function above is essentially what most translator functions are.

For example, here is the round function defined for postgresql. One key difference is that it casts the column to a numeric beforehand.

[3]:
from siuba.sql.dialects.postgresql import funcs

f_round = funcs['scalar']['round']
sql_expr = f_round(sel.columns.x, 2)

print(sql_expr)
round(CAST(x AS NUMERIC), :round_1)

Handling windows with custom Over clauses

[4]:
f_win_mean = funcs['window']['mean']

sql_over_expr = f_win_mean(sel.columns.x)

print(type(sql_over_expr))
print(sql_over_expr)
<class 'siuba.sql.translate.AggOver'>
avg(x) OVER ()

Notice that this window expression has an empty over clause. This clause needs to be able to include any variables we’ve grouped the data by.

Siuba handles this by implementing a set_over method on these custom sqlalchemy Over clauses, which takes grouping and ordering variables as arguments.

[5]:
group_by_clause = sql.elements.ClauseList(sel.columns.x, sel.columns.y)
print(sql_over_expr.set_over(group_by_clause))
avg(x) OVER (PARTITION BY x, y)

Call shaping

The section above discusses how SQL translators are functions that take a sqlalchemy column, and return a SQL expression. However, when using siuba we often have expressions like…

mutate(data, x = _.y.round(2))

In this case, before we can even use a SQL translator, we need to…

  • find the name and arguments of the method being called

  • find the column it is being called on

This is done by using the CallTreeLocal class to analyze the tree of operations for each expression.

[6]:
from siuba.siu import Lazy, CallTreeLocal, Call, strip_symbolic
from siuba import _

_.y.round(2)
[6]:
█─'__call__'
├─█─.
│ ├─█─.
│ │ ├─_
│ │ └─'y'
│ └─'round'
└─2

Example of translation with CallTreeLocal

[7]:
from siuba.sql.dialects.postgresql import funcs

local_funcs = {**funcs['scalar'], **funcs['window']}

call_shaper = CallTreeLocal(
    local_funcs,
    call_sub_attr = ('dt',)
    )
[8]:
symbol = _.id.mean()
call = strip_symbolic(symbol)
print(call)
_.id.mean()
[9]:
func_call = call_shaper.enter(call)
print(func_call(sel.columns))
avg(id) OVER ()

This is the same result as when we called the SQL translator for mean manually! In that section we also showed that we can set group information, so that it takes an average within each group.

In this case it’s easy to set group information to the Over clause. However, an additional challenge is when it’s part of a larger expression…

[10]:
call2 = strip_symbolic(_.id.mean() + 1)
func_call2 = call_shaper.enter(call2)

func_call2(sel.columns)
[10]:
<sqlalchemy.sql.elements.BinaryExpression object at 0x7fcb3d0bf550>

Handling partitions

While the first section showed how siuba’s custom Over clauses can add grouping info to a translation, it is missing one key detail: expressions that generate Over clauses, like _.id.mean(), can be part of larger expressions. For example _.id.mean() + 1.

In this case, if we look at the call tree for that expression, the top operation is the addition…

[11]:
_.id.mean() + 1
[11]:
█─+
├─█─'__call__'
│ └─█─.
│   ├─█─.
│   │ ├─_
│   │ └─'id'
│   └─'mean'
└─1

How can we create the appropriate expression…

avg(some_col) OVER (PARTITION BY x, y) + 1

when the piece that needs grouping info is not easily accessible? The answer is by using a tree visitor, which steps down every black rectangle in the call tree shown above, from top to bottom.

Full example

Below, we copy the code from the call shaping section..

[12]:
from siuba.sql.verbs import track_call_windows
from siuba import _
from siuba.sql.dialects.postgresql import funcs

local_funcs = {**funcs['scalar'], **funcs['window']}

call_shaper = CallTreeLocal(
    local_funcs,
    call_sub_attr = ('dt',)
    )

symbol3 = _.id.mean() + 1
call3 = strip_symbolic(symbol3)
func_call3 = call_shaper.enter(call3)

Finally, we pass the shaped call…

[13]:
col, windows = track_call_windows(
    func_call3,
    sel.columns,
    group_by = ['x', 'y'],
    order_by = []
    )

print(col)
print(windows)

avg(id) OVER (PARTITION BY x, y) + :param_1
[]

Edit page on github here. Interactive version: Binder badge