vgi-python 0.36.2__tar.gz → 0.37.1__tar.gz

{vgi_python-0.36.2 → vgi_python-0.37.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: vgi-python
-Version: 0.36.2
+Version: 0.37.1
 Summary: Vector Gateway Interface - Connect DuckDB to external programs via Apache Arrow
 Project-URL: Homepage, https://query.farm
 Project-URL: Repository, https://github.com/Query-farm/vgi-python

{vgi_python-0.36.2 → vgi_python-0.37.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vgi-python"
-version = "0.36.2"
+version = "0.37.1"
 description = "Vector Gateway Interface - Connect DuckDB to external programs via Apache Arrow"
 readme = "README.md"
 keywords = [

{vgi_python-0.36.2 → vgi_python-0.37.1}/tests/test_serve.py

@@ -24,7 +24,7 @@ from vgi_rpc import (
 from vgi_rpc.rpc import AuthContext, PeerResolutionContext, peer_identity_primary
 
 from vgi.arguments import Arguments, Auth, Param, Returns
-from vgi.auth import TokenIdentity
+from vgi.auth import IssuedGrant, TokenIdentity
 from vgi.client import Client
 from vgi.scalar_function import ScalarFunction
 from vgi.serve import (
@@ -486,6 +486,32 @@ class _IntrospectingWorker(Worker):
         return None
 
 
+class _MintingWorker(Worker):
+    """Worker that implements only the optional grant-minting hook."""
+
+    functions = [_DoubleFunc]
+
+    @classmethod
+    def mint_grant(cls, principal: str, purpose: str, scopes: list[str], ttl_seconds: int) -> IssuedGrant:
+        # A sealed self-contained credential: the worker owns the format and
+        # the framework never parses it. Capped TTL, to pin that a worker may
+        # issue SHORTER than asked.
+        ttl = min(ttl_seconds, 3600)
+        return IssuedGrant(
+            token=f"sealed:{principal}:{':'.join(scopes)}:{purpose}",
+            expires_at=1_000_000.0 + ttl,
+            grant_id="grant-1",
+        )
+
+
+class _BothWorker(_IntrospectingWorker):
+    """Worker that both resolves credentials and mints grants."""
+
+    @classmethod
+    def mint_grant(cls, principal: str, purpose: str, scopes: list[str], ttl_seconds: int) -> IssuedGrant:
+        return IssuedGrant(token=f"sealed:{principal}", expires_at=1_000_000.0, grant_id="g")
+
+
 class TestIntrospectResolverDetection:
     """``resolve_token`` presence is detected by override, not by a flag."""
 
@@ -516,6 +542,31 @@ def _introspector_ctx(principal: str) -> object:
     )
 
 
+def _minting_ctx(principal: str, *, auth_age_s: float = 0.0) -> object:
+    """Build a ``CallContext`` for a caller who authenticated *auth_age_s* ago.
+
+    Minting reads ``auth_time`` off the credential's claims, and a credential
+    without one cannot mint at all. That single rule is what stops a grant
+    minting another grant — a grant is not IdP-issued, so it carries no
+    ``auth_time`` — and what makes a static bearer token, which proves a
+    machine holds a secret rather than that a human just authenticated,
+    fail closed here.
+    """
+    import time
+
+    from vgi_rpc.rpc import AuthContext, CallContext
+
+    return CallContext(
+        auth=AuthContext(
+            domain="test",
+            authenticated=True,
+            principal=principal,
+            claims={"auth_time": time.time() - auth_age_s},
+        ),
+        emit_client_log=lambda *a, **k: None,
+    )
+
+
 def _hosted_protocols(worker_cls: type[Worker]) -> frozenset[str]:
     """Return the protocol names the server built for *worker_cls* hosts.
 
@@ -655,6 +706,111 @@ class TestTokenIntrospection:
             create_app(_IntrospectingWorker, prefix="/vgi", describe=False)
 
 
+class TestGrantMinting:
+    """``issue_grant`` — the other half of ``vgi_rpc.Identity.v1``.
+
+    Guarded very differently from ``introspect_token``, and deliberately so.
+    Introspection answers a question about *somebody else's* credential, so it
+    needs an allowlist. Minting produces a credential for the *caller* — there
+    is no subject on the wire, so cross-subject minting is closed by
+    construction and no allowlist applies. What this repo owns is the wiring:
+    a ``mint_grant`` override becomes an ``IdentityImpl`` that offers the
+    method, and a worker that never wrote one does not host it.
+    """
+
+    @staticmethod
+    def _identity(worker_cls: type[Worker], **env: str) -> object:
+        from vgi.serve import _build_identity
+
+        with pytest.MonkeyPatch.context() as mp:
+            mp.delenv("VGI_INTROSPECT_PRINCIPALS", raising=False)
+            for key, value in env.items():
+                mp.setenv(key, value)
+            return _build_identity(worker_cls, None)
+
+    def test_base_worker_has_no_minter(self) -> None:
+        """A worker that never wrote the minting gets no issuer."""
+        assert _SingleWorker._grant_minter() is None
+
+    def test_override_is_detected(self) -> None:
+        """Implementing the hook is what enables it."""
+        minter = _MintingWorker._grant_minter()
+        assert minter is not None
+
+    def test_minting_alone_hosts_the_protocol(self) -> None:
+        """A worker may mint without resolving — the two are independent."""
+        assert "vgi_rpc.Identity.v1" in _hosted_protocols(_MintingWorker)
+
+    def test_minting_alone_offers_only_issue_grant(self) -> None:
+        """The protocol describes what the worker does, not what it could do."""
+        identity = self._identity(_MintingWorker)
+        assert identity is not None
+        assert identity.offered_methods() == frozenset({"issue_grant"})  # type: ignore[attr-defined]
+
+    def test_minting_alone_needs_no_introspect_allowlist(self) -> None:
+        """No allowlist is required to mint: a grant is not an oracle.
+
+        Requiring one would be a startup error for a worker that never exposes
+        anybody else's identity — the exact opposite of the guard's purpose.
+        """
+        with pytest.MonkeyPatch.context() as mp:
+            mp.delenv("VGI_INTROSPECT_PRINCIPALS", raising=False)
+            create_app(_MintingWorker, prefix="/vgi", describe=False)
+
+    def test_both_hooks_offer_both_methods(self) -> None:
+        """Resolving and minting compose."""
+        identity = self._identity(_BothWorker, VGI_INTROSPECT_PRINCIPALS="proxy-a")
+        assert identity is not None
+        assert identity.offered_methods() == frozenset({"introspect_token", "issue_grant"})  # type: ignore[attr-defined]
+
+    def test_the_caller_is_the_subject(self) -> None:
+        """The framework supplies the principal; the wire method has no subject field.
+
+        This is the property that makes an allowlist unnecessary, so pin it
+        here rather than trusting the signature to stay that way.
+        """
+        identity = self._identity(_MintingWorker)
+        assert identity is not None
+        grant = identity.issue_grant(  # type: ignore[attr-defined]
+            "nightly-etl", ["read"], 600, _minting_ctx("alice")
+        )
+        assert grant.token.startswith("sealed:alice:")
+        assert grant.grant_id == "grant-1"
+
+    def test_a_worker_may_issue_a_shorter_ttl_than_asked(self) -> None:
+        """A worker may narrow the request; it must never widen it."""
+        identity = self._identity(_MintingWorker)
+        assert identity is not None
+        grant = identity.issue_grant(  # type: ignore[attr-defined]
+            "long", [], 86_400, _minting_ctx("alice")
+        )
+        assert grant.expires_at == 1_000_000.0 + 3600
+
+    def test_stale_authentication_cannot_mint(self) -> None:
+        """Minting while the user is present is enforced, not just documented."""
+        from vgi_rpc.rpc._token_identity import StaleAuthError
+
+        identity = self._identity(_MintingWorker)
+        assert identity is not None
+        with pytest.raises(StaleAuthError):
+            identity.issue_grant("x", [], 60, _minting_ctx("alice", auth_age_s=86_400))  # type: ignore[attr-defined]
+
+    def test_a_credential_with_no_auth_time_cannot_mint(self) -> None:
+        """A grant cannot mint another grant, and a static bearer cannot mint at all.
+
+        Neither carries ``auth_time`` — only an IdP-issued credential does — so
+        the lineage cannot escape the identity provider. This is the rule that
+        makes delegation accountable to a human login rather than to whoever
+        currently holds a machine secret.
+        """
+        from vgi_rpc.rpc._token_identity import StaleAuthError
+
+        identity = self._identity(_MintingWorker)
+        assert identity is not None
+        with pytest.raises(StaleAuthError):
+            identity.issue_grant("x", [], 60, _introspector_ctx("etl-service"))  # type: ignore[attr-defined]
+
+
 class TestAuthUnavailableReExport:
     """``AuthUnavailableError`` is reachable without naming a vgi_rpc module."""
 

{vgi_python-0.36.2 → vgi_python-0.37.1}/uv.lock

@@ -2580,7 +2580,7 @@ requires-dist = [
 
 [[package]]
 name = "vgi-python"
-version = "0.36.2"
+version = "0.37.1"
 source = { editable = "." }
 dependencies = [
     { name = "click" },

vgi_python-0.37.1/vgi/_test_fixtures/secret_cache.py

@@ -0,0 +1,237 @@
+# Copyright 2026 Query Farm LLC - https://query.farm
+
+"""Cacheable fixtures whose results depend on a secret.
+
+The C++ result cache keys a secret-dependent result on a fingerprint of the
+secrets its bind resolved (never their values), so a result is reused while the
+secret is unchanged and recomputed the moment it is rotated, re-scoped or
+dropped. Each fixture here reads the ``vgi_example`` secret's ``secret_string``
+and advertises cacheability, one per cache path the fingerprint has to reach:
+
+* ``secret_cache_nonce()`` — producer table function; the secret is declared in
+  ``Meta.required_secrets``. Also exposed as the ``data.secret_cache_nonce``
+  table with ``inline_bind=True``, which covers the client's inline-bind path
+  (no bind RPC; the secrets are resolved client-side for init).
+* ``secret_cached_scalar(x)`` — scalar, per-value memoized; secret declared via
+  a ``Secret()`` annotation.
+* ``secret_cached_lateral(x)`` — blended map, per-value memoized, called under
+  ``LATERAL``; the secret is requested in ``on_bind`` (the two-phase bind).
+
+Every output carries a ``nonce`` minted only when the worker really runs. It is
+random rather than a counter because a pooled worker may run several processes,
+and a per-process counter can repeat across them: equal nonces prove a cache
+HIT, different ones a MISS, on any pool size.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Annotated, Any, ClassVar, cast
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+from vgi_rpc.rpc import OutputCollector
+
+from vgi.arguments import Arg, OutputLength, Param, Returns, Secret, SecretLookupEntry
+from vgi.cache_control import CacheControl
+from vgi.invocation import BindResponse
+from vgi.metadata import FunctionExample
+from vgi.scalar_function import ScalarFunction
+from vgi.schema_utils import schema
+from vgi.table_function import (
+    BindParams,
+    ProcessParams,
+    TableFunctionGenerator,
+    bind_fixed_schema,
+    init_single_worker,
+)
+from vgi.table_in_out_function import RowTransformFunction
+
+if TYPE_CHECKING:
+    from vgi.protocol import VgiOutputCollector
+
+#: The secret type every fixture here reads. Registered by the C++ extension's
+#: test build (``CREATE SECRET ... (TYPE vgi_example, secret_string '...')``).
+SECRET_TYPE = "vgi_example"
+
+#: Long enough that TTL never lapses mid-test.
+_TTL_SECONDS = 300
+
+
+def _nonce() -> int:
+    """A value unique to this invocation, across every process in a worker pool."""
+    return int.from_bytes(os.urandom(7), "big")
+
+
+def _secret_string(secret: dict[str, Any]) -> str | None:
+    """The ``secret_string`` field of a resolved secret, or None when there is none."""
+    value = secret.get("secret_string")
+    if value is None:
+        return None
+    return str(value.as_py()) if isinstance(value, pa.Scalar) else str(value)
+
+
+# ---------------------------------------------------------------------------
+# secret_cache_nonce — producer
+# ---------------------------------------------------------------------------
+@dataclass(slots=True, frozen=True)
+class SecretCacheNonceArgs:
+    """Arguments for SecretCacheNonceFunction (none)."""
+
+
+@dataclass(kw_only=True)
+class _SecretCacheNonceState(ArrowSerializableDataclass):
+    """The one row to emit, minted on a real invocation."""
+
+    secret_string: str | None
+    nonce: int
+    done: bool = False
+
+
+@init_single_worker
+@bind_fixed_schema
+class SecretCacheNonceFunction(TableFunctionGenerator[SecretCacheNonceArgs, _SecretCacheNonceState]):
+    """One row: the secret's ``secret_string`` and a per-invocation nonce; cacheable.
+
+    ``initial_state`` runs only on a cache MISS, so the nonce is stable across
+    HITs. A rotated secret must MISS and report the new value; restoring the
+    original secret must HIT the entry it produced.
+    """
+
+    class Meta:
+        """Metadata for SecretCacheNonceFunction."""
+
+        name = "secret_cache_nonce"
+        description = "One row with a secret's value and a per-invocation nonce; cacheable per secret"
+        categories = ["generator", "cache", "secret", "testing"]
+        required_secrets = [SecretLookupEntry(secret_type=SECRET_TYPE)]
+        examples = [
+            FunctionExample(
+                sql="SELECT * FROM secret_cache_nonce()",
+                description="The nonce is stable while the vgi_example secret is unchanged",
+            ),
+        ]
+
+    FunctionArguments = SecretCacheNonceArgs
+    FIXED_SCHEMA: ClassVar[pa.Schema] = schema(secret_string=pa.string(), nonce=pa.int64())
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[SecretCacheNonceArgs]) -> _SecretCacheNonceState:
+        """Read the secret and mint a nonce for this (real) invocation."""
+        secret = next(iter(params.secrets.of_type(SECRET_TYPE)), {})
+        return _SecretCacheNonceState(secret_string=_secret_string(secret), nonce=_nonce())
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[SecretCacheNonceArgs],
+        state: _SecretCacheNonceState,
+        out: OutputCollector,
+    ) -> None:
+        """Emit the single row once, advertising a cache TTL."""
+        if state.done:
+            out.finish()
+            return
+        batch = pa.RecordBatch.from_pydict(
+            {"secret_string": [state.secret_string], "nonce": [state.nonce]},
+            schema=params.output_schema,
+        )
+        cast("VgiOutputCollector", out).emit(batch, cache_control=CacheControl(ttl=_TTL_SECONDS))
+        state.done = True
+
+
+# ---------------------------------------------------------------------------
+# secret_cached_scalar — scalar, per-value
+# ---------------------------------------------------------------------------
+class SecretCachedScalarFunction(ScalarFunction):
+    """``x`` -> ``'<secret_string>|<nonce>'``, memoized per value per secret.
+
+    With no secret resolved the label is ``'|<nonce>'``.
+
+    One nonce per ``compute`` call, shared by the batch, so a served value keeps
+    the nonce of the call that produced it. ``per_value`` is a test choice, as on
+    ``cached_double_scalar``: the point is coverage of the tier, not economics.
+    """
+
+    CACHE_CONTROL = CacheControl(ttl=_TTL_SECONDS, per_value=True)
+
+    class Meta:
+        """Function metadata."""
+
+        name = "secret_cached_scalar"
+        description = "Returns '<secret_string>|<nonce>' per value; memoized per value per secret"
+        examples = [
+            FunctionExample(
+                sql="SELECT secret_cached_scalar(1)",
+                description="Stable while the vgi_example secret is unchanged",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        value: Annotated[pa.Int64Array, Param(doc="Any value; the output ignores it")],
+        _length: Annotated[int, OutputLength()],
+        vgi_example: Annotated[dict[str, pa.Scalar[Any]] | None, Secret(SECRET_TYPE)] = None,
+    ) -> Annotated[pa.StringArray, Returns(pa.string())]:
+        """Label every row with the secret's value ('' when none) and this call's nonce.
+
+        The framework omits a ``Secret()`` argument when no such secret exists,
+        hence the default: a dropped secret is a state this fixture must serve.
+        """
+        label = f"{_secret_string(vgi_example or {}) or ''}|{_nonce()}"
+        return pa.array([label] * _length, type=pa.string())
+
+
+# ---------------------------------------------------------------------------
+# secret_cached_lateral — blended map, per-value, two-phase secret
+# ---------------------------------------------------------------------------
+@dataclass(slots=True, frozen=True, kw_only=True)
+class _SecretCachedLateralArgs:
+    """One positional input column; the output ignores its value."""
+
+    x: Annotated[int, Arg(0, doc="Input column")]
+
+
+class SecretCachedLateralFunction(RowTransformFunction[_SecretCachedLateralArgs]):
+    """1->1 map emitting the secret's ``secret_string`` and a per-call nonce.
+
+    Requests the secret from ``on_bind`` — the two-phase bind, so the secret is
+    discovered at bind time rather than declared — and advertises ``per_value``
+    so a correlated ``LATERAL`` call is memoized per input value per secret.
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "secret_cached_lateral"
+        description = "Blended map emitting a secret's value and a per-call nonce; memoized per secret"
+        categories = ["blended", "cache", "secret", "test"]
+
+    @classmethod
+    def on_bind(cls, params: BindParams[_SecretCachedLateralArgs]) -> BindResponse:
+        """Request the secret (two-phase) and declare the output columns."""
+        params.secrets.get(SECRET_TYPE)
+        return BindResponse(output_schema=schema(secret_string=pa.string(), nonce=pa.int64()))
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_SecretCachedLateralArgs],
+        state: None,
+        batch: pa.RecordBatch,
+        out: OutputCollector,
+    ) -> None:
+        """Emit one row per input row, all carrying this call's nonce."""
+        secret = next(iter(params.secrets.of_type(SECRET_TYPE)), {})
+        rows = batch.num_rows
+        cast("VgiOutputCollector", out).emit(
+            pa.record_batch(
+                {
+                    "secret_string": pa.array([_secret_string(secret)] * rows, type=pa.string()),
+                    "nonce": pa.array([_nonce()] * rows, type=pa.int64()),
+                }
+            ),
+            cache_control=CacheControl(ttl=_TTL_SECONDS, per_value=True),
+        )

{vgi_python-0.36.2 → vgi_python-0.37.1}/vgi/_test_fixtures/table/__init__.py

@@ -113,6 +113,7 @@ from vgi._test_fixtures.table.partition_columns import (
     OverlappingRangePartitionedFunction,
     PartitionedWithExplicitOverrideFunction,
     RegionYearPartitionedFunction,
+    TrailingPartitionSalesFunction,
 )
 from vgi._test_fixtures.table.partition_columns_broken import (
     BrokenMissingPartitionValuesFunction,
@@ -264,6 +265,7 @@ __all__ = [
     "RFF_SIMPLE_COLUMNS",
     "RFF_STRUCT_COLUMNS",
     "RegionYearPartitionedFunction",
+    "TrailingPartitionSalesFunction",
     "RepeatValueIntFunction",
     "RepeatValueStrFunction",
     "RffMultiScanFunction",

{vgi_python-0.36.2 → vgi_python-0.37.1}/vgi/_test_fixtures/table/partition_columns.py

@@ -23,6 +23,11 @@ Fixtures:
 * :class:`RegionYearPartitionedFunction` — multi-column SINGLE_VALUE.
   Each chunk has a single ``(region, year)`` tuple.
 
+* :class:`TrailingPartitionSalesFunction` — same contract as the
+  country fixture, but the partition column is declared LAST rather
+  than first, so the worker-schema and scan-local index spaces do not
+  coincide. Also exposed as a catalog table.
+
 * :class:`PartitionedWithProjectedOutColumnFunction` — declares
   partition on ``category`` but DOES NOT include ``category`` in the
   emitted batch. Uses the explicit ``partition_values=`` override on
@@ -49,7 +54,7 @@ from __future__ import annotations
 
 import struct
 from dataclasses import dataclass
-from typing import Annotated, ClassVar, cast
+from typing import Annotated, Any, ClassVar, cast
 
 import pyarrow as pa
 from vgi_rpc import ArrowSerializableDataclass
@@ -285,6 +290,126 @@ class RegionYearPartitionedFunction(TableFunctionGenerator[_RegionYearArgs, _Reg
         state.current_idx = rpp
 
 
+# =============================================================================
+# SINGLE_VALUE_PARTITIONS with the partition column NOT first
+# =============================================================================
+
+
+@dataclass(slots=True, frozen=True)
+class _TrailingPartitionArgs:
+    """Arguments for ``trailing_partition_sales``."""
+
+    rows_per_country: Annotated[int, Arg(0, doc="Rows to emit per country partition", ge=1)]
+
+
+@dataclass(kw_only=True)
+class _TrailingPartitionState(ArrowSerializableDataclass):
+    """Per-worker cursor over countries (same shape as the country fixture)."""
+
+    current_country: str | None = None
+    current_country_idx: int = -1
+    current_idx: int = 0
+
+
+@bind_fixed_schema
+@_cardinality_from_count
+class TrailingPartitionSalesFunction(TableFunctionGenerator[_TrailingPartitionArgs, _TrailingPartitionState]):
+    """``country`` is single-valued per chunk, but sits LAST in the schema.
+
+    Identical contract to :class:`CountryPartitionedSalesFunction`; the only
+    difference is where the partition column lives. That difference is the
+    point: every other partitioned fixture declares its partition column at
+    index 0, which makes two distinct index spaces accidentally agree.
+
+    ``CanUsePartitionedAggregate`` asks ``get_partition_info`` about
+    WORKER-SCHEMA indices, but the sink later asks ``get_partition_data``
+    about SCAN-LOCAL ones (positions in the scan's own ``column_ids``, after
+    projection pushdown). ``GROUP BY country`` projects just ``country`` and
+    ``sales``, so here the sink asks about scan-local 0 while the declared
+    index is 3 — and a client that compares them without mapping raises
+    "sink requested partition column 0 that is not in the declared partition
+    set", which is fatal. With the partition column at index 0 both spaces
+    say 0 and the bug is invisible.
+
+    Also registered as a catalog TABLE (``example.data.trailing_partition_sales``)
+    because the two paths install their scan functions separately: a client can
+    wire ``get_partition_info`` for direct function calls and miss it for
+    catalog tables, in which case the planner silently never picks
+    ``PARTITIONED_AGGREGATE`` for a table however it is declared.
+    """
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema(
+        cast(
+            "list[pa.Field[Any]]",
+            [
+                pa.field("seq", pa.int64()),
+                pa.field("label", pa.string()),
+                pa.field("sales", pa.int64()),
+                partition_field("country", pa.string()),
+            ],
+        )
+    )
+
+    class Meta:
+        name = "trailing_partition_sales"
+        projection_pushdown = True
+        description = (
+            "Per-country sales rows, one Arrow batch per country, with the "
+            "SINGLE_VALUE partition column declared LAST in the schema "
+            "instead of first."
+        )
+        categories = ["generator", "partitioning"]
+        partition_kind = PartitionKind.SINGLE_VALUE_PARTITIONS
+        examples = [
+            FunctionExample(
+                sql="SELECT country, SUM(sales) FROM trailing_partition_sales(100) GROUP BY country",
+                description="Partitioned aggregate over a non-leading partition column",
+            ),
+        ]
+
+    @classmethod
+    def on_init(cls, params: InitParams[_TrailingPartitionArgs]) -> GlobalInitResponse:
+        items = [struct.pack(_QUEUE_ITEM_FMT, i) for i in range(len(_COUNTRIES))]
+        params.storage.queue_push(items)
+        return GlobalInitResponse()
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_TrailingPartitionArgs]) -> _TrailingPartitionState:
+        return _TrailingPartitionState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_TrailingPartitionArgs],
+        state: _TrailingPartitionState,
+        out: OutputCollector,
+    ) -> None:
+        if state.current_country is None or state.current_idx >= params.args.rows_per_country:
+            item = params.storage.queue_pop()
+            if item is None:
+                out.finish()
+                return
+            (state.current_country_idx,) = struct.unpack(_QUEUE_ITEM_FMT, item)
+            state.current_country = _COUNTRIES[state.current_country_idx]
+            state.current_idx = 0
+
+        rpc = params.args.rows_per_country
+        # Same deterministic sales values as country_partitioned_sales, so a
+        # test can assert the two agree column-for-column despite the layout.
+        base = state.current_country_idx * 1_000_000
+        batch = pa.RecordBatch.from_pydict(
+            {
+                "seq": list(range(rpc)),
+                "label": [f"{state.current_country}-{i}" for i in range(rpc)],
+                "sales": [base + i for i in range(rpc)],
+                "country": [state.current_country] * rpc,
+            },
+            schema=cls.FIXED_SCHEMA,
+        )
+        out.emit(batch)
+        state.current_idx = rpc
+
+
 # =============================================================================
 # Projected-out partition column — exercises explicit override path
 # =============================================================================