vgi-python 0.8.0__py3-none-any.whl

vgi/_test_fixtures/table/__init__.py

@@ -0,0 +1,221 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Table-function fixtures.
+
+Originally a single 3,270-line module; split into cohesive sub-modules and
+re-exported here so existing import sites (worker.py, tests) keep working
+unchanged.
+
+If you're looking for a specific fixture, the module names below should
+point you at the right file:
+
+* :mod:`._common`         — ``CountdownState``, ``_BaseSequenceFunction``
+* :mod:`.sequence`        — sequence / partitioned / nested / row_id
+* :mod:`.make_series`     — make_series_count / range / step / csv / float
+* :mod:`.pairs`           — make_pairs_*, repeat_value_*, constant_columns
+* :mod:`.settings`        — settings_aware, struct_settings, secret_demo
+* :mod:`.filters`         — filter_echo, dynamic_filter_echo, expression_filter,
+                            spatial_filter
+* :mod:`.catalog_scans`   — colors / departments / employees / products / projects
+* :mod:`.versioned`       — versioned_data + versioned_constraints (time travel)
+* :mod:`.misc`            — projected_data, generator_exception,
+                            logging_generator, order_echo, sample_echo
+"""
+
+from vgi._test_fixtures.table.batch_index import (
+    PartitionedBatchIndexFunction,
+    PartitionedBatchIndexMarkedFunction,
+)
+from vgi._test_fixtures.table.batch_index_broken import (
+    BatchIndexOverflowFunction,
+    MissingBatchIndexTagFunction,
+    NonMonotoneBatchIndexFunction,
+)
+from vgi._test_fixtures.table.catalog_scans import (
+    ColorsScanFunction,
+    DepartmentsScanFunction,
+    EmployeesScanFunction,
+    ProductsScanFunction,
+    ProjectsScanFunction,
+)
+from vgi._test_fixtures.table.filters import (
+    DictFilterEchoFunction,
+    DynamicFilterEchoFunction,
+    ExpressionFilterTestFunction,
+    FilterEchoFunction,
+    FilterEchoPartitionedFunction,
+    FilterEchoTableScanFunction,
+    SpatialFilterExampleFunction,
+    ValuePruneFunction,
+)
+from vgi._test_fixtures.table.late_materialization import (
+    LateMaterializationFunction,
+)
+from vgi._test_fixtures.table.make_series import (
+    MakeSeriesCountFunction,
+    MakeSeriesCsvFunction,
+    MakeSeriesFloatFunction,
+    MakeSeriesRangeFunction,
+    MakeSeriesStepFunction,
+)
+from vgi._test_fixtures.table.misc import (
+    GeneratorExceptionFunction,
+    LoggingGeneratorFunction,
+    OrderEchoFunction,
+    ProjectedDataFunction,
+    SampleEchoFunction,
+)
+from vgi._test_fixtures.table.order_modes import (
+    PartitionedFixedOrderFunction,
+    PartitionedNoOrderGuaranteeFunction,
+    PartitionedPreservesOrderFunction,
+)
+from vgi._test_fixtures.table.pairs import (
+    ConstantColumnsFunction,
+    MakePairsIntFunction,
+    MakePairsIntStrFunction,
+    MakePairsStrFunction,
+    RepeatValueIntFunction,
+    RepeatValueStrFunction,
+)
+from vgi._test_fixtures.table.partition_columns import (
+    CountryPartitionedSalesFunction,
+    DisjointRangePartitionedFunction,
+    PartitionedWithExplicitOverrideFunction,
+    RegionYearPartitionedFunction,
+)
+from vgi._test_fixtures.table.partition_columns_broken import (
+    BrokenMissingPartitionValuesFunction,
+    BrokenPartitionColumnAbsentFromBatchFunction,
+    BrokenPartitionMinNeqMaxFunction,
+    BrokenPartitionValuesNoAnnotationFunction,
+)
+from vgi._test_fixtures.table.profiling_example import (
+    ProfilingDemoFunction,
+)
+from vgi._test_fixtures.table.required_filters import (
+    RFF_MULTI_COLUMNS,
+    RFF_NESTED_COLUMNS,
+    RFF_NONE_COLUMNS,
+    RFF_ROWID_COLUMNS,
+    RFF_SIMPLE_COLUMNS,
+    RFF_STRUCT_COLUMNS,
+    RffMultiScanFunction,
+    RffNestedScanFunction,
+    RffNoneScanFunction,
+    RffRowidScanFunction,
+    RffSimpleScanFunction,
+    RffStructScanFunction,
+)
+from vgi._test_fixtures.table.sequence import (
+    DoubleSequenceFunction,
+    NamedParamsEchoFunction,
+    NestedSequenceFunction,
+    PartitionedSequenceFunction,
+    RowIdSequenceFunction,
+    SequenceFunction,
+    TenThousandFunction,
+)
+from vgi._test_fixtures.table.settings import (
+    ScopedSecretDemoFunction,
+    SecretDemoFunction,
+    SettingsAwareFunction,
+    StructSettingsFunction,
+)
+from vgi._test_fixtures.table.transaction_storage import TxCachedValueFunction
+from vgi._test_fixtures.table.versioned import (
+    _CURRENT_VERSION,
+    _VERSIONED_CONSTRAINTS_CURRENT,
+    _VERSIONED_CONSTRAINTS_DATA,
+    _VERSIONED_CONSTRAINTS_SCHEMAS,
+    _VERSIONED_DATA,
+    _VERSIONED_SCHEMAS,
+    VersionedConstraintsScanFunction,
+    VersionedDataFunction,
+    resolve_version,
+    resolve_versioned_constraints_version,
+)
+
+__all__ = [
+    "_CURRENT_VERSION",
+    "_VERSIONED_CONSTRAINTS_CURRENT",
+    "_VERSIONED_CONSTRAINTS_DATA",
+    "_VERSIONED_CONSTRAINTS_SCHEMAS",
+    "_VERSIONED_DATA",
+    "_VERSIONED_SCHEMAS",
+    "BatchIndexOverflowFunction",
+    "BrokenMissingPartitionValuesFunction",
+    "BrokenPartitionColumnAbsentFromBatchFunction",
+    "BrokenPartitionMinNeqMaxFunction",
+    "BrokenPartitionValuesNoAnnotationFunction",
+    "ColorsScanFunction",
+    "ConstantColumnsFunction",
+    "CountryPartitionedSalesFunction",
+    "DisjointRangePartitionedFunction",
+    "DepartmentsScanFunction",
+    "DictFilterEchoFunction",
+    "DoubleSequenceFunction",
+    "DynamicFilterEchoFunction",
+    "EmployeesScanFunction",
+    "ExpressionFilterTestFunction",
+    "FilterEchoFunction",
+    "FilterEchoPartitionedFunction",
+    "FilterEchoTableScanFunction",
+    "GeneratorExceptionFunction",
+    "ValuePruneFunction",
+    "LateMaterializationFunction",
+    "LoggingGeneratorFunction",
+    "MakePairsIntFunction",
+    "MakePairsIntStrFunction",
+    "MakePairsStrFunction",
+    "MakeSeriesCountFunction",
+    "MakeSeriesCsvFunction",
+    "MakeSeriesFloatFunction",
+    "MakeSeriesRangeFunction",
+    "MakeSeriesStepFunction",
+    "MissingBatchIndexTagFunction",
+    "NamedParamsEchoFunction",
+    "NestedSequenceFunction",
+    "NonMonotoneBatchIndexFunction",
+    "OrderEchoFunction",
+    "PartitionedBatchIndexFunction",
+    "PartitionedBatchIndexMarkedFunction",
+    "PartitionedFixedOrderFunction",
+    "PartitionedNoOrderGuaranteeFunction",
+    "PartitionedPreservesOrderFunction",
+    "PartitionedSequenceFunction",
+    "PartitionedWithExplicitOverrideFunction",
+    "ProductsScanFunction",
+    "ProfilingDemoFunction",
+    "ProjectedDataFunction",
+    "ProjectsScanFunction",
+    "RegionYearPartitionedFunction",
+    "RepeatValueIntFunction",
+    "RepeatValueStrFunction",
+    "RFF_MULTI_COLUMNS",
+    "RFF_NESTED_COLUMNS",
+    "RFF_NONE_COLUMNS",
+    "RFF_ROWID_COLUMNS",
+    "RFF_SIMPLE_COLUMNS",
+    "RFF_STRUCT_COLUMNS",
+    "RffMultiScanFunction",
+    "RffNestedScanFunction",
+    "RffNoneScanFunction",
+    "RffRowidScanFunction",
+    "RffSimpleScanFunction",
+    "RffStructScanFunction",
+    "RowIdSequenceFunction",
+    "SampleEchoFunction",
+    "ScopedSecretDemoFunction",
+    "SecretDemoFunction",
+    "SequenceFunction",
+    "SettingsAwareFunction",
+    "SpatialFilterExampleFunction",
+    "StructSettingsFunction",
+    "TenThousandFunction",
+    "TxCachedValueFunction",
+    "VersionedConstraintsScanFunction",
+    "VersionedDataFunction",
+    "resolve_version",
+    "resolve_versioned_constraints_version",
+]

vgi/_test_fixtures/table/_common.py

@@ -0,0 +1,162 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Shared infrastructure for table fixture functions.
+
+Holds the cardinality decorator, the common ``CountdownState``, the
+``CountBatchArgs`` base for fixtures that take ``(count, batch_size)``,
+and the ``_BaseSequenceFunction`` template-method base class for
+countdown-style generators.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Annotated, Any, ClassVar
+
+import numpy as np
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+from vgi_rpc.rpc import OutputCollector
+
+from vgi.arguments import Arg
+from vgi.catalog.catalog_interface import ColumnStatistics
+from vgi.table_function import (
+    BindParams,
+    ProcessParams,
+    TableCardinality,
+    TableFunctionGenerator,
+)
+
+
+def _cardinality_from_count[T: TableFunctionGenerator[Any, Any]](cls: type[T]) -> type[T]:
+    """Class decorator to implement cardinality() based on a 'count' argument."""
+    if "cardinality" not in cls.__dict__:  # only inject if subclass hasn't overridden
+
+        def cardinality_impl(cls_: type[T], params: BindParams[Any]) -> TableCardinality:
+            count = getattr(params.args, "count", None)
+            if not isinstance(count, int) or count < 0:
+                raise ValueError(f"Expected a non-negative integer 'count' argument for {cls_.__name__}")
+            return TableCardinality(estimate=count, max=count)
+
+        cls.cardinality = classmethod(cardinality_impl)  # type: ignore[assignment]
+
+    return cls
+
+
+@dataclass(kw_only=True)
+class CountdownState(ArrowSerializableDataclass):
+    """Mutable state tracking remaining rows and current position."""
+
+    remaining: int
+    current_index: int = 0
+
+
+@dataclass(frozen=True)
+class CountBatchArgs:
+    """Standard ``(count, batch_size)`` argument pair for countdown-style fixtures.
+
+    Subclass this to add fixture-specific knobs without re-declaring the two
+    common fields. Note: ``slots=True`` is intentionally omitted so subclasses
+    can extend cleanly without slot-conflict gymnastics.
+    """
+
+    count: Annotated[int, Arg(0, doc="Number of rows to generate", ge=0)]
+    batch_size: Annotated[int, Arg("batch_size", default=1000, doc="Batch size for output", ge=1)]
+
+
+@dataclass(slots=True, frozen=True)
+class _EmptyArgs:
+    """No arguments."""
+
+
+@dataclass(kw_only=True)
+class _OneShotState(ArrowSerializableDataclass):
+    """State that emits data once."""
+
+    done: bool = False
+
+
+class _BaseSequenceFunction(TableFunctionGenerator[Any, CountdownState]):
+    """Template-method base for countdown-style fixture generators.
+
+    Provides ``initial_state``, the countdown bookkeeping in ``process``, and
+    a default numpy-arange ``_emit_chunk`` used by SequenceFunction /
+    DoubleSequenceFunction. Subclasses with non-arange output (e.g. echoes,
+    nested types, row-id sequences) override ``_emit_chunk``.
+
+    ``BATCH_SIZE_FALLBACK`` is used when ``params.args`` has no ``batch_size``
+    field — i.e. fixtures that want a fixed batch size rather than a user knob.
+    """
+
+    NUMPY_DTYPE: ClassVar[type[np.generic]] = np.int64
+    STATS_ARROW_TYPE: ClassVar[pa.DataType] = pa.int64()
+    STATS_COLUMN_NAME: ClassVar[str] = "n"
+    BATCH_SIZE_FALLBACK: ClassVar[int] = 1000
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[Any]) -> CountdownState:
+        """Create initial state with remaining count."""
+        return CountdownState(remaining=params.args.count)
+
+    @classmethod
+    def statistics(cls, params: BindParams[Any]) -> list[ColumnStatistics] | None:
+        """Exact per-column stats derived from the user's bind args.
+
+        For sequence(count, increment=k): the output column spans
+        [0, (count - 1) * increment] with no nulls and count distinct values.
+        Returns ``None`` (no stats) for fixtures whose output isn't a single
+        ``int64`` arange — they should override.
+        """
+        count = getattr(params.args, "count", None)
+        increment = getattr(params.args, "increment", 1)
+        if not isinstance(count, int) or count <= 0:
+            return []
+        max_value = (count - 1) * increment
+        return [
+            ColumnStatistics(
+                column_name=cls.STATS_COLUMN_NAME,
+                min=pa.scalar(0, cls.STATS_ARROW_TYPE),
+                max=pa.scalar(max_value, cls.STATS_ARROW_TYPE),
+                has_null=False,
+                has_not_null=True,
+                distinct_count=count,
+            )
+        ]
+
+    @classmethod
+    def process(cls, params: ProcessParams[Any], state: CountdownState, out: OutputCollector) -> None:
+        """Run the standard countdown loop; delegate batch contents to ``_emit_chunk``."""
+        if state.remaining <= 0:
+            out.finish()
+            return
+
+        batch_size = getattr(params.args, "batch_size", cls.BATCH_SIZE_FALLBACK)
+        size = min(state.remaining, batch_size)
+        cls._emit_chunk(params, state, out, state.current_index, size)
+        state.current_index += size
+        state.remaining -= size
+
+    @classmethod
+    def _emit_chunk(
+        cls,
+        params: ProcessParams[Any],
+        state: CountdownState,
+        out: OutputCollector,
+        start: int,
+        size: int,
+    ) -> None:
+        """Default implementation: numpy arange × increment.
+
+        Subclasses with non-arange output override this hook. ``state`` is
+        passed in case subclasses want to track additional info; the standard
+        countdown bookkeeping (``current_index``/``remaining``) is handled by
+        ``process`` itself, so subclass hooks should NOT mutate them.
+        """
+        increment = params.args.increment
+        values = np.arange(
+            start * increment,
+            (start + size) * increment,
+            increment,
+            dtype=cls.NUMPY_DTYPE,
+        )
+        out.emit(pa.RecordBatch.from_arrays([pa.array(values)], schema=params.output_schema))

vgi/_test_fixtures/table/batch_index.py

@@ -0,0 +1,283 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Partitioned-queue fixtures that opt in to ``supports_batch_index``.
+
+These exist so SQL integration tests can verify the batch_index feature:
+
+* ``partitioned_batch_index(count)`` — single-column ``n int64`` output;
+  parallel scan with FIXED_ORDER preservation. Each queue item is tagged
+  with a stable partition_id; the worker emits Arrow batches tagged with
+  that id via ``out.emit(batch, batch_index=partition_id)``. The DuckDB
+  extension reads the tag from each batch's KeyValueMetadata, threads it
+  through ``TableFunction::get_partition_data``, and ordered sinks
+  (``PhysicalBatchCollector``, ``PhysicalBatchInsert``,
+  ``PhysicalBatchCopyToFile``, ``PhysicalLimit``) reassemble output in
+  partition_id order. The FIXED_ORDER ``MaxThreads=1`` clamp is dropped
+  for opted-in functions.
+
+* ``partitioned_batch_index_marked(count, chunk_size)`` — two-column
+  ``(partition_id int64, seq int64)`` output. Lets tests directly
+  observe partition boundaries in the output stream (e.g. "no row with
+  partition_id=N appears after a row with partition_id=N+1"). Projection
+  pushdown is disabled so the ``partition_id`` column survives even
+  ``SELECT seq FROM …`` queries.
+
+The worker uses the existing in-process ``state`` to carry per-worker
+cursor information across ``process()`` calls — same approach as
+``_BasePartitionedOrderMode`` in ``order_modes.py``. HTTP transport's
+existing STATE_KEY mechanism (in vgi_rpc.http) round-trips this state
+across requests; nothing new is added for HTTP resumption.
+"""
+
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass
+from typing import Annotated, 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.invocation import GlobalInitResponse
+from vgi.metadata import FunctionExample, OrderPreservation
+from vgi.protocol import VgiOutputCollector
+from vgi.schema_utils import schema
+from vgi.table_function import (
+    InitParams,
+    ProcessParams,
+    TableFunctionGenerator,
+    bind_fixed_schema,
+)
+
+# Queue-item encoding: (partition_id, start, end) packed as three uint64s.
+# Decoded by ``process()`` on the worker; partition_id is what the worker
+# emits to DuckDB via the batch_index= kwarg.
+_ITEM_FMT = ">QQQ"
+_ITEM_SIZE = struct.calcsize(_ITEM_FMT)
+
+
+# =============================================================================
+# Single-column variant: partitioned_batch_index(count)
+# =============================================================================
+
+
+@dataclass(slots=True, frozen=True)
+class _BatchIndexArgs:
+    """Arguments for ``partitioned_batch_index``."""
+
+    count: Annotated[int, Arg(0, doc="Total number of integers to generate", ge=0)]
+
+
+@dataclass(kw_only=True)
+class _BatchIndexState(ArrowSerializableDataclass):
+    """Per-worker cursor state.
+
+    ``partition_id`` is the queue-push order of the current work item; emitted
+    on every Arrow batch via the batch_index= kwarg. ``current_idx`` advances
+    through ``[current_start, current_end)`` as the worker produces batches.
+    All three reset to None at the moment a partition is exhausted; the next
+    ``process()`` call pops a fresh item.
+    """
+
+    partition_id: int | None = None
+    current_start: int | None = None
+    current_end: int | None = None
+    current_idx: int = 0
+
+
+@bind_fixed_schema
+@_cardinality_from_count
+class PartitionedBatchIndexFunction(TableFunctionGenerator[_BatchIndexArgs, _BatchIndexState]):
+    """Parallel-scan sequence with batch_index ordering.
+
+    The primary worker enqueues N work items at on_init, each encoding
+    ``(partition_id, start, end)``. Any worker pulls the next item via
+    ``queue_pop``; emits a stream of Arrow batches tagged with
+    partition_id; advances to the next item when exhausted. DuckDB's
+    ordered sinks reassemble output in partition_id order — final output
+    matches a single-threaded scan, but the source itself fans out across
+    threads.
+    """
+
+    # NOTE: left at the original small chunk for now. Capping the partition
+    # count (as partitioned_sequence / filter_echo_partitioned do) once made
+    # batch_index.test segfault (exit 139) — but that turned out to be a
+    # pre-existing use-after-free in the extension's async cancel path (the
+    # cancel dispatcher logged through the destroyed query ClientContext), not
+    # a batch_index bug. The cap just triggered more stream cancellations and
+    # exposed it. With that fixed in the extension, capping this fixture is now
+    # safe to re-enable (verified 15/15 clean under UBSan); deferred only so the
+    # resize lands as its own change.
+    CHUNK_SIZE: ClassVar[int] = 1000
+    BATCH_SIZE: ClassVar[int] = 1000
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = schema(n=pa.int64())
+
+    class Meta:
+        name = "partitioned_batch_index"
+        description = (
+            "Multi-worker partitioned sequence with per-batch batch_index "
+            "tagging; parallel scan + ordered sink reassembly."
+        )
+        categories = ["generator", "utility"]
+        preserves_order = OrderPreservation.FIXED_ORDER
+        supports_batch_index = True
+        examples = [
+            FunctionExample(
+                sql="SELECT * FROM partitioned_batch_index(100)",
+                description=(
+                    "Generate 0..99 in parallel; DuckDB sinks reassemble output in partition_id (insertion) order."
+                ),
+            ),
+        ]
+
+    @classmethod
+    def on_init(cls, params: InitParams[_BatchIndexArgs]) -> GlobalInitResponse:
+        work_items: list[bytes] = []
+        for partition_id, start_idx in enumerate(range(0, params.args.count, cls.CHUNK_SIZE)):
+            end_idx = min(start_idx + cls.CHUNK_SIZE, params.args.count)
+            work_items.append(struct.pack(_ITEM_FMT, partition_id, start_idx, end_idx))
+        params.storage.queue_push(work_items)
+        return GlobalInitResponse()
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_BatchIndexArgs]) -> _BatchIndexState:
+        return _BatchIndexState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_BatchIndexArgs],
+        state: _BatchIndexState,
+        out: OutputCollector,
+    ) -> None:
+        if state.partition_id is None or state.current_idx >= (state.current_end or 0):
+            work_data = params.storage.queue_pop()
+            if work_data is None:
+                out.finish()
+                return
+            partition_id, start, end = struct.unpack(_ITEM_FMT, work_data)
+            state.partition_id = partition_id
+            state.current_start = start
+            state.current_end = end
+            state.current_idx = start
+
+        batch_end_idx = min(state.current_idx + cls.BATCH_SIZE, state.current_end or 0)
+        values = list(range(state.current_idx, batch_end_idx))
+        cast(VgiOutputCollector, out).emit(
+            pa.RecordBatch.from_pydict({"n": values}, schema=params.output_schema),
+            batch_index=state.partition_id,
+        )
+        state.current_idx = batch_end_idx
+
+
+# =============================================================================
+# Two-column variant: partitioned_batch_index_marked(count, chunk_size)
+# =============================================================================
+
+
+@dataclass(slots=True, frozen=True)
+class _BatchIndexMarkedArgs:
+    """Arguments for ``partitioned_batch_index_marked``."""
+
+    count: Annotated[int, Arg(0, doc="Total number of rows to generate", ge=0)]
+    chunk_size: Annotated[int, Arg("chunk_size", default=1000, doc="Rows per partition", ge=1)]
+
+
+@dataclass(kw_only=True)
+class _BatchIndexMarkedState(ArrowSerializableDataclass):
+    partition_id: int | None = None
+    current_start: int | None = None
+    current_end: int | None = None
+    current_idx: int = 0
+
+
+@bind_fixed_schema
+@_cardinality_from_count
+class PartitionedBatchIndexMarkedFunction(TableFunctionGenerator[_BatchIndexMarkedArgs, _BatchIndexMarkedState]):
+    """Two-column batch_index fixture for direct ordering observation.
+
+    Output rows are ``(partition_id, seq)`` where ``partition_id`` is the
+    queue-push order (matches the emitted batch_index) and ``seq`` counts
+    up within each partition starting at 0. Tests assert that no row with
+    a higher partition_id appears before a row with a lower one — proving
+    that DuckDB's sink-side reassembly worked.
+
+    Projection pushdown is OFF on this fixture so ``SELECT seq FROM …``
+    still gets the partition_id column emitted by the worker; the C++
+    extension's projection then drops it after the ordering metadata has
+    been observed.
+    """
+
+    BATCH_SIZE: ClassVar[int] = 256
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = schema(
+        partition_id=pa.int64(),
+        seq=pa.int64(),
+    )
+
+    class Meta:
+        name = "partitioned_batch_index_marked"
+        description = (
+            "Two-column batch_index demo: rows are (partition_id, seq). Tests "
+            "assert that DuckDB's ordered sinks reassemble output in "
+            "partition_id order under parallel execution."
+        )
+        categories = ["generator", "utility", "testing"]
+        preserves_order = OrderPreservation.FIXED_ORDER
+        supports_batch_index = True
+        projection_pushdown = False
+        examples = [
+            FunctionExample(
+                sql="SELECT * FROM partitioned_batch_index_marked(100, chunk_size := 25) LIMIT 5",
+                description="First 5 rows of partition 0",
+            ),
+        ]
+
+    @classmethod
+    def on_init(cls, params: InitParams[_BatchIndexMarkedArgs]) -> GlobalInitResponse:
+        work_items: list[bytes] = []
+        chunk_size = params.args.chunk_size
+        for partition_id, start_idx in enumerate(range(0, params.args.count, chunk_size)):
+            end_idx = min(start_idx + chunk_size, params.args.count)
+            work_items.append(struct.pack(_ITEM_FMT, partition_id, start_idx, end_idx))
+        params.storage.queue_push(work_items)
+        return GlobalInitResponse()
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_BatchIndexMarkedArgs]) -> _BatchIndexMarkedState:
+        return _BatchIndexMarkedState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_BatchIndexMarkedArgs],
+        state: _BatchIndexMarkedState,
+        out: OutputCollector,
+    ) -> None:
+        if state.partition_id is None or state.current_idx >= (state.current_end or 0):
+            work_data = params.storage.queue_pop()
+            if work_data is None:
+                out.finish()
+                return
+            partition_id, start, end = struct.unpack(_ITEM_FMT, work_data)
+            state.partition_id = partition_id
+            state.current_start = start
+            state.current_end = end
+            state.current_idx = start
+
+        batch_end_idx = min(state.current_idx + cls.BATCH_SIZE, state.current_end or 0)
+        rows = batch_end_idx - state.current_idx
+        partition_ids = [state.partition_id] * rows
+        seqs = list(range(state.current_idx - (state.current_start or 0), batch_end_idx - (state.current_start or 0)))
+        cast(VgiOutputCollector, out).emit(
+            pa.RecordBatch.from_pydict(
+                {"partition_id": partition_ids, "seq": seqs},
+                schema=params.output_schema,
+            ),
+            batch_index=state.partition_id,
+        )
+        state.current_idx = batch_end_idx