vgi-python 0.8.7__py3-none-any.whl → 0.8.8__py3-none-any.whl

vgi/_test_fixtures/copy_from.py

@@ -0,0 +1,99 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Fixture ``COPY ... FROM`` format reader for VGI integration tests.
+
+``ExampleLinesCopyFromFunction`` registers the SQL format ``example_lines`` — a
+toy delimited-text reader. It exercises the full COPY-FROM path plus the option
+machinery: a defaulted option (``delimiter``), an ``INTEGER`` option with a range
+constraint (``skip_rows``), a required option (``null_string``), and an
+enum/``choices`` option (``on_error``).
+
+Usage::
+
+    CREATE TABLE t (a INTEGER, b VARCHAR);
+    COPY t FROM '/path/data.txt' (FORMAT example_lines, null_string 'NA');
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Annotated, ClassVar
+
+import pyarrow as pa
+
+from vgi.arguments import Arg
+from vgi.copy_from_function import CopyFromFunction
+
+if TYPE_CHECKING:
+    from vgi_rpc.rpc import OutputCollector
+
+    from vgi.table_function import ProcessParams
+
+__all__ = ["ExampleLinesCopyFromFunction"]
+
+
+@dataclass(slots=True, frozen=True, kw_only=True)
+class ExampleLinesCopyFromArgs:
+    """Options for the ``example_lines`` COPY format."""
+
+    null_string: Annotated[str, Arg("null_string", doc="Token parsed as SQL NULL")]
+    delimiter: Annotated[str, Arg("delimiter", default=",", doc="Field separator")] = ","
+    skip_rows: Annotated[int, Arg("skip_rows", default=0, ge=0, doc="Leading lines to skip before data")] = 0
+    on_error: Annotated[
+        str,
+        Arg(
+            "on_error",
+            default="fail",
+            choices=["fail", "skip"],
+            doc="Behavior on a row whose column count does not match the target",
+        ),
+    ] = "fail"
+
+
+class ExampleLinesCopyFromFunction(CopyFromFunction[ExampleLinesCopyFromArgs]):
+    """Toy delimited-text ``COPY ... FROM`` reader (test fixture)."""
+
+    COPY_FROM_FORMAT: ClassVar[str] = "example_lines"
+    COPY_FROM_COMMENT: ClassVar[str | None] = "Toy delimited-text reader for tests"
+
+    class Meta:
+        name = "example_lines_copy_reader"
+        description = "Read a delimited text file into the COPY target table"
+        categories = ["copy", "test"]
+        tags = {"category": "copy_from", "stability": "test"}
+
+    @classmethod
+    def read(
+        cls,
+        *,
+        path: str,
+        options: ExampleLinesCopyFromArgs,
+        expected_schema: pa.Schema,
+        params: ProcessParams[ExampleLinesCopyFromArgs],
+        out: OutputCollector,
+    ) -> None:
+        """Parse ``path`` line-by-line and emit one batch matching ``expected_schema``."""
+        with open(path, encoding="utf-8") as fh:
+            lines = fh.read().splitlines()
+        lines = lines[options.skip_rows :]
+
+        ncols = len(expected_schema)
+        rows: list[list[str]] = []
+        for line in lines:
+            if line == "":
+                continue
+            cells = line.split(options.delimiter)
+            if len(cells) != ncols:
+                if options.on_error == "skip":
+                    continue
+                raise ValueError(f"example_lines: row has {len(cells)} fields, expected {ncols}: {line!r}")
+            rows.append(cells)
+
+        # Column-major string arrays, NULL where the cell equals null_string,
+        # then cast each column to the target type (DuckDB inserts no cast).
+        columns = list(zip(*rows)) if rows else [() for _ in range(ncols)]
+        arrays = []
+        for idx, field in enumerate(expected_schema):
+            raw = [None if v == options.null_string else v for v in columns[idx]]
+            arrays.append(pa.array(raw, type=pa.string()).cast(field.type))
+        out.emit(pa.RecordBatch.from_arrays(arrays, schema=expected_schema))

vgi/_test_fixtures/copy_to.py

@@ -0,0 +1,160 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Fixture ``COPY ... TO`` format writer for VGI integration tests.
+
+``ExampleLinesCopyToFunction`` registers the SQL format ``example_lines_out`` — a
+toy delimited-text writer, the symmetric counterpart of the ``example_lines``
+reader. It exercises the COPY-TO Sink+Combine path plus the option machinery: a
+required option (``null_string``), a defaulted option (``delimiter``), a BOOLEAN
+option (``header``), and an enum/``choices`` option (``on_exists``).
+
+Shards are buffered in ``params.storage`` (``execution_id``-scoped) by ``write()``
+and concatenated to the destination by ``close()`` — the cross-process-safe
+pattern, so it works under pool rotation / HTTP.
+
+Usage::
+
+    COPY (SELECT * FROM t) TO '/path/out.txt' (FORMAT 'acme.example_lines_out', null_string 'NA');
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Annotated, ClassVar
+
+import pyarrow as pa
+
+from vgi.arguments import Arg
+from vgi.copy_to_function import CopyToFunction
+
+if TYPE_CHECKING:
+    from vgi.table_buffering_function import TableBufferingParams
+
+__all__ = ["ExampleLinesCopyToFunction", "ExampleLinesOrderedCopyToFunction"]
+
+_SHARD_NS = b"copy_to_shard"
+
+
+@dataclass(slots=True, frozen=True, kw_only=True)
+class ExampleLinesCopyToArgs:
+    """Options for the ``example_lines_out`` COPY format."""
+
+    null_string: Annotated[str, Arg("null_string", doc="Token written for SQL NULL")]
+    delimiter: Annotated[str, Arg("delimiter", default=",", doc="Field separator")] = ","
+    header: Annotated[bool, Arg("header", default=False, doc="Write a header row of column names")] = False
+    header_repeat: Annotated[
+        int,
+        Arg("header_repeat", default=1, ge=0, le=3, doc="When header=true, write the header line this many times"),
+    ] = 1
+    on_exists: Annotated[
+        str,
+        Arg(
+            "on_exists",
+            default="overwrite",
+            choices=["overwrite", "error"],
+            doc="Behavior when the destination file already exists",
+        ),
+    ] = "overwrite"
+    fail_on_value: Annotated[
+        str,
+        Arg("fail_on_value", default="", doc="If non-empty, fail mid-write when a cell equals this value"),
+    ] = ""
+
+
+class ExampleLinesCopyToFunction(CopyToFunction[ExampleLinesCopyToArgs]):
+    """Toy delimited-text ``COPY ... TO`` writer (test fixture)."""
+
+    COPY_TO_FORMAT: ClassVar[str] = "example_lines_out"
+    COPY_TO_COMMENT: ClassVar[str | None] = "Toy delimited-text writer for tests"
+
+    class Meta:
+        name = "example_lines_writer"
+        description = "Write the COPY source to a delimited text file"
+        categories = ["copy", "test"]
+        tags = {"category": "copy_to", "stability": "test"}
+
+    @classmethod
+    def write(
+        cls,
+        *,
+        batch: pa.RecordBatch,
+        options: ExampleLinesCopyToArgs,
+        file_path: str,
+        params: TableBufferingParams[ExampleLinesCopyToArgs],
+    ) -> None:
+        """Buffer one input batch as an IPC blob in execution-scoped storage."""
+        # Mid-sink failure trigger: raise during a process() call when a cell
+        # matches fail_on_value. Exercises the in-flight teardown/recovery path.
+        if options.fail_on_value:
+            for col in batch.columns:
+                for value in col.to_pylist():
+                    if value is not None and str(value) == options.fail_on_value:
+                        raise ValueError(f"example_lines_out: fail_on_value hit: {options.fail_on_value!r}")
+        sink = pa.BufferOutputStream()
+        with pa.ipc.new_stream(sink, batch.schema) as writer:
+            writer.write_batch(batch)
+        # state_append is atomic + race-safe across parallel sink threads/workers.
+        params.storage.state_append(_SHARD_NS, b"", sink.getvalue().to_pybytes())
+
+    @classmethod
+    def close(
+        cls,
+        *,
+        options: ExampleLinesCopyToArgs,
+        file_path: str,
+        params: TableBufferingParams[ExampleLinesCopyToArgs],
+    ) -> int:
+        """Concatenate every shard and write the delimited destination file (once)."""
+        import os
+
+        if options.on_exists == "error" and os.path.exists(file_path):
+            raise FileExistsError(f"example_lines_out: destination already exists: {file_path}")
+
+        shards = params.storage.state_log_scan(_SHARD_NS, b"", after_id=-1)
+
+        def fmt(value: object) -> str:
+            return options.null_string if value is None else str(value)
+
+        def write_header(fh: object, names: list[str]) -> None:
+            # header=true writes the column-name line `header_repeat` times.
+            if options.header:
+                for _ in range(options.header_repeat):
+                    fh.write(options.delimiter.join(names) + "\n")  # type: ignore[attr-defined]
+
+        rows_written = 0
+        with open(file_path, "w", encoding="utf-8") as fh:
+            wrote_header = False
+            for _log_id, blob in shards:
+                table = pa.ipc.open_stream(blob).read_all()
+                if not wrote_header:
+                    write_header(fh, list(table.schema.names))
+                    wrote_header = True
+                for row in table.to_pylist():
+                    fh.write(options.delimiter.join(fmt(row[name]) for name in table.schema.names) + "\n")
+                    rows_written += 1
+            # Empty COPY with header=true still emits the header row(s). We need the
+            # source column names; they ride the bind's input_schema.
+            if not wrote_header:
+                assert params.init_call is not None
+                in_schema = params.init_call.bind_call.input_schema
+                if in_schema is not None:
+                    write_header(fh, list(in_schema.names))
+        return rows_written
+
+
+class ExampleLinesOrderedCopyToFunction(ExampleLinesCopyToFunction):
+    """Ordered variant of :class:`ExampleLinesCopyToFunction`.
+
+    ``Meta.ordered = True`` makes the extension use a single-threaded sink, so the
+    worker receives every batch in source order and writes the file in order.
+    """
+
+    COPY_TO_FORMAT: ClassVar[str] = "example_lines_ordered_out"
+    COPY_TO_COMMENT: ClassVar[str | None] = "Toy delimited-text writer (ordered, single-thread sink)"
+
+    class Meta:
+        name = "example_lines_ordered_writer"
+        description = "Write the COPY source to a delimited file, preserving source order"
+        categories = ["copy", "test"]
+        tags = {"category": "copy_to", "stability": "test"}
+        sink_order_dependent = True  # ordered COPY TO → single-thread sink

vgi/_test_fixtures/table/__init__.py

@@ -78,6 +78,7 @@ from vgi._test_fixtures.table.pairs import (
     MakePairsStrFunction,
     RepeatValueIntFunction,
     RepeatValueStrFunction,
+    UnionVarargsFunction,
 )
 from vgi._test_fixtures.table.partition_columns import (
     CountryPartitionedSalesFunction,
@@ -119,6 +120,7 @@ from vgi._test_fixtures.table.sequence import (
     TenThousandFunction,
 )
 from vgi._test_fixtures.table.settings import (
+    MultiSecretDemoFunction,
     ScopedSecretDemoFunction,
     SecretDemoFunction,
     SettingsAwareFunction,
@@ -198,6 +200,7 @@ __all__ = [
     "RegionYearPartitionedFunction",
     "RepeatValueIntFunction",
     "RepeatValueStrFunction",
+    "UnionVarargsFunction",
     "RFF_MULTI_COLUMNS",
     "RFF_NESTED_COLUMNS",
     "RFF_NONE_COLUMNS",
@@ -212,6 +215,7 @@ __all__ = [
     "RffStructScanFunction",
     "RowIdSequenceFunction",
     "SampleEchoFunction",
+    "MultiSecretDemoFunction",
     "ScopedSecretDemoFunction",
     "SecretDemoFunction",
     "SequenceFunction",

vgi/_test_fixtures/table/pairs.py

@@ -14,7 +14,7 @@ from vgi_rpc.rpc import OutputCollector
 from vgi._test_fixtures.table._common import (
     _cardinality_from_count,
 )
-from vgi.arguments import Arg
+from vgi.arguments import Arg, TaggedUnion
 from vgi.invocation import BindResponse
 from vgi.metadata import FunctionExample
 from vgi.schema_utils import schema
@@ -447,3 +447,109 @@ class RepeatValueStrFunction(TableFunctionGenerator[RepeatValueStrArgs, RepeatVa
         data = {f"v{i}": col for i, col in enumerate(state.rows)}
         out_schema = schema({f"v{i}": pa.string() for i in range(len(state.rows))})
         out.emit(pa.RecordBatch.from_pydict(data, schema=out_schema))
+
+
+# ============================================================================
+
+# Sparse union shared by every union_varargs argument. DuckDB only ever emits
+# sparse unions (+us:) over Arrow, so this round-trips end-to-end.
+UNION_VARARGS_TYPE = pa.sparse_union([pa.field("i", pa.int64()), pa.field("s", pa.string())])
+
+UNION_VARARGS_SCHEMA = schema(idx=pa.int64(), tag=pa.string(), value=pa.string())
+
+
+@dataclass(kw_only=True)
+class UnionVarargsArgs:
+    """Arguments for union_varargs."""
+
+    configs: Annotated[
+        tuple[TaggedUnion, ...],
+        Arg(
+            0,
+            varargs=True,
+            arrow_type=UNION_VARARGS_TYPE,
+            doc="Union values whose active member tag is echoed back",
+        ),
+    ]
+
+
+@dataclass(kw_only=True)
+class UnionVarargsState(ArrowSerializableDataclass):
+    """State for union_varargs."""
+
+    idx: list[int] = field(default_factory=list)
+    tags: list[str | None] = field(default_factory=list)
+    values: list[str] = field(default_factory=list)
+    done: bool = False
+
+
+@init_single_worker
+@bind_fixed_schema
+class UnionVarargsFunction(TableFunctionGenerator[UnionVarargsArgs, UnionVarargsState]):
+    """Echo the active member tag and value of each union vararg.
+
+    USE CASE
+    --------
+    Exercises union-typed varargs: each argument arrives as a
+    [`TaggedUnion`][vgi.arguments.TaggedUnion] so the active member
+    discriminator (which a plain ``Scalar.as_py()`` would drop) is preserved.
+    Emits one row per vararg with its positional index, the active member
+    name, and the member value stringified into a single fixed column.
+
+    SCHEMA
+    ------
+    Fixed: ``{"idx": int64, "tag": string, "value": string}``.
+
+    Example:
+        SELECT * FROM union_varargs(
+            union_value(i := 1)::UNION(i INT, s VARCHAR),
+            union_value(s := 'x')::UNION(i INT, s VARCHAR))
+        Returns: (0, 'i', '1'), (1, 's', 'x')
+
+    Attributes:
+        FIXED_SCHEMA: The fixed Arrow output schema this function always produces.
+
+    """
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = UNION_VARARGS_SCHEMA
+
+    class Meta:
+        """Function metadata."""
+
+        name = "union_varargs"
+        description = "Echo the active member tag and value of each union vararg"
+        categories = ["generator", "utility"]
+        examples = [
+            FunctionExample(
+                sql=(
+                    "SELECT * FROM union_varargs("
+                    "union_value(i := 1)::UNION(i INT, s VARCHAR), "
+                    "union_value(s := 'x')::UNION(i INT, s VARCHAR))"
+                ),
+                description="Echo the tag and value of two union arguments",
+            ),
+        ]
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[UnionVarargsArgs]) -> UnionVarargsState:
+        """Decompose each union vararg into (idx, tag, value) rows."""
+        configs = params.args.configs
+        return UnionVarargsState(
+            idx=list(range(len(configs))),
+            tags=[cfg.tag for cfg in configs],
+            values=[str(cfg.value) for cfg in configs],
+        )
+
+    @classmethod
+    def process(cls, params: ProcessParams[UnionVarargsArgs], state: UnionVarargsState, out: OutputCollector) -> None:
+        """Emit one row per union vararg."""
+        if state.done:
+            out.finish()
+            return
+        state.done = True
+        out.emit(
+            pa.RecordBatch.from_pydict(
+                {"idx": state.idx, "tag": state.tags, "value": state.values},
+                schema=UNION_VARARGS_SCHEMA,
+            )
+        )

vgi/_test_fixtures/table/settings.py

@@ -321,7 +321,7 @@ class SecretDemoFunction(TableFunctionGenerator[None, SecretDemoState]):
     @classmethod
     def initial_state(cls, params: ProcessParams[None]) -> SecretDemoState:
         """Build initial state from secret key-value pairs."""
-        secret = params.secrets.get("vgi_example", {})
+        secret = next(iter(params.secrets.of_type("vgi_example")), {})
         keys = list(secret.keys())
         values = [str(v.as_py()) for v in secret.values()]
         types = [str(v.type) for v in secret.values()]
@@ -406,7 +406,7 @@ class ScopedSecretDemoFunction(TableFunctionGenerator[ScopedSecretDemoArgs, Scop
     @classmethod
     def initial_state(cls, params: ProcessParams[ScopedSecretDemoArgs]) -> ScopedSecretDemoState:
         """Build state from resolved secrets."""
-        secret = params.secrets.get("vgi_example", {})
+        secret = next(iter(params.secrets.of_type("vgi_example")), {})
         return ScopedSecretDemoState(
             found=bool(secret),
             secret_keys=",".join(secret.keys()) if secret else "",
@@ -430,3 +430,58 @@ class ScopedSecretDemoFunction(TableFunctionGenerator[ScopedSecretDemoArgs, Scop
         )
         out.emit(batch)
         out.finish()
+
+
+@dataclass(kw_only=True)
+class MultiSecretDemoState(ArrowSerializableDataclass):
+    """State for MultiSecretDemoFunction."""
+
+    api_key: str = ""
+
+
+@init_single_worker
+class MultiSecretDemoFunction(TableFunctionGenerator[ScopedSecretDemoArgs, MultiSecretDemoState]):
+    """Resolve TWO same-type scoped secrets in one bind, then select per path.
+
+    Requests the ``vgi_example`` secret for both ``s3://bucket-a/`` and
+    ``s3://bucket-b/`` scopes in a single bind. Because resolved secrets are keyed
+    by name, both survive; ``for_scope_of_type`` then picks the one whose scope
+    matches the ``path`` argument and returns its ``api_key``.
+    """
+
+    class Meta:
+        """Metadata for MultiSecretDemoFunction."""
+
+        name = "multi_secret_demo"
+        description = "Demo: two same-type scoped secrets resolved in one bind"
+
+    @classmethod
+    def on_bind(cls, params: BindParams[ScopedSecretDemoArgs]) -> BindResponse:
+        """Request the secret for two distinct scopes of the same type."""
+        params.secrets.get("vgi_example", scope="s3://bucket-a/")
+        params.secrets.get("vgi_example", scope="s3://bucket-b/")
+        return BindResponse(output_schema=schema({"api_key": pa.string()}))
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[ScopedSecretDemoArgs]) -> MultiSecretDemoState:
+        """Select the resolved secret matching the path and return its api_key."""
+        secret = params.secrets.for_scope_of_type(params.args.path, "vgi_example") or {}
+        api_key = secret.get("api_key")
+        if api_key is not None and hasattr(api_key, "as_py"):
+            api_key = api_key.as_py()
+        return MultiSecretDemoState(api_key="" if api_key is None else str(api_key))
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[ScopedSecretDemoArgs],
+        state: MultiSecretDemoState,
+        out: OutputCollector,
+    ) -> None:
+        """Emit the resolved api_key for the path-matched secret."""
+        batch = pa.RecordBatch.from_pydict(
+            {"api_key": [state.api_key]},
+            schema=params.output_schema,
+        )
+        out.emit(batch)
+        out.finish()

vgi/_test_fixtures/worker.py

@@ -59,6 +59,8 @@ from vgi._test_fixtures.cancellable import (
     SlowCancellableFunction,
     SlowCancellableInOutFunction,
 )
+from vgi._test_fixtures.copy_from import ExampleLinesCopyFromFunction
+from vgi._test_fixtures.copy_to import ExampleLinesCopyToFunction, ExampleLinesOrderedCopyToFunction
 from vgi._test_fixtures.nest_tensor import NestTensorFunction, UnnestTensorFunction, UnnestTensorRowsFunction
 from vgi._test_fixtures.scalar import (
     AddValuesFunction,
@@ -143,6 +145,7 @@ from vgi._test_fixtures.table import (
     MakeSeriesRangeFunction,
     MakeSeriesStepFunction,
     MissingBatchIndexTagFunction,
+    MultiSecretDemoFunction,
     NamedParamsEchoFunction,
     NestedSequenceFunction,
     NonMonotoneBatchIndexFunction,
@@ -179,6 +182,7 @@ from vgi._test_fixtures.table import (
     TenThousandFunction,
     TxCachedValueFunction,
     TypedProbeFunction,
+    UnionVarargsFunction,
     ValuePruneFunction,
     VersionedConstraintsScanFunction,
     VersionedDataFunction,
@@ -349,6 +353,11 @@ _EXAMPLE_CATALOG = Catalog(
                 EchoBufferingFunction,
                 BufferEmitWideFunction,
                 SlowCancellableBufferingFunction,
+                # CopyFromFunction - custom COPY ... FROM format reader
+                ExampleLinesCopyFromFunction,
+                # CopyToFunction - custom COPY ... TO format writer
+                ExampleLinesCopyToFunction,
+                ExampleLinesOrderedCopyToFunction,
                 # TableFunctionGenerator - generate output without input
                 ConstantColumnsFunction,
                 SlowCancellableFunction,
@@ -374,6 +383,7 @@ _EXAMPLE_CATALOG = Catalog(
                 MakePairsStrFunction,
                 RepeatValueIntFunction,
                 RepeatValueStrFunction,
+                UnionVarargsFunction,
                 NamedParamsEchoFunction,
                 NestedSequenceFunction,
                 ProfilingDemoFunction,
@@ -409,6 +419,7 @@ _EXAMPLE_CATALOG = Catalog(
                 SampleEchoFunction,
                 RowIdSequenceFunction,
                 SecretDemoFunction,
+                MultiSecretDemoFunction,
                 ScopedSecretDemoFunction,
                 ExpressionFilterTestFunction,
                 SequenceFunction,