vgi-python 0.8.0__py3-none-any.whl

vgi/_test_fixtures/table/partition_columns_broken.py

@@ -0,0 +1,304 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Deliberately-broken PartitionColumns fixtures for v2 contract testing.
+
+Each fixture violates one specific clause of the PartitionColumns contract
+documented at ``vgi/_test_fixtures/table/partition_columns.py`` /
+``vgi/src/vgi_table_function_impl.cpp::InstallBatch``.
+
+* :class:`BrokenMissingPartitionValuesFunction` — declares
+  ``partition_kind = SINGLE_VALUE_PARTITIONS`` and an annotated bind-
+  schema field, but bypasses the framework's wrapper validation by
+  reaching the inner OutputCollector directly. The C++ extension's
+  ``InstallBatch`` catches the missing ``vgi_partition_values#b64``
+  metadata.
+
+* :class:`BrokenPartitionMinNeqMaxFunction` — declares
+  ``SINGLE_VALUE_PARTITIONS`` but emits a chunk whose partition
+  column has multiple distinct values. The framework's auto-extract
+  path would catch this client-side, so the fixture supplies an
+  explicit ``partition_values={"col": (min, max)}`` with min != max
+  to defeat the worker check and reach the C++ defense-in-depth
+  validation in ``InstallBatch``. The C++ check is what guarantees
+  this fires on release builds where DuckDB's own
+  ``BatchedDataCollection::Append`` assertion is compiled out.
+
+* :class:`BrokenPartitionValuesNoAnnotationFunction` — no
+  ``vgi.partition_column`` annotation on any bind-schema field and
+  ``partition_kind = NOT_PARTITIONED``, but the worker passes
+  ``partition_values=`` on ``out.emit`` anyway. The framework
+  rejects with RuntimeError at the emit site.
+
+* :class:`BrokenPartitionColumnAbsentFromBatchFunction` — declares
+  ``partition_kind`` and annotates a bind-schema field, but the
+  worker emits a batch that DOES NOT include that column AND does
+  not supply an explicit ``partition_values=`` override. The
+  framework's ``_merge_partition_values`` raises RuntimeError at
+  the emit site (auto-extract can't find the column).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Annotated, Any, ClassVar, cast
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+from vgi_rpc.rpc import OutputCollector
+
+from vgi._test_fixtures.table._common import _cardinality_from_count
+from vgi.arguments import Arg
+from vgi.metadata import PartitionKind
+from vgi.protocol import VgiOutputCollector
+from vgi.schema_utils import partition_field
+from vgi.table_function import (
+    ProcessParams,
+    TableFunctionGenerator,
+    bind_fixed_schema,
+)
+
+
+@dataclass(slots=True, frozen=True)
+class _BrokenArgs:
+    count: Annotated[int, Arg(0, doc="Rows to attempt to emit", ge=1)]
+
+
+@dataclass(kw_only=True)
+class _BrokenState(ArrowSerializableDataclass):
+    emitted: bool = False
+
+
+# =============================================================================
+# 1. Missing partition_values metadata (C++ side raises)
+# =============================================================================
+
+
+@bind_fixed_schema
+@_cardinality_from_count
+class BrokenMissingPartitionValuesFunction(TableFunctionGenerator[_BrokenArgs, _BrokenState]):
+    """Opt-in declared, but worker bypasses framework metadata merge."""
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
+        [
+            partition_field("country", pa.string()),
+            pa.field("sales", pa.int64()),
+        ]
+    )
+
+    class Meta:
+        name = "broken_missing_partition_values"
+        description = (
+            "DELIBERATELY BROKEN: declares partition_kind + partition-annotated "
+            "field but emits a data batch without vgi_partition_values#b64 "
+            "metadata. C++ extension's contract check raises."
+        )
+        categories = ["testing", "broken"]
+        partition_kind = PartitionKind.SINGLE_VALUE_PARTITIONS
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_BrokenArgs]) -> _BrokenState:
+        return _BrokenState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_BrokenArgs],
+        state: _BrokenState,
+        out: OutputCollector,
+    ) -> None:
+        if state.emitted:
+            out.finish()
+            return
+        batch = pa.RecordBatch.from_pydict(
+            {"country": ["US"] * params.args.count, "sales": list(range(params.args.count))},
+            schema=cls.FIXED_SCHEMA,
+        )
+        # Reach into the wrapper stack and call the innermost inner
+        # directly. This is what makes the fixture "broken": the
+        # framework's _merge_partition_values validator never runs, so
+        # the data batch has no vgi_partition_values#b64 metadata and
+        # the C++ extension's InstallBatch contract check fires.
+        # Same pattern as v1's broken_missing_batch_index_tag fixture.
+        inner = out
+        while hasattr(inner, "_inner"):
+            inner = inner._inner
+        inner.emit(batch)
+        state.emitted = True
+
+
+# =============================================================================
+# 2. SINGLE_VALUE with min != max (C++ defense-in-depth raises)
+# =============================================================================
+
+
+@bind_fixed_schema
+@_cardinality_from_count
+class BrokenPartitionMinNeqMaxFunction(TableFunctionGenerator[_BrokenArgs, _BrokenState]):
+    """SINGLE_VALUE_PARTITIONS but emit min != max via explicit override."""
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
+        [
+            partition_field("country", pa.string()),
+            pa.field("sales", pa.int64()),
+        ]
+    )
+
+    class Meta:
+        name = "broken_partition_min_neq_max"
+        description = (
+            "DELIBERATELY BROKEN: declares SINGLE_VALUE_PARTITIONS but "
+            "supplies an explicit partition_values override with "
+            "min != max. The framework's wrapper validation doesn't "
+            "compare min vs max for SINGLE_VALUE; the C++ extension's "
+            "defense-in-depth check in InstallBatch raises."
+        )
+        categories = ["testing", "broken"]
+        partition_kind = PartitionKind.SINGLE_VALUE_PARTITIONS
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_BrokenArgs]) -> _BrokenState:
+        return _BrokenState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_BrokenArgs],
+        state: _BrokenState,
+        out: OutputCollector,
+    ) -> None:
+        if state.emitted:
+            out.finish()
+            return
+        # Single-valued country column at the data level (so the
+        # framework's auto-extract WOULD pass), but the explicit
+        # override forces min != max — defeats the framework check
+        # and reaches C++ defense-in-depth.
+        batch = pa.RecordBatch.from_pydict(
+            {"country": ["US"] * params.args.count, "sales": list(range(params.args.count))},
+            schema=cls.FIXED_SCHEMA,
+        )
+        cast(VgiOutputCollector, out).emit(
+            batch,
+            partition_values={
+                "country": (
+                    pa.scalar("US", type=pa.string()),
+                    pa.scalar("BR", type=pa.string()),  # max != min — bug
+                ),
+            },
+        )
+        state.emitted = True
+
+
+# =============================================================================
+# 3. partition_values kwarg without any annotated field (worker-side raise)
+# =============================================================================
+
+
+@bind_fixed_schema
+@_cardinality_from_count
+class BrokenPartitionValuesNoAnnotationFunction(TableFunctionGenerator[_BrokenArgs, _BrokenState]):
+    """No partition annotation, but worker passes partition_values=."""
+
+    # No partition_field() — bind schema has no partition columns.
+    # cast: mypy joins Field[StringType] + Field[Int64Type] to Field[object];
+    # the runtime list is a plain list of pa.Field.
+    FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
+        cast(
+            "list[pa.Field[Any]]",
+            [pa.field("country", pa.string()), pa.field("sales", pa.int64())],
+        )
+    )
+
+    class Meta:
+        name = "broken_partition_values_no_annotation"
+        description = (
+            "DELIBERATELY BROKEN: no field carries vgi.partition_column "
+            "metadata (and partition_kind defaults to NOT_PARTITIONED), "
+            "but the worker passes partition_values= on out.emit. The "
+            "framework rejects with RuntimeError before the wire."
+        )
+        categories = ["testing", "broken"]
+        # No partition_kind setting — defaults to NOT_PARTITIONED.
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_BrokenArgs]) -> _BrokenState:
+        return _BrokenState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_BrokenArgs],
+        state: _BrokenState,
+        out: OutputCollector,
+    ) -> None:
+        if state.emitted:
+            out.finish()
+            return
+        batch = pa.RecordBatch.from_pydict(
+            {"country": ["US"] * params.args.count, "sales": list(range(params.args.count))},
+            schema=cls.FIXED_SCHEMA,
+        )
+        cast(VgiOutputCollector, out).emit(
+            batch,
+            partition_values={
+                "country": (
+                    pa.scalar("US", type=pa.string()),
+                    pa.scalar("US", type=pa.string()),
+                ),
+            },
+        )
+        state.emitted = True
+
+
+# =============================================================================
+# 4. Annotated column missing from batch, no explicit override (worker-side raise)
+# =============================================================================
+
+
+@bind_fixed_schema
+@_cardinality_from_count
+class BrokenPartitionColumnAbsentFromBatchFunction(TableFunctionGenerator[_BrokenArgs, _BrokenState]):
+    """Annotated partition column not in emitted batch, no override."""
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
+        [
+            partition_field("category", pa.string()),
+            pa.field("revenue", pa.int64()),
+        ]
+    )
+
+    class Meta:
+        name = "broken_partition_column_absent_from_batch"
+        description = (
+            "DELIBERATELY BROKEN: declares partition_kind on "
+            "'category' but emits a batch without 'category' AND "
+            "doesn't supply an explicit partition_values override. The "
+            "framework's auto-extract fails with RuntimeError before "
+            "the wire."
+        )
+        categories = ["testing", "broken"]
+        partition_kind = PartitionKind.SINGLE_VALUE_PARTITIONS
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_BrokenArgs]) -> _BrokenState:
+        return _BrokenState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_BrokenArgs],
+        state: _BrokenState,
+        out: OutputCollector,
+    ) -> None:
+        if state.emitted:
+            out.finish()
+            return
+        # Emit a batch WITHOUT 'category'. Framework's auto-extract
+        # tries to read batch.column('category') and raises.
+        batch_schema = pa.schema([pa.field("revenue", pa.int64())])
+        batch = pa.RecordBatch.from_pydict(
+            {"revenue": list(range(params.args.count))},
+            schema=batch_schema,
+        )
+        out.emit(batch)
+        state.emitted = True

vgi/_test_fixtures/table/profiling_example.py

@@ -0,0 +1,195 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Example function exercising the dynamic_to_string callback.
+
+``ProfilingDemoFunction`` demonstrates the recommended persistence
+pattern for diagnostics that should surface under ``EXPLAIN ANALYZE``:
+
+1. ``process()`` keeps per-stream counters in user state (rows,
+   batches, start time), and after every tick writes a serialized
+   snapshot via ``params.storage.put(bytes)``.
+2. ``dynamic_to_string()`` constructs a ``BoundStorage`` for the
+   given ``execution_id``, calls ``collect()`` to gather every
+   worker's last snapshot, and sums them.
+
+``BoundStorage`` defaults to the sqlite-backed shared storage (see
+CLAUDE.md → ``VGI_WORKER_SHARED_STORAGE``), so the pattern works across
+worker processes — both subprocess transport and HTTP transport with
+``max_workers > 1``. No in-memory class state is involved.
+"""
+
+from __future__ import annotations
+
+import struct
+import time
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Annotated, ClassVar
+
+import numpy as np
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+from vgi_rpc.rpc import OutputCollector
+
+from vgi._test_fixtures.table._common import CountBatchArgs
+from vgi.arguments import Arg
+from vgi.function_storage import BoundStorage
+from vgi.metadata import FunctionExample
+from vgi.schema_utils import schema
+from vgi.table_function import (
+    BindParams,
+    ProcessParams,
+    TableCardinality,
+    TableFunctionGenerator,
+    bind_fixed_schema,
+    init_single_worker,
+)
+
+
+@dataclass(frozen=True)
+class ProfilingDemoArgs(CountBatchArgs):
+    """Arguments for ProfilingDemoFunction."""
+
+    increment: Annotated[int, Arg("increment", default=1, doc="Step between values", ge=1)]
+
+
+@dataclass(kw_only=True)
+class ProfilingState(ArrowSerializableDataclass):
+    """Per-stream counters."""
+
+    remaining: int
+    current_index: int = 0
+    rows_emitted: int = 0
+    batches_emitted: int = 0
+    started_at_ns: int = 0
+
+
+# Serialized snapshot wire format: three little-endian uint64s
+# (rows, batches, elapsed_us). Compact; survives multi-worker collect().
+_SNAPSHOT = struct.Struct("<QQQ")
+
+
+def _pack_snapshot(rows: int, batches: int, elapsed_us: int) -> bytes:
+    return _SNAPSHOT.pack(rows, batches, elapsed_us)
+
+
+def _unpack_snapshot(data: bytes) -> tuple[int, int, int]:
+    return _SNAPSHOT.unpack(data)
+
+
+@init_single_worker
+@bind_fixed_schema
+class ProfilingDemoFunction(TableFunctionGenerator[ProfilingDemoArgs, ProfilingState]):
+    """Sequence generator that publishes per-execution metrics under EXPLAIN ANALYZE.
+
+    Output is identical to ``sequence(count, batch_size, increment)``.
+    Additionally tracks ``rows_produced``, ``batches_emitted``, and
+    ``elapsed_ms`` and surfaces them via ``dynamic_to_string``.
+    """
+
+    FunctionArguments = ProfilingDemoArgs
+
+    class Meta:
+        """Metadata for ProfilingDemoFunction."""
+
+        name = "profiling_demo"
+        description = "Sequence generator publishing diagnostics under EXPLAIN ANALYZE"
+        categories = ["generator", "utility"]
+        examples = [
+            FunctionExample(
+                sql="EXPLAIN ANALYZE SELECT count(*) FROM profiling_demo(500)",
+                description="Run with diagnostics surfaced as Extra Info",
+            ),
+        ]
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = schema(n=pa.int64())
+    NUMPY_DTYPE: ClassVar[type[np.generic]] = np.int64
+
+    @classmethod
+    def cardinality(cls, params: BindParams[ProfilingDemoArgs]) -> TableCardinality:
+        count = params.args.count
+        return TableCardinality(estimate=count, max=count)
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[ProfilingDemoArgs]) -> ProfilingState:
+        return ProfilingState(
+            remaining=params.args.count,
+            started_at_ns=time.monotonic_ns(),
+        )
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[ProfilingDemoArgs],
+        state: ProfilingState,
+        out: OutputCollector,
+    ) -> None:
+        if state.remaining <= 0:
+            # Final write so dynamic_to_string sees the totals even after
+            # the stream finishes. One row per OS pid via state_put under
+            # namespace b"profile" — dynamic_to_string drains them all.
+            elapsed_us = (time.monotonic_ns() - state.started_at_ns) // 1000
+            import os as _os
+
+            params.storage.state_put(
+                b"profile",
+                BoundStorage.pack_int_key(_os.getpid()),
+                _pack_snapshot(state.rows_emitted, state.batches_emitted, elapsed_us),
+            )
+            out.finish()
+            return
+        batch_size = params.args.batch_size
+        size = min(state.remaining, batch_size)
+        increment = params.args.increment
+        values = np.arange(
+            state.current_index * increment,
+            (state.current_index + size) * increment,
+            increment,
+            dtype=cls.NUMPY_DTYPE,
+        )
+        out.emit(pa.RecordBatch.from_arrays([pa.array(values)], schema=params.output_schema))
+        state.current_index += size
+        state.remaining -= size
+        state.rows_emitted += size
+        state.batches_emitted += 1
+
+        # Per-tick snapshot — overwrites this worker's slot. The dispatcher's
+        # state_drain on dynamic_to_string sums one snapshot per worker pid.
+        elapsed_us = (time.monotonic_ns() - state.started_at_ns) // 1000
+        import os as _os
+
+        params.storage.state_put(
+            b"profile",
+            BoundStorage.pack_int_key(_os.getpid()),
+            _pack_snapshot(state.rows_emitted, state.batches_emitted, elapsed_us),
+        )
+
+    @classmethod
+    def dynamic_to_string(
+        cls,
+        params: BindParams[ProfilingDemoArgs],
+        execution_id: bytes,
+    ) -> Mapping[str, str]:
+        # BindParams doesn't carry a BoundStorage (no execution_id at bind
+        # time). Construct one with the execution_id we received.
+        storage = BoundStorage(cls.storage, execution_id, request=params.bind_call)
+        try:
+            # state_drain returns (key, value) pairs; we only want the values.
+            snapshots = [v for _, v in storage.state_drain(b"profile")]
+        except Exception:
+            return {}
+        if not snapshots:
+            return {}
+        rows = 0
+        batches = 0
+        elapsed_us = 0
+        for blob in snapshots:
+            r, b, e = _unpack_snapshot(blob)
+            rows += r
+            batches += b
+            elapsed_us = max(elapsed_us, e)
+        return {
+            "rows_produced": str(rows),
+            "batches_emitted": str(batches),
+            "elapsed_ms": f"{elapsed_us / 1000.0:.2f}",
+        }