在SQLite中加入一些Numba

重要提示

此库执行不安全操作，例如将函数指针地址作为整数传递。 风险自负。

如果您不熟悉为什么将函数指针的地址作为整数传递可能不安全，则不应使用此库。

需求

• Python >=3.7
• numba

使用存储库中的nix-shell以避免依赖问题。

安装

• poetry install

示例

标量函数

这些几乎与用numba.jit装饰Python函数相同。

from typing import Optional

from numbsql import sqlite_udf


@sqlite_udf
def add_one(x: Optional[int]) -> Optional[int]:
    """Add one to `x` if `x` is not NULL."""

    if x is not None:
        return x + 1
    return None

聚合函数

这些遵循Python标准库中sqlite3.Connection.create_aggregate方法的API。与numbsql聚合函数的不同之处在于，它们需要两个装饰器：numba.experimental.jit_class和numbsql.sqlite_udaf。让我们定义64位浮点数的avg（算术平均值）函数。

from typing import Optional

from numba.experimental import jitclass

from numbsql import sqlite_udaf


@sqlite_udaf
@jitclass
class Avg:
    total: float
    count: int

    def __init__(self):
        self.total = 0.0
        self.count = 0

    def step(self, value: Optional[float]) -> None:
        if value is not None:
            self.total += value
            self.count += 1

    def finalize(self) -> Optional[float]:
        if not self.count:
            return None
        return self.total / self.count

窗口函数

您还可以定义用于SQLite的OVER构造的窗口函数。

from typing import Optional

from numba.experimental import jitclass

from numbsql import sqlite_udaf


@sqlite_udaf
@jitclass
class WinAvg:  # pragma: no cover
    total: float
    count: int

    def __init__(self) -> None:
        self.total = 0.0
        self.count = 0

    def step(self, value: Optional[float]) -> None:
        if value is not None:
            self.total += value
            self.count += 1

    def finalize(self) -> Optional[float]:
        count = self.count
        if count:
            return self.total / count
        return None

    def value(self) -> Optional[float]:
        return self.finalize()

    def inverse(self, value: Optional[float]) -> None:
        if value is not None:
            self.total -= value
            self.count -= 1

调用聚合函数

与标量函数类似，我们将函数注册到sqlite3.Connection对象。

>>> import sqlite3
>>> from numbsql import create_aggregate, create_function
>>> con = sqlite3.connect(":memory:")
>>> create_function(con, "add_one", 1, add_one)
>>> con.execute("SELECT add_one(1)").fetchall()
[(2,)]