vgi-python 0.8.0__py3-none-any.whl

vgi/_test_fixtures/aggregate/streaming.py

@@ -0,0 +1,192 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Streaming-partitioned aggregate fixtures.
+
+Exercise the ``streaming_partitioned`` opt-in: ``streaming_open`` /
+``streaming_chunk`` / ``streaming_close``. These are routed through the
+VGI DuckDB extension's custom streaming operator, which pipes input
+chunks straight to the worker without materialising the partition on
+the DuckDB side. State is bounded by partitions × per-partition state,
+not by row count — the structural answer to "running aggregate over
+unbounded ordered input."
+
+These fixtures are reference implementations for the protocol. Real
+production aggregates (e.g. ``portfolio_agg``) follow the same shape
+but with domain-specific state and I/O optimisations (Decimal128 buffer
+tricks, etc.).
+
+When the optimizer rule rejects a query (non-cumulative frame, EXCLUDE
+clause, DISTINCT/FILTER, etc.) DuckDB falls back to the standard
+windowed path — so all three of these classes also implement
+update/combine/finalize for plain GROUP BY usage and (optionally) the
+windowed callbacks. The streaming methods are additive.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Any
+
+import pyarrow as pa
+
+from vgi._test_fixtures.aggregate._common import SumState
+from vgi.aggregate_function import AggregateFunction
+from vgi.arguments import Param, Returns
+from vgi.metadata import DistinctDependence, NullHandling, OrderDependence
+from vgi.table_function import ProcessParams
+
+
+class StreamingSumFunction(AggregateFunction[SumState]):
+    """Streaming-partitioned running sum.
+
+    Cumulative across each `(PARTITION BY key)` group, in `ORDER BY` order.
+    For every input row, emits the running sum of the value column at that
+    row's position in its partition.
+
+    Also wired for ``GROUP BY`` via ``update`` / ``combine`` / ``finalize``,
+    so the same function works in both shapes::
+
+        -- streaming-partitioned (one running value per fill row):
+        SELECT k, v, vgi_streaming_sum(v) OVER (PARTITION BY k ORDER BY ts)
+        FROM trades;
+
+        -- group-by (one final value per partition):
+        SELECT k, vgi_streaming_sum(v) FROM trades GROUP BY k;
+
+    State persistence: the per-partition dict lives in worker memory in an
+    in-process LRU and is also persisted to ``FunctionStorage`` after each
+    chunk so a follow-up chunk landing on a different worker pool entry
+    can rehydrate. No special handling required from this class — the
+    framework does it.
+    """
+
+    class Meta:
+        name = "vgi_streaming_sum"
+        description = (
+            "Running sum across PARTITION BY keys via the streaming-partitioned "
+            "protocol. Each input row emits the cumulative sum at its position."
+        )
+        null_handling = NullHandling.DEFAULT
+        order_dependent = OrderDependence.NOT_ORDER_DEPENDENT
+        distinct_dependent = DistinctDependence.NOT_DISTINCT_DEPENDENT
+        supports_window = False
+        # Opt into the streaming-partitioned operator. The optimizer rule
+        # will route eligible OVER queries through it; ineligible shapes
+        # (sliding frames, EXCLUDE, DISTINCT, FILTER) fall back to the
+        # standard windowed path automatically.
+        streaming_partitioned = True
+
+    # ------------------------------------------------------------------
+    # GROUP BY path — required for plain aggregation queries.
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[None]) -> SumState:
+        return SumState()
+
+    @classmethod
+    def update(
+        cls,
+        states: dict[int, SumState],
+        group_ids: pa.Int64Array,
+        value: Annotated[pa.Int64Array, Param(doc="Column to sum")],
+    ) -> None:
+        table = pa.table({"gid": group_ids, "value": value})
+        grouped = table.group_by("gid").aggregate([("value", "sum")])
+        for i in range(grouped.num_rows):
+            gid: int = grouped.column("gid")[i].as_py()
+            v = grouped.column("value_sum")[i].as_py()
+            if v is not None:
+                states[gid] = SumState(total=states[gid].total + v)
+
+    @classmethod
+    def combine(cls, source: SumState, target: SumState, params: ProcessParams[None]) -> SumState:
+        return SumState(total=source.total + target.total)
+
+    @classmethod
+    def finalize(
+        cls,
+        group_ids: pa.Int64Array,
+        states: dict[int, SumState],
+        params: ProcessParams[None],
+    ) -> Annotated[pa.RecordBatch, Returns(pa.int64())]:
+        results = [(s.total if (s := states.get(gid.as_py())) is not None else None) for gid in group_ids]
+        return pa.record_batch({"result": pa.array(results, type=pa.int64())})
+
+    # ------------------------------------------------------------------
+    # Streaming-partitioned path.
+    # ------------------------------------------------------------------
+    #
+    # Three callbacks: open / chunk / close. The framework handles session
+    # lifecycle (allocates execution_id, persists to FunctionStorage,
+    # rehydrates across pool workers); user code only owns the in-memory
+    # state object.
+
+    @classmethod
+    def streaming_open(cls, params: ProcessParams[None]) -> dict[str, Any]:
+        # Session state. Free shape — anything picklable. The framework
+        # passes this object back to streaming_chunk and streaming_close
+        # unchanged. For multi-partition aggregates, hold a per-partition
+        # dict here; for single-partition, just hold the running scalar.
+        return {
+            # partition_key_tuple -> running int sum
+            "partition_states": {},
+        }
+
+    @classmethod
+    def streaming_chunk(
+        cls,
+        chunk: pa.RecordBatch,
+        streaming_state: dict[str, Any],
+        partition_key_count: int,
+        order_key_count: int,
+        params: ProcessParams[None],
+    ) -> pa.Array[Any]:
+        # Column layout from the operator:
+        #   [partition_key_cols..., order_key_cols..., value_cols...]
+        # We don't actually need the order keys at runtime here — the
+        # input arrives in (partition, order) order already, so cumulative
+        # state is naturally correct.
+        n = chunk.num_rows
+        value_idx = partition_key_count + order_key_count
+
+        if partition_key_count > 0:
+            pk_columns = [chunk.column(i).to_pylist() for i in range(partition_key_count)]
+        else:
+            pk_columns = []
+        values = chunk.column(value_idx).to_pylist()
+
+        partition_states: dict[Any, int] = streaming_state["partition_states"]
+
+        # Returns one cumulative-sum int per input row. NULL value rows
+        # leave state unchanged but still emit the current sum (matches
+        # the GROUP BY path's NullHandling.DEFAULT semantics).
+        out: list[int] = [0] * n
+        for i in range(n):
+            if partition_key_count == 0:
+                key: Any = ()
+            elif partition_key_count == 1:
+                key = pk_columns[0][i]
+            else:
+                key = tuple(col[i] for col in pk_columns)
+
+            running = partition_states.get(key, 0)
+            v = values[i]
+            if v is not None:
+                running += v
+                partition_states[key] = running
+            out[i] = running
+
+        return pa.array(out, type=pa.int64())
+
+    @classmethod
+    def streaming_close(
+        cls,
+        streaming_state: dict[str, Any],
+        params: ProcessParams[None],
+    ) -> None:
+        # Cleanup hook. For this fixture there's nothing to release;
+        # state is plain Python objects that GC collects when the
+        # session is dropped from the framework's cache. Real
+        # implementations might release file handles, close DB
+        # connections, or flush logs here.
+        return None

vgi/_test_fixtures/aggregate/varargs.py

@@ -0,0 +1,75 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""SumAllFunction — varargs aggregate (sums any number of numeric columns)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Annotated
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass, ArrowType
+
+from vgi.aggregate_function import AggregateFunction
+from vgi.arguments import Param, Returns
+from vgi.metadata import DistinctDependence, NullHandling, OrderDependence
+from vgi.table_function import ProcessParams
+
+
+@dataclass(kw_only=True)
+class SumAllState(ArrowSerializableDataclass):
+    total: Annotated[float, ArrowType(pa.float64())] = 0.0
+
+
+class SumAllFunction(AggregateFunction[SumAllState]):
+    """Sum all numeric columns — demonstrates varargs aggregate.
+
+    Accepts any number of numeric columns and sums them all together.
+    SQL: ``SELECT vgi_sum_all(a, b, c) FROM t GROUP BY category``
+    """
+
+    class Meta:
+        name = "vgi_sum_all"
+        description = "Sum all numeric columns"
+        null_handling = NullHandling.DEFAULT
+        order_dependent = OrderDependence.NOT_ORDER_DEPENDENT
+        distinct_dependent = DistinctDependence.NOT_DISTINCT_DEPENDENT
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[None]) -> SumAllState:
+        return SumAllState()
+
+    @classmethod
+    def update(
+        cls,
+        states: dict[int, SumAllState],
+        group_ids: pa.Int64Array,
+        columns: Annotated[pa.Array, Param(doc="Numeric columns to sum", varargs=True)],  # type: ignore[type-arg]
+    ) -> None:
+        for i in range(len(group_ids)):
+            gid: int = group_ids[i].as_py()
+            row_total = 0.0
+            for col in columns:
+                val = col[i].as_py()
+                if val is not None:
+                    row_total += float(val)
+            states[gid] = SumAllState(total=states[gid].total + row_total)
+
+    @classmethod
+    def combine(cls, source: SumAllState, target: SumAllState, params: ProcessParams[None]) -> SumAllState:
+        return SumAllState(total=source.total + target.total)
+
+    @classmethod
+    def finalize(
+        cls,
+        group_ids: pa.Int64Array,
+        states: dict[int, SumAllState],
+        params: ProcessParams[None],
+    ) -> Annotated[pa.RecordBatch, Returns(pa.float64())]:
+        results = [s.total if (s := states[gid.as_py()]) is not None else None for gid in group_ids]
+        return pa.record_batch({"result": pa.array(results, type=pa.float64())})
+
+
+# ---------------------------------------------------------------------------
+# DynamicAggregateFunction — aggregate behavior defined by Python code string
+# ---------------------------------------------------------------------------

vgi/_test_fixtures/aggregate/window.py

@@ -0,0 +1,380 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Windowed aggregate fixtures (window_sum, window_median, window_listagg)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Annotated, Any
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+
+from vgi._test_fixtures.aggregate._common import ListAggState, SumState
+from vgi.aggregate_function import AggregateFunction, WindowPartition
+from vgi.arguments import Param, Returns
+from vgi.metadata import DistinctDependence, NullHandling, OrderDependence
+from vgi.table_function import ProcessParams
+
+
+@dataclass(kw_only=True)
+class _EmptyWindowState(ArrowSerializableDataclass):
+    """Placeholder for functions that don't need derived per-partition state."""
+
+    pass
+
+
+class WindowSumFunction(AggregateFunction[SumState]):
+    """Windowed running-sum — demonstrates a simple window() callback.
+
+    Also implements update/combine/finalize so the function still works in
+    plain ``GROUP BY`` contexts (DuckDB picks the window path automatically
+    via ``WindowCustomAggregator::CanAggregate``).
+
+    SQL::
+
+        SELECT x, vgi_window_sum(x) OVER (ORDER BY x ROWS BETWEEN 2 PRECEDING AND CURRENT ROW)
+        FROM generate_series(1, 10) t(x);
+    """
+
+    class Meta:
+        name = "vgi_window_sum"
+        description = "Windowed sum that uses the per-partition window() callback"
+        null_handling = NullHandling.DEFAULT
+        order_dependent = OrderDependence.NOT_ORDER_DEPENDENT
+        distinct_dependent = DistinctDependence.NOT_DISTINCT_DEPENDENT
+        supports_window = True
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[None]) -> SumState:
+        return SumState()
+
+    @classmethod
+    def update(
+        cls,
+        states: dict[int, SumState],
+        group_ids: pa.Int64Array,
+        value: Annotated[pa.Int64Array, Param(doc="Column to sum")],
+    ) -> None:
+        table = pa.table({"gid": group_ids, "value": value})
+        grouped = table.group_by("gid").aggregate([("value", "sum")])
+        for i in range(grouped.num_rows):
+            gid: int = grouped.column("gid")[i].as_py()
+            val = grouped.column("value_sum")[i].as_py()
+            if val is not None:
+                states[gid] = SumState(total=states[gid].total + val)
+
+    @classmethod
+    def combine(cls, source: SumState, target: SumState, params: ProcessParams[None]) -> SumState:
+        return SumState(total=source.total + target.total)
+
+    @classmethod
+    def finalize(
+        cls,
+        group_ids: pa.Int64Array,
+        states: dict[int, SumState],
+        params: ProcessParams[None],
+    ) -> Annotated[pa.RecordBatch, Returns(pa.int64())]:
+        results = [s.total if (s := states[gid.as_py()]) is not None else None for gid in group_ids]
+        return pa.record_batch({"result": pa.array(results, type=pa.int64())})
+
+    # --- Window path ---
+
+    @classmethod
+    def window(
+        cls,
+        rid: int,
+        subframes: list[tuple[int, int]],
+        partition: WindowPartition,
+        window_state: Any,
+        params: ProcessParams[None],
+    ) -> int | None:
+        import pyarrow.compute as pc
+
+        value_col = partition.inputs.column(0)
+        total = 0
+        any_valid = False
+        for begin, end in subframes:
+            if end <= begin:
+                continue
+            slice_ = value_col.slice(begin, end - begin)
+            if partition.filter_mask is not None:
+                mask = partition.filter_mask.slice(begin, end - begin)
+                slice_ = slice_.filter(mask)
+            s = pc.sum(slice_)
+            if s.is_valid:
+                total += s.as_py()
+                any_valid = True
+        return total if any_valid else None
+
+
+class WindowMedianFunction(AggregateFunction[_EmptyWindowState]):
+    """Windowed median — non-incremental, benefits from caching the partition.
+
+    Uses the window() callback exclusively (no incremental update path makes
+    sense for median). Falls back to a naive GROUP BY implementation via
+    update/combine/finalize that collects values in a single string field.
+
+    SQL::
+
+        SELECT x, vgi_window_median(x) OVER (ORDER BY x ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING)
+        FROM generate_series(1, 20) t(x);
+    """
+
+    class Meta:
+        name = "vgi_window_median"
+        description = "Windowed median (window() callback demonstrates non-incremental aggregates)"
+        null_handling = NullHandling.DEFAULT
+        order_dependent = OrderDependence.NOT_ORDER_DEPENDENT
+        distinct_dependent = DistinctDependence.NOT_DISTINCT_DEPENDENT
+        supports_window = True
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[None]) -> _EmptyWindowState:
+        return _EmptyWindowState()
+
+    @classmethod
+    def update(
+        cls,
+        states: dict[int, _EmptyWindowState],
+        group_ids: pa.Int64Array,
+        value: Annotated[pa.DoubleArray, Param(doc="Column to compute median of")],
+    ) -> None:
+        # GROUP BY path not the primary use — kept only so the function works
+        # when used outside an OVER clause. Caller must not expect exact
+        # semantics for huge groups.
+        pass
+
+    @classmethod
+    def combine(
+        cls, source: _EmptyWindowState, target: _EmptyWindowState, params: ProcessParams[None]
+    ) -> _EmptyWindowState:
+        return target
+
+    @classmethod
+    def finalize(
+        cls,
+        group_ids: pa.Int64Array,
+        states: dict[int, _EmptyWindowState],
+        params: ProcessParams[None],
+    ) -> Annotated[pa.RecordBatch, Returns(pa.float64())]:
+        results = [None] * len(group_ids)
+        return pa.record_batch({"result": pa.array(results, type=pa.float64())})
+
+    @classmethod
+    def window(
+        cls,
+        rid: int,
+        subframes: list[tuple[int, int]],
+        partition: WindowPartition,
+        window_state: Any,
+        params: ProcessParams[None],
+    ) -> float | None:
+        value_col = partition.inputs.column(0)
+        values: list[float] = []
+        for begin, end in subframes:
+            if end <= begin:
+                continue
+            slice_ = value_col.slice(begin, end - begin)
+            if partition.filter_mask is not None:
+                mask = partition.filter_mask.slice(begin, end - begin)
+                slice_ = slice_.filter(mask)
+            for v in slice_.to_pylist():
+                if v is not None:
+                    values.append(float(v))
+        if not values:
+            return None
+        values.sort()
+        n = len(values)
+        mid = n // 2
+        if n % 2 == 1:
+            return values[mid]
+        return (values[mid - 1] + values[mid]) / 2.0
+
+
+class WindowListAggFunction(AggregateFunction[ListAggState]):
+    """Windowed ORDER_DEPENDENT aggregate — demonstrates the fallback handoff.
+
+    For ``vgi_window_listagg(s) OVER (ORDER BY x ...)`` DuckDB picks our
+    ``window()`` callback (arg_orders is empty; frame ordering comes from
+    the OVER clause).
+
+    For ``vgi_window_listagg(s ORDER BY x) OVER (...)`` DuckDB's
+    ``WindowCustomAggregator::CanAggregate`` rejects the window path
+    because ``wexpr.arg_orders`` is non-empty, and falls back to
+    update/combine/finalize. The result is still correct — just slower.
+    """
+
+    class Meta:
+        name = "vgi_window_listagg"
+        description = "Windowed string concat (ORDER_DEPENDENT; tests fallback handoff)"
+        null_handling = NullHandling.DEFAULT
+        order_dependent = OrderDependence.ORDER_DEPENDENT
+        distinct_dependent = DistinctDependence.DISTINCT_DEPENDENT
+        supports_window = True
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[None]) -> ListAggState:
+        return ListAggState()
+
+    @classmethod
+    def update(
+        cls,
+        states: dict[int, ListAggState],
+        group_ids: pa.Int64Array,
+        value: Annotated[pa.StringArray, Param(doc="String column")],
+    ) -> None:
+        for i in range(len(group_ids)):
+            gid: int = group_ids[i].as_py()
+            val = value[i].as_py()
+            if val is not None:
+                s = states[gid]
+                if s.values:
+                    states[gid] = ListAggState(values=s.values + "," + val)
+                else:
+                    states[gid] = ListAggState(values=val)
+
+    @classmethod
+    def combine(cls, source: ListAggState, target: ListAggState, params: ProcessParams[None]) -> ListAggState:
+        if source.values and target.values:
+            return ListAggState(values=target.values + "," + source.values)
+        return ListAggState(values=target.values or source.values)
+
+    @classmethod
+    def finalize(
+        cls,
+        group_ids: pa.Int64Array,
+        states: dict[int, ListAggState],
+        params: ProcessParams[None],
+    ) -> Annotated[pa.RecordBatch, Returns(pa.string())]:
+        results = [s.values or None if (s := states[gid.as_py()]) is not None else None for gid in group_ids]
+        return pa.record_batch({"result": pa.array(results, type=pa.string())})
+
+    @classmethod
+    def window(
+        cls,
+        rid: int,
+        subframes: list[tuple[int, int]],
+        partition: WindowPartition,
+        window_state: Any,
+        params: ProcessParams[None],
+    ) -> str | None:
+        value_col = partition.inputs.column(0)
+        parts: list[str] = []
+        for begin, end in subframes:
+            if end <= begin:
+                continue
+            slice_ = value_col.slice(begin, end - begin)
+            if partition.filter_mask is not None:
+                mask = partition.filter_mask.slice(begin, end - begin)
+                slice_ = slice_.filter(mask)
+            for v in slice_.to_pylist():
+                if v is not None:
+                    parts.append(v)
+        return ",".join(parts) if parts else None
+
+
+class WindowSumBatchFunction(AggregateFunction[SumState]):
+    """Windowed running-sum returning a pre-built ``pa.Array``.
+
+    Overrides ``window_batch`` to return a pre-built ``pa.Array`` rather
+    than a Python list.
+
+    Functionally equivalent to :class:`WindowSumFunction`. The point of this
+    fixture is to exercise the framework's polymorphic batch return: when
+    user code returns a ``pa.Array``, the worker should ship it directly
+    without round-tripping through ``pa.array(list, type=...)``.
+
+    Used by the unit tests for ``window_batch`` to confirm the dispatcher
+    accepts both a list and a pa.Array, and that the pa.Array path
+    produces identical answers.
+    """
+
+    class Meta:
+        name = "vgi_window_sum_batch"
+        description = "Windowed sum demonstrating window_batch returning pa.Array"
+        null_handling = NullHandling.DEFAULT
+        order_dependent = OrderDependence.NOT_ORDER_DEPENDENT
+        distinct_dependent = DistinctDependence.NOT_DISTINCT_DEPENDENT
+        supports_window = True
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[None]) -> SumState:
+        return SumState()
+
+    @classmethod
+    def update(
+        cls,
+        states: dict[int, SumState],
+        group_ids: pa.Int64Array,
+        value: Annotated[pa.Int64Array, Param(doc="Column to sum")],
+    ) -> None:
+        table = pa.table({"gid": group_ids, "value": value})
+        grouped = table.group_by("gid").aggregate([("value", "sum")])
+        for i in range(grouped.num_rows):
+            gid: int = grouped.column("gid")[i].as_py()
+            val = grouped.column("value_sum")[i].as_py()
+            if val is not None:
+                states[gid] = SumState(total=states[gid].total + val)
+
+    @classmethod
+    def combine(cls, source: SumState, target: SumState, params: ProcessParams[None]) -> SumState:
+        return SumState(total=source.total + target.total)
+
+    @classmethod
+    def finalize(
+        cls,
+        group_ids: pa.Int64Array,
+        states: dict[int, SumState],
+        params: ProcessParams[None],
+    ) -> Annotated[pa.RecordBatch, Returns(pa.int64())]:
+        results = [s.total if (s := states[gid.as_py()]) is not None else None for gid in group_ids]
+        return pa.record_batch({"result": pa.array(results, type=pa.int64())})
+
+    @classmethod
+    def window(
+        cls,
+        rid: int,
+        subframes: list[tuple[int, int]],
+        partition: WindowPartition,
+        window_state: Any,
+        params: ProcessParams[None],
+    ) -> int | None:
+        # Single-row fallback (still required so plain window() invocations
+        # work in unit tests). Production callers go through window_batch.
+        return cls._sum_one(subframes, partition)
+
+    @classmethod
+    def window_batch(
+        cls,
+        row_ids: list[int],
+        subframes: list[list[tuple[int, int]]],
+        partition: WindowPartition,
+        window_state: Any,
+        params: ProcessParams[None],
+    ) -> pa.Array[Any]:
+        out = [cls._sum_one(frames, partition) for frames in subframes]
+        return pa.array(out, type=pa.int64())
+
+    @staticmethod
+    def _sum_one(
+        subframes: list[tuple[int, int]],
+        partition: WindowPartition,
+    ) -> int | None:
+        import pyarrow.compute as pc
+
+        value_col = partition.inputs.column(0)
+        total = 0
+        any_valid = False
+        for begin, end in subframes:
+            if end <= begin:
+                continue
+            slice_ = value_col.slice(begin, end - begin)
+            if partition.filter_mask is not None:
+                mask = partition.filter_mask.slice(begin, end - begin)
+                slice_ = slice_.filter(mask)
+            s = pc.sum(slice_)
+            if s.is_valid:
+                total += s.as_py()
+                any_valid = True
+        return total if any_valid else None