datafusion_pywindowframe.py

import pyarrow
import pyarrow.compute
import datafusion
from datafusion import (
    col,
    udaf,
    Accumulator,
    SessionContext,
)
import datafusion.functions as F


def make_udaf(cls, typ=pyarrow.float64, volatility="stable", *, name=None):
    return udaf(
        cls,
        typ(),
        typ(),
        [typ()],
        volatility,
        name=name or cls.__name__.lower(),
    )


class CumulativeSum(Accumulator):

    def __init__(self):
        self._value = pyarrow.scalar(0.)

    def update(self, values: pyarrow.Array) -> None:
        # not nice since pyarrow scalars can't be summed yet. This breaks on `None`
        self._value = pyarrow.scalar(
            self._value.as_py() + pyarrow.compute.sum(values).as_py()
        )

    def merge(self, states: pyarrow.Array) -> None:
        # not nice since pyarrow scalars can't be summed yet. This breaks on `None`
        self._value = pyarrow.scalar(
            self._value.as_py() + pyarrow.compute.sum(states).as_py()
        )

    def state(self) -> pyarrow.Array:
        return pyarrow.array([self._value.as_py()])

    def evaluate(self) -> pyarrow.Scalar:
        return self._value

    def retract_batch(self, values: pyarrow.Array) -> None:
        self._value = pyarrow.scalar(
            self._value.as_py() - pyarrow.compute.sum(values).as_py()
        )

    def supports_retract_batch(self) -> bool:
        return True


class CumulativeMax(Accumulator):

    def __init__(self):
        self._value = pyarrow.scalar(float('-inf'))

    def update(self, values: pyarrow.Array) -> None:
        # not nice since pyarrow scalars can't be summed yet. This breaks on `None`
        self._value = pyarrow.scalar(
            max(self._value.as_py(), pyarrow.compute.max(values).as_py())
        )

    def merge(self, states: pyarrow.Array) -> None:
        # not nice since pyarrow scalars can't be summed yet. This breaks on `None`
        self._value = pyarrow.scalar(
            max(self._value.as_py(), pyarrow.compute.max(states).as_py())
        )

    def state(self) -> pyarrow.Array:
        return pyarrow.array([self._value.as_py()])

    def evaluate(self) -> pyarrow.Scalar:
        return self._value


def make_ctx():
    ctx = SessionContext()

    csum_udaf = make_udaf(CumulativeSum)
    cmax_udaf = make_udaf(CumulativeMax)

    ctx.register_udaf(csum_udaf)
    ctx.register_udaf(cmax_udaf)

    df = (
        ctx
        .create_dataframe(
            [[
                pyarrow.RecordBatch.from_arrays(
                    [pyarrow.array([1., 1., 1., 1., 2., 0.]), pyarrow.array([4, 2, -1, 10, 5, 6])],
                    names=["a", "b"],
                )
            ]],
            name='t',
        )
    )
    return ctx


def main(cls=CumulativeSum, preceding=None, following=0):
    ctx = make_ctx()
    fname = cls.__name__.lower()
    query_string = f"""
    SELECT a, b, {fname}(b) OVER (
        PARTITION BY a
        ORDER BY a, b
        ROWS BETWEEN {"UNBOUNDED" if preceding is None else preceding} PRECEDING and {following if following is not None else "UNBOUNDED"} FOLLOWING
    ) AS 'csum-b'
    FROM t"""
    # import pdb; pdb.set_trace()
    from_string = ctx.sql(query_string)
    from_df = (
        ctx.table('t')
        .select(
            col('a'),
            col('b'),
            F.alias(
                F.window(
                    fname,
                    [col('b')],
                    partition_by=[F.col('a')],
                    order_by=[F.order_by(col('a')), F.order_by(col('b'))],
                    window_frame=datafusion._internal.WindowFrame("rows", preceding, following),
                    ctx=ctx,
                ),
                'csum-b',
            ),
        )
    )
    print(from_string)
    print(from_df)


if __name__ == '__main__':
    cls = CumulativeSum
    main(cls, None, None)
    main(cls, 1, None)
    main(cls, None, 1)
    main(cls, 1, 1)
    main(cls, 0, 0)
    main(cls, 0, None)
    main(cls, None, 0)
    main(cls, 0, 1)
    main(cls, 1, 0)