Put some Numba in your SQLite

Fair Warning

This library does unsafe things like pass around function pointer addresses as integers. Use at your own risk.

If you're unfamiliar with why passing function pointers' addresses around as integers might be unsafe, then you shouldn't use this library.

Requirements

• Python >=3.10
• numba
• llvmlite

Use nix develop from the repository root to avoid dependency hell.

Installation

• pip install numbsql

Examples

Scalar Functions

These are almost the same as decorating a Python function with numba.jit.

fromtypingimportOptionalfromnumbsqlimportsqlite_udf@sqlite_udfdefadd_one(x: Optional[int]) ->Optional[int]:
"""Add one to `x` if `x` is not NULL."""ifxisnotNone:
returnx+1returnNone

Calling your scalar function

Similar to scalar functions, we register the function with a sqlite3.Connection object:

>>>importsqlite3>>>fromnumbsqlimportcreate_function>>>con=sqlite3.connect(":memory:")
>>>create_function(con, "add_one", 1, add_one)
>>>con.execute("SELECT add_one(1)").fetchall()
[(2,)]

Aggregate Functions

These follow the API of the Python standard library's sqlite3.Connection.create_aggregate method. The difference with numbsql aggregates is that they require two decorators: numba.experimental.jit_class and numbsql.sqlite_udaf. Let's define the avg (arithmetic mean) function for 64-bit floating point numbers.

fromtypingimportOptionalfromnumba.experimentalimportjitclassfromnumbsqlimportsqlite_udaf@sqlite_udaf@jitclassclassAvg:
total: floatcount: intdef__init__(self):
self.total=0.0self.count=0defstep(self, value: Optional[float]) ->None:
ifvalueisnotNone:
self.total+=valueself.count+=1deffinalize(self) ->Optional[float]:
ifnotself.count:
returnNonereturnself.total/self.count

Window Functions

You can also define window functions for use with SQLite's OVER construct:

fromtypingimportOptionalfromnumba.experimentalimportjitclassfromnumbsqlimportsqlite_udaf@sqlite_udaf@jitclassclassWinAvg: # pragma: no covertotal: floatcount: intdef__init__(self) ->None:
self.total=0.0self.count=0defstep(self, value: Optional[float]) ->None:
ifvalueisnotNone:
self.total+=valueself.count+=1deffinalize(self) ->Optional[float]:
count=self.countifcount:
returnself.total/countreturnNonedefvalue(self) ->Optional[float]:
returnself.finalize()
definverse(self, value: Optional[float]) ->None:
ifvalueisnotNone:
self.total-=valueself.count-=1

Calling your aggregate function

Similar to scalar functions, we register the function with a sqlite3.Connection object:

>>>importsqlite3>>>fromnumbsqlimportcreate_aggregate>>>con=sqlite3.connect(":memory:")
>>>create_aggregate(con, "winavg", 1, WinAvg)
>>>con.execute("CREATE TABLE t (x INTEGER, y TEXT)")
>>>con.execute("INSERT INTO t VALUES (1, 'a'), (2, 'a'), (3, 'b')")
>>>con.execute("SELECT winavg(x) FROM t").fetchall()
[(2.0,)]
>>>con.execute("SELECT winavg(x) OVER (PARTITION BY y) FROM t").fetchall()
[(1.5,), (3.0,)]

Goodies

Some string operations are available:

fromtypingimportOptionalfromnumbsqlimportsqlite_udf@sqlite_udfdefnumbsql_len(s: Optional[str]) ->Optional[int]:
returnlen(s) ifsisnotNoneelseNone

>>>importsqlite3>>>fromnumbsqlimportcreate_function>>>con=sqlite3.connect(":memory:")
>>>create_function(con, "numbsql_len", 1, numbsql_len)
>>>con.execute("CREATE TABLE t (name TEXT)")
>>>con.execute("INSERT INTO t VALUES ('Alice', 'Bob', 'Susan', 'Joe')")
>>>con.execute("SELECT numbsql_len(x) FROM t").fetchall()
[(5,), (3,), (5,), (3,)]