vgi-python 0.8.0__py3-none-any.whl

vgi/_test_fixtures/table/tt_pushdown.py

@@ -0,0 +1,191 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Time-travel + filter-pushdown fixtures.
+
+These back ``test/sql/integration/table/time_travel_pushdown.test`` in the C++
+repo, which asserts that a table can be partition-pruned (filter pushdown) AND
+time-travelled (``AT (VERSION|TIMESTAMP ...)``) in the *same* query — for tables
+declared **both** ways:
+
+- ``tt_pushdown_fn`` — **function-backed** (``Table(function=...)``). It reads the
+  AT clause from the init request (``params.at_value`` →
+  ``init_call.bind_call.at_value``), which only works once the framework threads
+  AT onto the bind request. Before that fix this table cannot see AT at all, so
+  ``seen_version`` collapses to the current version — the regression guard.
+- ``tt_pushdown_cols`` — **columns-based** (``Table(columns=...)`` routed via
+  ``table_scan_function_get``). It gets the resolved version as a scan-function
+  **argument** (the ``versioned_data`` mechanism) — the native columns-based AT
+  path, here to prove the bind-side change didn't regress it.
+
+Both echo ``seen_version`` (the version they actually scanned) and
+``pushed_filters`` (the SQL-like predicate DuckDB pushed down), so one query can
+assert both signals at once. ``auto_apply_filters`` keeps the result set correct.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Annotated, Any, ClassVar
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+from vgi_rpc.rpc import OutputCollector
+
+from vgi._test_fixtures.table._common import _EmptyArgs
+from vgi._test_fixtures.table.filters import _format_pushed_filters
+from vgi.arguments import Arg
+from vgi.schema_utils import schema
+from vgi.table_function import (
+    ProcessParams,
+    TableFunctionGenerator,
+    bind_fixed_schema,
+    init_single_worker,
+)
+
+# Output schema is version-INDEPENDENT (no schema evolution in scope): only the
+# row *data* changes per version, so the function-backed table stays inline-bound.
+_TT_SCHEMA: pa.Schema = schema(
+    id=pa.int64(),
+    val=pa.int64(),
+    seen_version=pa.int64(),
+    pushed_filters=pa.string(),
+)
+
+# Per-version row ids (val = id * 10). v2 is a strict superset of v1, so a row
+# count difference cleanly proves which version was scanned.
+_TT_VERSION_IDS: dict[int, list[int]] = {
+    1: [1, 2, 3, 4, 5],
+    2: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+}
+_TT_CURRENT_VERSION = 2  # default when there is no AT clause
+
+
+def resolve_tt_version(at_unit: str | None, at_value: str | None) -> int:
+    """Resolve an AT clause to one of this fixture's versions (1 or 2).
+
+    - ``None``         → current version (2)
+    - ``VERSION => n`` → ``int(n)`` (must be 1 or 2)
+    - ``TIMESTAMP``    → year <= 2020 → 1, else 2
+    """
+    if not at_unit:
+        return _TT_CURRENT_VERSION
+    unit = at_unit.upper()
+    if unit == "VERSION":
+        version = int(at_value)  # type: ignore[arg-type]
+        if version not in _TT_VERSION_IDS:
+            raise ValueError(f"Unknown version {version}; valid: {sorted(_TT_VERSION_IDS)}")
+        return version
+    if unit == "TIMESTAMP":
+        year = int(str(at_value)[:4])
+        return 1 if year <= 2020 else 2
+    raise ValueError(f"Unsupported at_unit: {at_unit!r}")
+
+
+@dataclass(kw_only=True)
+class _TtState(ArrowSerializableDataclass):
+    """State for the time-travel + pushdown fixtures.
+
+    ``seen_version`` / ``pushed_filters`` are serialized (NOT transient): the HTTP
+    state-token rehydrate path deserializes state without re-running
+    ``initial_state``, so they must survive that round-trip to echo correctly.
+    """
+
+    seen_version: int = 0
+    pushed_filters: str = "(none)"
+    done: bool = False
+
+
+def _emit_version(params: ProcessParams[Any], state: _TtState, out: OutputCollector) -> None:
+    """Emit one batch for ``state.seen_version``, projected to the output schema."""
+    if state.done:
+        out.finish()
+        return
+    state.done = True
+    ids = _TT_VERSION_IDS[state.seen_version]
+    full: dict[str, list[Any]] = {
+        "id": ids,
+        "val": [i * 10 for i in ids],
+        "seen_version": [state.seen_version] * len(ids),
+        "pushed_filters": [state.pushed_filters] * len(ids),
+    }
+    # projection_pushdown=True: emit only the requested columns.
+    columns = {f.name: full[f.name] for f in params.output_schema}
+    out.emit(pa.RecordBatch.from_pydict(columns, schema=params.output_schema))
+
+
+def _pushed_filters_str(params: ProcessParams[Any]) -> str:
+    assert params.init_call is not None
+    pf = params.init_call.pushdown_filters
+    jk = params.init_call.join_keys
+    filters = TableFunctionGenerator.pushdown_filters(pf, join_keys=jk) if pf is not None else None
+    return _format_pushed_filters(filters)
+
+
+@init_single_worker
+@bind_fixed_schema
+class TimeTravelPushdownFunction(TableFunctionGenerator[_EmptyArgs, _TtState]):
+    """Function-backed time-travel + pushdown scan.
+
+    Reads the AT clause from the **init** request (``params.at_value``) — proving
+    the framework now threads AT onto the bind request embedded in init. No
+    arguments: the version comes from AT, not from a scan-function argument.
+    """
+
+    class Meta:
+        name = "tt_pushdown_scan"
+        description = "Function-backed time-travel + filter-pushdown scan (reads AT at init)."
+        categories = ["generator", "diagnostic", "testing"]
+        filter_pushdown = True
+        auto_apply_filters = True
+        projection_pushdown = True
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = _TT_SCHEMA
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_EmptyArgs]) -> _TtState:
+        version = resolve_tt_version(params.at_unit, params.at_value)
+        return _TtState(seen_version=version, pushed_filters=_pushed_filters_str(params))
+
+    @classmethod
+    def process(cls, params: ProcessParams[_EmptyArgs], state: _TtState, out: OutputCollector) -> None:
+        _emit_version(params, state, out)
+
+
+@dataclass(slots=True, frozen=True)
+class _TtColsArgs:
+    """Argument for the columns-based scan.
+
+    The resolved version (injected by the worker's
+    ``table_scan_function_get`` from the AT clause).
+    """
+
+    version: Annotated[int, Arg(0, doc="Resolved data version")]
+
+
+@init_single_worker
+@bind_fixed_schema
+class TtPushdownColsScanFunction(TableFunctionGenerator[_TtColsArgs, _TtState]):
+    """Columns-based time-travel + pushdown scan.
+
+    Receives the version as a scan-function **argument** (the native columns-based
+    AT mechanism: the worker resolves AT → version in ``table_scan_function_get``).
+    Backs ``example.data.tt_pushdown_cols``.
+    """
+
+    class Meta:
+        name = "tt_pushdown_cols_scan"
+        description = "Columns-based time-travel + filter-pushdown scan (version via arg)."
+        categories = ["generator", "diagnostic", "testing"]
+        filter_pushdown = True
+        auto_apply_filters = True
+        projection_pushdown = True
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = _TT_SCHEMA
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_TtColsArgs]) -> _TtState:
+        return _TtState(seen_version=params.args.version, pushed_filters=_pushed_filters_str(params))
+
+    @classmethod
+    def process(cls, params: ProcessParams[_TtColsArgs], state: _TtState, out: OutputCollector) -> None:
+        _emit_version(params, state, out)

vgi/_test_fixtures/table/versioned.py

@@ -0,0 +1,230 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Versioned data fixtures (versioned_data, versioned_constraints) used by time-travel tests."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Annotated, Any
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+from vgi_rpc.rpc import OutputCollector
+
+from vgi._test_fixtures.table._common import _OneShotState
+from vgi.arguments import Arg
+from vgi.invocation import BindResponse
+from vgi.schema_utils import schema
+from vgi.table_function import (
+    BindParams,
+    ProcessParams,
+    TableFunctionGenerator,
+    init_single_worker,
+)
+
+# ============================================================================
+
+# Version definitions: schema and data per version
+_VERSIONED_SCHEMAS: dict[int, pa.Schema] = {
+    1: schema(id=pa.int64()),
+    2: schema(id=pa.int64(), name=pa.string(), score=pa.float64(), active=pa.bool_()),
+    3: schema(id=pa.int64(), score=pa.float64()),
+}
+
+_VERSIONED_DATA: dict[int, dict[str, list[Any]]] = {
+    1: {"id": [1, 2, 3]},
+    2: {
+        "id": [1, 2, 3, 4, 5],
+        "name": ["alice", "bob", "carol", "dave", "eve"],
+        "score": [10.0, 20.0, 30.0, 40.0, 50.0],
+        "active": [True, False, True, False, True],
+    },
+    3: {"id": [1, 2, 3, 4], "score": [15.0, 25.0, 35.0, 45.0]},
+}
+
+# Current version (default when no AT clause)
+_CURRENT_VERSION = 3
+
+
+def resolve_version(at_unit: str | None, at_value: str | None) -> int:
+    """Resolve AT clause to a version number.
+
+    - ``VERSION``: direct integer version (must exist in ``_VERSIONED_SCHEMAS``)
+    - ``TIMESTAMP``: year-based mapping (<=2020→1, <=2021→2, >=2022→3)
+    - ``None``: current version (3)
+
+    Raises ``ValueError`` for unknown versions or unsupported AT units.
+    """
+    if not at_unit:
+        return _CURRENT_VERSION
+
+    if at_unit.upper() == "VERSION":
+        version = int(at_value)  # type: ignore[arg-type]
+        if version not in _VERSIONED_SCHEMAS:
+            raise ValueError(f"Unknown version: {version}. Valid versions: {sorted(_VERSIONED_SCHEMAS)}")
+        return version
+
+    if at_unit.upper() == "TIMESTAMP":
+        # Parse year from timestamp string (e.g. "2020-06-15 00:00:00")
+        year = int(str(at_value)[:4])
+        if year < 2020:
+            raise ValueError(f"No version exists at timestamp {at_value!r}: table did not exist before 2020")
+        if year <= 2020:
+            return 1
+        if year <= 2021:
+            return 2
+        return 3
+
+    raise ValueError(f"Unsupported at_unit: {at_unit!r}")
+
+
+@dataclass(slots=True, frozen=True)
+class VersionedDataFunctionArgs:
+    """Arguments for VersionedDataFunction."""
+
+    version: Annotated[int, Arg(0, doc="Data version to return")]
+
+
+@dataclass(kw_only=True)
+class VersionedDataState(ArrowSerializableDataclass):
+    """State for VersionedDataFunction."""
+
+    done: bool = False
+
+
+@init_single_worker
+class VersionedDataFunction(TableFunctionGenerator[VersionedDataFunctionArgs, VersionedDataState]):
+    """Returns version-specific data demonstrating time travel with schema evolution.
+
+    Each version has a different schema and different data:
+
+    - **Version 1**: ``(id int64)`` — 3 rows
+    - **Version 2**: ``(id int64, name string, score double, active bool)`` — 5 rows
+    - **Version 3** (current): ``(id int64, score double)`` — 4 rows
+
+    """
+
+    class Meta:
+        """Metadata for VersionedDataFunction."""
+
+        name = "versioned_data_scan"
+        description = "Returns versioned data with schema evolution"
+        categories = ["generator", "testing"]
+
+    @classmethod
+    def on_bind(cls, params: BindParams[VersionedDataFunctionArgs]) -> BindResponse:
+        """Return version-specific output schema."""
+        version = params.args.version
+        if version not in _VERSIONED_SCHEMAS:
+            raise ValueError(f"Unknown version: {version}. Valid versions: {sorted(_VERSIONED_SCHEMAS)}")
+        return BindResponse(output_schema=_VERSIONED_SCHEMAS[version])
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[VersionedDataFunctionArgs]) -> VersionedDataState:
+        """Create initial state."""
+        return VersionedDataState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[VersionedDataFunctionArgs],
+        state: VersionedDataState,
+        out: OutputCollector,
+    ) -> None:
+        """Emit all rows for the requested version in one batch."""
+        if state.done:
+            out.finish()
+            return
+        state.done = True
+        version = params.args.version
+        data = _VERSIONED_DATA[version]
+        out.emit(pa.RecordBatch.from_pydict(data, schema=params.output_schema))
+
+
+# ============================================================================
+
+# Version 1: simple users table (id, name) — NOT NULL on id only
+# Version 2: adds email column, PK on id, UNIQUE on email
+# Version 3: adds department_id column, FK to departments
+
+_VERSIONED_CONSTRAINTS_SCHEMAS: dict[int, pa.Schema] = {
+    1: schema(id=pa.int64(), name=pa.string()),
+    2: schema(id=pa.int64(), name=pa.string(), email=pa.string()),
+    3: schema(id=pa.int64(), name=pa.string(), email=pa.string(), department_id=pa.int64()),
+}
+
+_VERSIONED_CONSTRAINTS_DATA: dict[int, dict[str, list[Any]]] = {
+    1: {"id": [1, 2], "name": ["Alice", "Bob"]},
+    2: {"id": [1, 2, 3], "name": ["Alice", "Bob", "Carol"], "email": ["a@co", "b@co", "c@co"]},
+    3: {
+        "id": [1, 2, 3],
+        "name": ["Alice", "Bob", "Carol"],
+        "email": ["a@co", "b@co", "c@co"],
+        "department_id": [1, 2, 1],
+    },
+}
+
+_VERSIONED_CONSTRAINTS_CURRENT = 3
+
+
+def resolve_versioned_constraints_version(at_unit: str | None, at_value: str | None) -> int:
+    """Resolve AT clause for versioned_constraints table."""
+    if not at_unit:
+        return _VERSIONED_CONSTRAINTS_CURRENT
+
+    if at_unit.upper() == "VERSION":
+        version = int(at_value)  # type: ignore[arg-type]
+        if version not in _VERSIONED_CONSTRAINTS_SCHEMAS:
+            raise ValueError(f"Unknown version: {version}. Valid versions: {sorted(_VERSIONED_CONSTRAINTS_SCHEMAS)}")
+        return version
+
+    raise ValueError(f"Unsupported at_unit: {at_unit!r}")
+
+
+@dataclass(slots=True, frozen=True)
+class _VersionedConstraintsArgs:
+    """Arguments for VersionedConstraintsScanFunction."""
+
+    version: Annotated[int, Arg(0, doc="Data version")]
+
+
+@init_single_worker
+class VersionedConstraintsScanFunction(
+    TableFunctionGenerator[_VersionedConstraintsArgs, _OneShotState],
+):
+    """Returns version-specific data for constraint evolution testing."""
+
+    class Meta:
+        """Metadata for VersionedConstraintsScanFunction."""
+
+        name = "versioned_constraints_scan"
+        description = "Scan versioned constraints table"
+
+    @classmethod
+    def on_bind(cls, params: BindParams[_VersionedConstraintsArgs]) -> BindResponse:
+        """Return output schema."""
+        version = params.args.version
+        if version not in _VERSIONED_CONSTRAINTS_SCHEMAS:
+            raise ValueError(f"Unknown version: {version}")
+        return BindResponse(output_schema=_VERSIONED_CONSTRAINTS_SCHEMAS[version])
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_VersionedConstraintsArgs]) -> _OneShotState:
+        """Create initial state."""
+        return _OneShotState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_VersionedConstraintsArgs],
+        state: _OneShotState,
+        out: OutputCollector,
+    ) -> None:
+        """Emit data."""
+        if state.done:
+            out.finish()
+            return
+        state.done = True
+        version = params.args.version
+        data = _VERSIONED_CONSTRAINTS_DATA[version]
+        out.emit(pa.RecordBatch.from_pydict(data, schema=params.output_schema))