Run Numba functions in SQLite: WTF?

https://slides.com/phillipcloud/numsql

Important legal information

The information presented here is offered for informational purposes only and should not be used for any other purpose (including, without limitation, the making of investment decisions). Examples provided herein are for illustrative purposes only and are not necessarily based on actual data. Nothing herein constitutes: an offer to sell or the solicitation of any offer to buy any security or other interest; tax advice; or investment advice. This presentation shall remain the property of Two Sigma Investments, LP (“Two Sigma”) and Two Sigma reserves the right to require the return of this presentation at any time.

Some of the images, logos or other material used herein may be protected by copyright and/or trademark. If so, such copyrights and/or trademarks are most likely owned by the entity that created the material and are used purely for identification and comment as fair use under international copyright and/or trademark laws.  Use of such image, copyright or trademark does not imply any association with such organization (or endorsement of such organization) by Two Sigma, nor vice versa.

Copyright © 2018 TWO SIGMA INVESTMENTS, LP. All rights reserved

 

About me

WARNING:
Hacks AHEAD

TAKEAWAY:

NUMBA iS POWEFUL

TAKEAWAY:

not Hard to extend

Extending SQlite

  • UDF: User-Defined Function
  • UDAF: User-Defined Aggregate Function

USING PYTHON

UDFS

In [6]: import sqlite3

In [7]: con = sqlite3.connect(':memory:')

In [8]: def incr(x):
   ...:     return x + 1
   ...: 

In [9]: con.create_function('incr', 1, incr)

In [10]: list(con.execute('select incr(1)'))
Out[10]: [(2,)]

UDAFS

In [12]: class Average:
    ...:     def __init__(self):
    ...:         self.total = 0.0
    ...:         self.count = 0
    ...:
    ...:     def step(self, value):
    ...:         if value is not None:
    ...:             self.total += value
    ...:             self.count += 1
    ...:
    ...:     def finalize(self):
    ...:         if self.count:
    ...:             return self.total / self.count
    ...:         return None

UDAF FLOW

Scratch space, step and finalize

  • Initialize scratch space to hold intermediate aggregate values
  • For each row
    • step(scratch_space, row['col'])
  • After all rows have been traversed
    • finalize(scratch_space)
k v
a 1
b 2
c NULL
d 4
e 5
total = 0.0
count = 0

UDAF FLOW

Scratch space, step and finalize

SELECT AVG(v)
FROM t
total = 1.0
count = 1
total = 3.0
count = 2
total = 3.0
count = 2
total = 7.0
count = 3
total = 12.0
count = 4
total / count

step(1)

finalize

step(2)

step(3)

step(4)

step(5)

UDAFS

In [13]: con.create_aggregate('my_average', 1, Average)

In [14]: list(con.execute("""
    ...:     with t as (
    ...:         select 1 as c union
    ...:         select 2 as c union
    ...:         select 3 as c
    ...:     )
    ...:     select my_average(c) from t
    ...: """))
Out[14]: [(2.0,)]

USING C

SQLITE C API - STEP

(Called once for each row)

void step(sqlite3_context* ctx, int argc, sqlite3_value** argv) {
  static const size_t nbytes = sizeof(Average);
  void* raw_ctx = sqlite3_aggregate_context(ctx, nbytes);
  Average* agg_ctx = (Average*) raw_ctx;
  if (agg_ctx != NULL) {
    sqlite3_value* argument = argv[0];
    int argument_type = sqlite3_value_type(argument);

    if (argument_type != SQLITE_NULL) {
      double value = sqlite3_value_double(argument);
      agg_ctx->total += value;
      agg_ctx->count += 1;
    }
  }
}
typedef struct {
  double total;
  int64_t count;
} Average;

SQLITE C API - FINALIZE

(Called exactly once, after all rows have been traversed)

void finalize(sqlite3_context* ctx) {
    static const size_t nbytes = sizeof(Average);
    void* raw_ctx = sqlite3_aggregate_context(ctx, nbytes);
    Average* agg_ctx = (Average*) raw_ctx;
    if (agg_ctx != NULL) {
	if (agg_ctx->count > 0) {
	    double average = agg_ctx->total / agg_ctx->count;
	    sqlite3_result_double(ctx, average);
	} else {
	    sqlite3_result_null(ctx);
	}
    }
}

GOAL:
PYTHON API + C SPEED

(Preserve the stdlib API)

C refresher

  • Function pointers are pointers
  • Pointers are integers
  • => Function pointers are integers
  • Not all integers are function pointers

But some are

😈

EVIL C refresher

never call UNTRUSTED functioN pointers

SERiOUSLY, DON'T

Numba

NUmba

  • JIT compiler for numeric Python
  • LLVM
from numba import njit

@njit
def fast_sum(array):
    total = 0.0
    for el in array:
        total += el
    return total
  • jitclasses
  • Extensibility
  • Helpful devs
  • cfunc

Jit classes

In [12]: class Average:
    ...:     def __init__(self):
    ...:         self.value = 0.0
    ...:         self.count = 0
    ...:
    ...:     def step(self, value):
    ...:         if value is not None:
    ...:             self.value += value
    ...:             self.count += 1
    ...:
    ...:     def finalize(self):
    ...:         if self.count:
    ...:             return self.value / self.count
    ...:         return None
In [12]: @jitclass([
    ...:     ('value', 'float64'),
    ...:     ('count', 'int64'),
    ...: ])
    ...: class Average:
    ...:     def __init__(self):
    ...:         self.value = 0.0
    ...:         self.count = 0
    ...:
    ...:     def step(self, value):
    ...:         if value is not None:
    ...:             self.value += value
    ...:             self.count += 1
    ...:
    ...:     def finalize(self):
    ...:         if self.count:
    ...:             return self.value / self.count
    ...:         return None
typedef struct {
    double value;
    int64_t count;
} Average;

void step(Average* self, double value);
double finalize(Average* self);

What they look like in C

more specific GOAL:

call A JITted METHOD From A C FUNCTION defined in PYTHON

... that matches the SQLite C API

@CFUNC

Kind of evil

😈

In [33]: from numba import cfunc, float64

In [34]: @cfunc(float64(float64))
    ...: def my_incr(x):
    ...:     return x + 1
    ...:
In [33]: from numba import cfunc, float64

In [34]: @cfunc(float64(float64))
    ...: def my_incr(x):
    ...:     return x + 1
    ...:

In [35]: my_incr.address
Out[35]: 139623106748448

wat

SLUMBA

  • Numba provides LLVM codegen
  • User provides step/finalize methods
  • What the heck does slumba do?

slumba
NUMSQL

Important REnaming information

Tying it ALL together

1. USER DEFINES A JITClass

from numba import int64, float64, jitclass




@jitclass(dict(total=float64, count=int64))
class Avg:
    def __init__(self):
        self.total = 0.0
        self.count = 0

    def step(self, value):
        if value is not None:
	    self.total += value
	    self.count += 1

    def finalize(self):
	if not self.count:
	    return None
        return self.total / self.count
from numba import int64, float64, jitclass
from slumba import sqlite_udaf


@sqlite_udaf(float64(float64))
@jitclass(dict(total=float64, count=int64))
class Avg:
    def __init__(self):
        self.total = 0.0
        self.count = 0

    def step(self, value):
        if value is not None:
	    self.total += value
	    self.count += 1

    def finalize(self):
	if not self.count:
	    return None
        return self.total / self.count

1. USER DEFINES A JITClass

2. SLUMBA DEFINES WRAPPERS

@cfunc(void(voidptr, intc, CPointer(voidptr)))
def step(ctx, argc, argv):
    raw_pointer = sqlite3_aggregate_context(ctx, sizeof(cls))
    agg_ctx = unsafe_cast(raw_pointer, cls)
    if not_null(agg_ctx):
        agg_ctx.step(*make_arg_tuple(step_func, argv))
@cfunc(void(voidptr))
def finalize(ctx):
    raw_pointer = sqlite3_aggregate_context(ctx, sizeof(cls))
    agg_ctx = unsafe_cast(raw_pointer, cls)
    if not_null(agg_ctx):
        result = agg_ctx.finalize()
        if result is None:
            sqlite3_result_null(ctx)
        else:
            result_setter = get_result_setter(result)
            result_setter(ctx, result)
@cfunc(void(voidptr, intc, CPointer(voidptr)))
def step(ctx, argc, argv):
    raw_pointer = sqlite3_aggregate_context(ctx, sizeof(cls))
    agg_ctx = unsafe_cast(raw_pointer, cls)
    if not_null(agg_ctx):
        agg_ctx.step(*make_arg_tuple(step_func, argv))

STEp

@cfunc(void(voidptr, intc, CPointer(voidptr)))
def step(ctx, argc, argv):
    raw_pointer = sqlite3_aggregate_context(ctx, sizeof(cls))
    agg_ctx = unsafe_cast(raw_pointer, cls)
    if not_null(agg_ctx):
        agg_ctx.step(*make_arg_tuple(step_func, argv))
  1. Allocate space for the class
  2. Cast pointer to our type
  3. If allocation didn't fail
    1. construct an argument tuple given a function signature and arguments
  4. Call the step method

TYPE INFERENCE IS HARD ...

# double my_func(sqlite3_context* ctx,
#                int64_t narg,
#                sqlite3_value**);

@cfunc(float64(voidptr, int64, CPointer(voidptr)))
def step(ctx, argc, argv):
    args = ()  # tuple of *shrug*

    for i in range(argc):
        # sqlite3_value* argument
        arg = argv[i]

        # note: *runtime* type
        value_type = sqlite3_value_type(arg)

        # check the *runtime* type
        if value_type == SQLITE_FLOAT:
            args += (sqlite3_value_double(arg),)

        elif value_type == SQLITE_INTEGER:
            args += (sqlite3_value_int64(arg),)

    # this is our jitclass method, and we must
    # know the number of arguments and their
    # types at *compile* time
    return agg_ctx.my_jitclass_method(*args)

... BUT we require

argument types

WhaT are we solving??

TypE inference is hard not necessary

def codegen(context, builder, signature, args):
    _, argv = args
    converted_args = []
    for i, argtype in enumerate(argtypes):


        instr =


        converted_args.append(instr)

    res = context.make_tuple(
        builder, tuple_type, converted_args)
    return imputils.impl_ret_borrowed(
        context, builder, tuple_type, res)
make_arg_tuple

Face melting calls into llvmlite

make_arg_tuple

# get the appropriate ctypes extraction routine
ctypes_function = VALUE_EXTRACTORS[argtype]

# create a numba function type for the converter
converter = ctypes_utils.make_function_type(ctypes_function)

# get the function pointer instruction out
fn = context.get_constant_generic(
    builder, converter, ctypes_function)

# get a pointer to the ith argument
element_pointer = cgutils.gep(builder, argv, i)

# deref that pointer
element = builder.load(element_pointer)

# call the value extraction routine
instr = builder.call(fn, [element])

3. NUMBA JITS all the things

Numba

Python

LLVM

C

4. USER REGISTERS THE AGG WITH THE DB

>>> con = sqlite3.connect(...)
>>> create_aggregate(con, 'myavg', 1, Avg)

What does SLUMBA PROVIDE?

Slumba adds the following Numba primitives

  • sizeof (like C sizeof)
  • unsafe_cast (cast a void* to whatever)
  • pointer not null (check if a pointer is not NULL)
  • make_arg_tuple

Some disappointing benchmarks

loW CARD. GROUP BY

WHole-TABLE AGG

HIGH CARD. GROUP BY

no reason to talk about scalar functions :(

What Slumba doESN'T DO

HEAP STUFF

# if your data look anything like this:

           "I am a string"

# you'll need to extend slumba, and 
# probably numba too

... lots of things look like that

Pray to THe DEMO Gods

Thank you!

numsql

By Phillip Cloud

numsql

Melt your mind by learning how to run jit-compiled numba functions inside of SQLite. Not for the faint of heart.

  • 39

More from Phillip Cloud