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 statement for convenience

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

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

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.

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

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.

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.

from siuba.siu import Lazy, CallTreeLocal, Call, strip_symbolic
from siuba import _

_.y.round(2)
█─'__call__'
├─█─.
│ ├─█─.
│ │ ├─_
│ │ └─'y'
│ └─'round'
└─2

Example of translation with CallTreeLocal

from siuba.sql.dialects.postgresql import funcs

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

call_shaper = CallTreeLocal(
    local_funcs,
    call_sub_attr = ('dt',)
    )
symbol = _.id.mean()
call = strip_symbolic(symbol)
print(call)
_.id.mean()
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…

call2 = strip_symbolic(_.id.mean() + 1)
func_call2 = call_shaper.enter(call2)

func_call2(sel.columns)
<sqlalchemy.sql.elements.BinaryExpression object at 0x11889ada0>

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…

_.id.mean() + 1
█─+
├─█─'__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..

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…

col, windows, window_cte = track_call_windows(
    func_call3,
    sel.columns,
    group_by = ['x', 'y'],
    order_by = [],
    # note that this is optional, and results in window_cte being a
    # copy of this select that contains the window clauses
    window_cte = sel.select()
    )

print(col)
print(windows)
print(window_cte)
avg(id) OVER (PARTITION BY x, y) + :param_1
[<sqlalchemy.sql.elements.Label object at 0x11889ec50>, <sqlalchemy.sql.elements.Label object at 0x11889e1d0>]
SELECT id, x, y, avg(id) OVER (PARTITION BY x, y) AS win1, avg(id) OVER (PARTITION BY x, y) + :param_1 AS win2 
FROM (SELECT id, x, y)