Skip to content

Commit 6b51e61

Browse files
rustyconoverclaude
andcommitted
feat(copy-to): worker-side COPY ... TO custom format support
Mirror the COPY FROM feature for the write direction: catalogs advertise custom COPY ... TO formats and a worker writes the streamed rows to a destination. - protocol.py: CopyToContext on BindRequest (format + path; additive, no version bump). Source columns ride the existing input_schema. - copy_to_function.py: CopyToFunction base — a buffered (Sink+Combine) function with no Source phase: write() persists a shard (execution_id-scoped storage), close() does the terminal write. Ordered output via Meta.sink_order_dependent. - catalog_interface.py: catalog_copy_from_formats discovery also enumerates CopyToFunction subclasses (direction='to'); CopyFromFormatInfo gains `ordered`. - _test_fixtures/copy_to.py: example_lines_out + example_lines_ordered_out writers. - Reuses the existing table_buffering_process/combine worker handlers + dispatch. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent d5ba35e commit 6b51e61

6 files changed

Lines changed: 505 additions & 15 deletions

File tree

‎tests/test_copy_to_function.py‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
# Copyright 2025, 2026 Query Farm LLC - https://query.farm
2+
3+
"""Worker-side unit tests for CopyToFunction and the example_lines_out format."""
4+
5+
from __future__ import annotations
6+
7+
import tempfile
8+
import types
9+
10+
import pyarrow as pa
11+
12+
from vgi._test_fixtures.copy_to import ExampleLinesCopyToArgs, ExampleLinesCopyToFunction
13+
from vgi._test_fixtures.worker import ExampleCatalog
14+
15+
SCHEMA = pa.schema([("a", pa.int64()), ("b", pa.string())])
16+
17+
18+
class _Store:
19+
"""Minimal in-memory BoundStorage stub (append + ordered log scan)."""
20+
21+
def __init__(self) -> None:
22+
self.log: list[tuple[int, bytes]] = []
23+
24+
def state_append(self, ns: bytes, key: bytes, val: bytes) -> None:
25+
self.log.append((len(self.log), val))
26+
27+
def state_log_scan(self, ns: bytes, key: bytes, after_id: int = -1, limit: int | None = None) -> list:
28+
rows = [(i, v) for (i, v) in self.log if i > after_id]
29+
return rows if limit is None else rows[:limit]
30+
31+
32+
def _tmp_path() -> str:
33+
with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as fh:
34+
return fh.name
35+
36+
37+
def _read(path: str) -> str:
38+
with open(path, encoding="utf-8") as fh:
39+
return fh.read()
40+
41+
42+
def _params(store: _Store) -> types.SimpleNamespace:
43+
bind_call = types.SimpleNamespace(input_schema=SCHEMA)
44+
init_call = types.SimpleNamespace(bind_call=bind_call)
45+
return types.SimpleNamespace(storage=store, init_call=init_call, execution_id=b"x", args=None)
46+
47+
48+
def test_write_then_close_round_trips_with_null_string() -> None:
49+
"""write() buffers shards; close() concatenates them to a delimited file."""
50+
store = _Store()
51+
params = _params(store)
52+
opts = ExampleLinesCopyToArgs(null_string="NA")
53+
out_name = _tmp_path()
54+
55+
ExampleLinesCopyToFunction.write(
56+
batch=pa.record_batch({"a": [1, 2], "b": ["foo", None]}, schema=SCHEMA),
57+
options=opts,
58+
file_path=out_name,
59+
params=params,
60+
)
61+
ExampleLinesCopyToFunction.write(
62+
batch=pa.record_batch({"a": [3], "b": ["baz"]}, schema=SCHEMA),
63+
options=opts,
64+
file_path=out_name,
65+
params=params,
66+
)
67+
n = ExampleLinesCopyToFunction.close(options=opts, file_path=out_name, params=params)
68+
assert n == 3
69+
assert _read(out_name) == "1,foo\n2,NA\n3,baz\n"
70+
71+
72+
def test_close_honors_delimiter_and_header() -> None:
73+
"""Non-default delimiter + header row are applied."""
74+
store = _Store()
75+
params = _params(store)
76+
opts = ExampleLinesCopyToArgs(null_string="NA", delimiter="|", header=True)
77+
out_name = _tmp_path()
78+
ExampleLinesCopyToFunction.write(
79+
batch=pa.record_batch({"a": [1], "b": ["x"]}, schema=SCHEMA),
80+
options=opts,
81+
file_path=out_name,
82+
params=params,
83+
)
84+
n = ExampleLinesCopyToFunction.close(options=opts, file_path=out_name, params=params)
85+
assert n == 1
86+
assert _read(out_name) == "a|b\n1|x\n"
87+
88+
89+
def test_close_empty_input_with_header_writes_header_only() -> None:
90+
"""An empty COPY with header=true still emits the header row (0 data rows)."""
91+
store = _Store()
92+
params = _params(store)
93+
out_name = _tmp_path()
94+
n = ExampleLinesCopyToFunction.close(
95+
options=ExampleLinesCopyToArgs(null_string="NA", header=True),
96+
file_path=out_name,
97+
params=params,
98+
)
99+
assert n == 0
100+
assert _read(out_name) == "a,b\n"
101+
102+
103+
def test_catalog_advertises_copy_to_format() -> None:
104+
"""The example catalog advertises example_lines_out with direction='to'."""
105+
formats = ExampleCatalog().copy_from_formats(attach_opaque_data=b"", transaction_opaque_data=None)
106+
by = {(f.direction, f.format_name): f for f in formats}
107+
assert ("to", "example_lines_out") in by
108+
fmt = by[("to", "example_lines_out")]
109+
assert fmt.handler == "example_lines_writer"
110+
assert fmt.comment == "Toy delimited-text writer for tests"
111+
assert fmt.tags.get("category") == "copy_to"
112+
opt_schema = pa.ipc.read_schema(pa.py_buffer(fmt.options))
113+
assert set(opt_schema.names) == {"delimiter", "null_string", "header", "on_exists"}

‎vgi/_test_fixtures/copy_to.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
# Copyright 2025, 2026 Query Farm LLC - https://query.farm
2+
3+
"""Fixture ``COPY ... TO`` format writer for VGI integration tests.
4+
5+
``ExampleLinesCopyToFunction`` registers the SQL format ``example_lines_out`` — a
6+
toy delimited-text writer, the symmetric counterpart of the ``example_lines``
7+
reader. It exercises the COPY-TO Sink+Combine path plus the option machinery: a
8+
required option (``null_string``), a defaulted option (``delimiter``), a BOOLEAN
9+
option (``header``), and an enum/``choices`` option (``on_exists``).
10+
11+
Shards are buffered in ``params.storage`` (``execution_id``-scoped) by ``write()``
12+
and concatenated to the destination by ``close()`` — the cross-process-safe
13+
pattern, so it works under pool rotation / HTTP.
14+
15+
Usage::
16+
17+
COPY (SELECT * FROM t) TO '/path/out.txt' (FORMAT 'acme.example_lines_out', null_string 'NA');
18+
"""
19+
20+
from __future__ import annotations
21+
22+
from dataclasses import dataclass
23+
from typing import TYPE_CHECKING, Annotated, ClassVar
24+
25+
import pyarrow as pa
26+
27+
from vgi.arguments import Arg
28+
from vgi.copy_to_function import CopyToFunction
29+
30+
if TYPE_CHECKING:
31+
from vgi.table_buffering_function import TableBufferingParams
32+
33+
__all__ = ["ExampleLinesCopyToFunction", "ExampleLinesOrderedCopyToFunction"]
34+
35+
_SHARD_NS = b"copy_to_shard"
36+
37+
38+
@dataclass(slots=True, frozen=True, kw_only=True)
39+
class ExampleLinesCopyToArgs:
40+
"""Options for the ``example_lines_out`` COPY format."""
41+
42+
null_string: Annotated[str, Arg("null_string", doc="Token written for SQL NULL")]
43+
delimiter: Annotated[str, Arg("delimiter", default=",", doc="Field separator")] = ","
44+
header: Annotated[bool, Arg("header", default=False, doc="Write a header row of column names")] = False
45+
on_exists: Annotated[
46+
str,
47+
Arg(
48+
"on_exists",
49+
default="overwrite",
50+
choices=["overwrite", "error"],
51+
doc="Behavior when the destination file already exists",
52+
),
53+
] = "overwrite"
54+
55+
56+
class ExampleLinesCopyToFunction(CopyToFunction[ExampleLinesCopyToArgs]):
57+
"""Toy delimited-text ``COPY ... TO`` writer (test fixture)."""
58+
59+
COPY_TO_FORMAT: ClassVar[str] = "example_lines_out"
60+
COPY_TO_COMMENT: ClassVar[str | None] = "Toy delimited-text writer for tests"
61+
62+
class Meta:
63+
name = "example_lines_writer"
64+
description = "Write the COPY source to a delimited text file"
65+
categories = ["copy", "test"]
66+
tags = {"category": "copy_to", "stability": "test"}
67+
68+
@classmethod
69+
def write(
70+
cls,
71+
*,
72+
batch: pa.RecordBatch,
73+
options: ExampleLinesCopyToArgs,
74+
file_path: str,
75+
params: TableBufferingParams[ExampleLinesCopyToArgs],
76+
) -> None:
77+
"""Buffer one input batch as an IPC blob in execution-scoped storage."""
78+
sink = pa.BufferOutputStream()
79+
with pa.ipc.new_stream(sink, batch.schema) as writer:
80+
writer.write_batch(batch)
81+
# state_append is atomic + race-safe across parallel sink threads/workers.
82+
params.storage.state_append(_SHARD_NS, b"", sink.getvalue().to_pybytes())
83+
84+
@classmethod
85+
def close(
86+
cls,
87+
*,
88+
options: ExampleLinesCopyToArgs,
89+
file_path: str,
90+
params: TableBufferingParams[ExampleLinesCopyToArgs],
91+
) -> int:
92+
"""Concatenate every shard and write the delimited destination file (once)."""
93+
import os
94+
95+
if options.on_exists == "error" and os.path.exists(file_path):
96+
raise FileExistsError(f"example_lines_out: destination already exists: {file_path}")
97+
98+
shards = params.storage.state_log_scan(_SHARD_NS, b"", after_id=-1)
99+
100+
def fmt(value: object) -> str:
101+
return options.null_string if value is None else str(value)
102+
103+
rows_written = 0
104+
with open(file_path, "w", encoding="utf-8") as fh:
105+
wrote_header = False
106+
for _log_id, blob in shards:
107+
table = pa.ipc.open_stream(blob).read_all()
108+
if options.header and not wrote_header:
109+
fh.write(options.delimiter.join(table.schema.names) + "\n")
110+
wrote_header = True
111+
for row in table.to_pylist():
112+
fh.write(options.delimiter.join(fmt(row[name]) for name in table.schema.names) + "\n")
113+
rows_written += 1
114+
# Empty COPY with header=true still emits the header row. We need the
115+
# source column names; they ride the bind's input_schema.
116+
if options.header and not wrote_header:
117+
assert params.init_call is not None
118+
in_schema = params.init_call.bind_call.input_schema
119+
if in_schema is not None:
120+
fh.write(options.delimiter.join(in_schema.names) + "\n")
121+
return rows_written
122+
123+
124+
class ExampleLinesOrderedCopyToFunction(ExampleLinesCopyToFunction):
125+
"""Ordered variant of :class:`ExampleLinesCopyToFunction`.
126+
127+
``Meta.ordered = True`` makes the extension use a single-threaded sink, so the
128+
worker receives every batch in source order and writes the file in order.
129+
"""
130+
131+
COPY_TO_FORMAT: ClassVar[str] = "example_lines_ordered_out"
132+
COPY_TO_COMMENT: ClassVar[str | None] = "Toy delimited-text writer (ordered, single-thread sink)"
133+
134+
class Meta:
135+
name = "example_lines_ordered_writer"
136+
description = "Write the COPY source to a delimited file, preserving source order"
137+
categories = ["copy", "test"]
138+
tags = {"category": "copy_to", "stability": "test"}
139+
sink_order_dependent = True # ordered COPY TO → single-thread sink

‎vgi/_test_fixtures/worker.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
SlowCancellableInOutFunction,
6161
)
6262
from vgi._test_fixtures.copy_from import ExampleLinesCopyFromFunction
63+
from vgi._test_fixtures.copy_to import ExampleLinesCopyToFunction, ExampleLinesOrderedCopyToFunction
6364
from vgi._test_fixtures.nest_tensor import NestTensorFunction, UnnestTensorFunction, UnnestTensorRowsFunction
6465
from vgi._test_fixtures.scalar import (
6566
AddValuesFunction,
@@ -354,6 +355,9 @@ def _build_enum_stats() -> dict[str, ColumnStatisticsInput]:
354355
SlowCancellableBufferingFunction,
355356
# CopyFromFunction - custom COPY ... FROM format reader
356357
ExampleLinesCopyFromFunction,
358+
# CopyToFunction - custom COPY ... TO format writer
359+
ExampleLinesCopyToFunction,
360+
ExampleLinesOrderedCopyToFunction,
357361
# TableFunctionGenerator - generate output without input
358362
ConstantColumnsFunction,
359363
SlowCancellableFunction,

‎vgi/catalog/catalog_interface.py‎

Lines changed: 47 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -705,18 +705,25 @@ class CopyFromFormatInfo(CatalogObject, ArrowSerializableDataclass):
705705
:attr:`FunctionInfo.arguments`); each field's metadata carries the
706706
option type / default / ``vgi_doc`` description. The reserved
707707
``file_path`` positional is excluded.
708-
direction: ``"from"`` — the only direction supported today. Reserved
709-
``"to"`` for a future ``COPY ... TO``; surfaced so the C++
710-
``vgi_copy_formats()`` diagnostic can split FROM vs TO.
708+
direction: ``"from"`` | ``"to"`` | ``"both"`` — which COPY direction(s)
709+
this format serves; surfaced so the C++ ``vgi_copy_formats()``
710+
diagnostic can split FROM vs TO and so registration wires the right
711+
callbacks.
711712
description: Intrinsic documentation from the handler's
712713
``Meta.description``.
714+
ordered: COPY ... TO only — when true the worker requires rows in source
715+
order, so the extension uses a single-threaded sink
716+
(``REGULAR_COPY_TO_FILE``) instead of the default parallel sharded
717+
write. Set via ``Meta.ordered = True`` on a ``CopyToFunction``.
718+
Ignored for readers.
713719
"""
714720

715721
format_name: str
716722
handler: str
717723
options: SerializedSchema
718724
direction: str = "from"
719725
description: str = ""
726+
ordered: bool = False
720727

721728

722729
@dataclass(frozen=True)
@@ -2863,17 +2870,20 @@ def copy_from_formats(
28632870
attach_opaque_data: AttachOpaqueData,
28642871
transaction_opaque_data: TransactionOpaqueData | None,
28652872
) -> list[CopyFromFormatInfo]:
2866-
"""Advertise every ``CopyFromFunction`` registered in this catalog.
2873+
"""Advertise every custom COPY format registered in this catalog.
28672874
28682875
Introspects the catalog's function list for
2869-
:class:`vgi.copy_from_function.CopyFromFunction` subclasses and converts
2870-
each into a :class:`CopyFromFormatInfo`. The option schema reuses the same
2871-
argument serialization as :meth:`_function_to_info`, so option types /
2872-
defaults / ``doc`` descriptions surface identically to
2873-
``vgi_function_arguments()``.
2876+
:class:`vgi.copy_from_function.CopyFromFunction` (``direction='from'``) and
2877+
:class:`vgi.copy_to_function.CopyToFunction` (``direction='to'``) subclasses
2878+
and converts each into a :class:`CopyFromFormatInfo`. (The RPC name is
2879+
historical; it returns *all* directions, distinguished by ``direction``.)
2880+
The option schema reuses the same argument serialization as
2881+
:meth:`_function_to_info`, so option types / defaults / ``doc`` descriptions
2882+
surface identically to ``vgi_function_arguments()``.
28742883
"""
28752884
from vgi.argument_spec import argument_specs_to_schema, extract_argument_specs
28762885
from vgi.copy_from_function import CopyFromFunction
2886+
from vgi.copy_to_function import CopyToFunction
28772887
from vgi.metadata import resolve_metadata
28782888

28792889
self._build_registries()
@@ -2883,23 +2893,45 @@ def copy_from_formats(
28832893
seen: set[str] = set()
28842894
for (_, _), func_classes in self._function_registry.items():
28852895
for func_cls in func_classes:
2886-
if not (isinstance(func_cls, type) and issubclass(func_cls, CopyFromFunction)):
2896+
if not isinstance(func_cls, type):
28872897
continue
2888-
fmt = getattr(func_cls, "COPY_FROM_FORMAT", None)
2889-
if not fmt or fmt in seen:
2898+
# A format is identified by COPY_FROM_FORMAT (readers) or
2899+
# COPY_TO_FORMAT (writers); direction comes from the matching
2900+
# *_DIRECTION marker.
2901+
is_to = False
2902+
if issubclass(func_cls, CopyFromFunction):
2903+
fmt = getattr(func_cls, "COPY_FROM_FORMAT", None)
2904+
direction = getattr(func_cls, "COPY_FROM_DIRECTION", "from")
2905+
comment = getattr(func_cls, "COPY_FROM_COMMENT", None)
2906+
elif issubclass(func_cls, CopyToFunction):
2907+
fmt = getattr(func_cls, "COPY_TO_FORMAT", None)
2908+
direction = getattr(func_cls, "COPY_TO_DIRECTION", "to")
2909+
comment = getattr(func_cls, "COPY_TO_COMMENT", None)
2910+
is_to = True
2911+
else:
28902912
continue
2891-
seen.add(fmt)
2913+
# Key by (format, direction): a from-format and a to-format may
2914+
# legitimately share a name; the extension scopes by alias anyway.
2915+
dedup_key = f"{direction}:{fmt}"
2916+
if not fmt or dedup_key in seen:
2917+
continue
2918+
seen.add(dedup_key)
28922919
meta = resolve_metadata(func_cls)
2920+
# COPY TO writers that require source order declare it via
2921+
# Meta.sink_order_dependent = True (single-thread sink ingest);
2922+
# the extension maps that to REGULAR_COPY_TO_FILE. Ignored for readers.
2923+
ordered = bool(is_to and meta.sink_order_dependent)
28932924
args_schema = argument_specs_to_schema(extract_argument_specs(func_cls))
28942925
formats.append(
28952926
CopyFromFormatInfo(
2896-
comment=getattr(func_cls, "COPY_FROM_COMMENT", None),
2927+
comment=comment,
28972928
tags=meta.tags,
28982929
format_name=fmt,
28992930
handler=meta.name,
29002931
options=SerializedSchema(args_schema.serialize().to_pybytes()),
2901-
direction=getattr(func_cls, "COPY_FROM_DIRECTION", "from"),
2932+
direction=direction,
29022933
description=meta.description or "",
2934+
ordered=ordered,
29032935
)
29042936
)
29052937
return formats

0 commit comments

Comments
 (0)