Skip to content
Query.Farm
Talk with Us

Add a custom COPY format

How a VGI catalog registers its own COPY formats, so users can COPY … FROM a source your worker knows how to parse and COPY … TO a destination it knows how to write — a proprietary format, a remote API, or a custom sink.

Prerequisites

Section titled “Prerequisites”
  • A catalog (see Expose a catalog) — COPY formats are catalog-level, so they need an ATTACH.
  • vgi-python 0.8.8+ for COPY … FROM and COPY … TO, 0.9.0+ for credential forwarding (on_secrets).

The model

Section titled “The model”

Both directions reuse machinery you already have, which is why they impose so few new rules:

  • A CopyFromFunction is an ordinary producer-mode table function. It reuses the whole table function bind/init/scan path.
  • A CopyToFunction is a buffered (Sink + Combine) function with no Source phase. It reuses the whole table_buffering_process / table_buffering_combine machinery.

What makes either one a COPY format is that it names a SQL FORMAT identifier and the catalog advertises it. Register the class in the catalog’s function list like any other function; the declarative Catalog introspects the list for both base classes and advertises what it finds, so there is normally nothing else to wire. The extension registers one DuckDB CopyFunction per entry at ATTACH time.

The COPY statement’s path is not an option — it arrives on the bind context. Everything else you declare as ordinary Arg-annotated arguments, and their doc strings become the option descriptions surfaced by vgi_copy_formats().

Reading: COPY … FROM

Section titled “Reading: COPY … FROM”

Set COPY_FROM_FORMAT, declare your options, and implement read:

from vgi.copy_from_function import CopyFromFunction

class ReadTsvLite(CopyFromFunction[ReadOptions]):
  COPY_FROM_FORMAT = "tsvlite"
  COPY_FROM_COMMENT = "Tab-separated, one record per line, no quoting"

  @classmethod
  def read(cls, *, path, options, expected_schema, params, out) -> None:
      ...              # parse `path`, build batches matching `expected_schema`
      out.emit(batch)
      out.finish()

Here is the whole worker — a TSV reader and writer in one catalog:

copy_format_worker.py 
# /// script
# requires-python = ">=3.13"
# dependencies = ["vgi-python"]
# ///
"""Custom ``COPY ... FROM`` and ``COPY ... TO`` formats, in one catalog.

A VGI catalog can register its own ``COPY`` formats, so users read and write a
format DuckDB has never heard of. Both directions reuse machinery that already
exists, which is why they add so few new rules:

- a :class:`CopyFromFunction` is an ordinary producer-mode **table function**;
- a :class:`CopyToFunction` is a **buffered** function with no Source phase —
  ``write`` is the sink (per batch, parallel), ``close`` is the combine (once).

The format here is tab-separated values, one record per line, no header, no
quoting. Deliberately trivial — the point is the wiring, not the parser.

Note the reader and writer declare **different** format names (``tsvlite`` and
``tsvlite_out``). A reader and a writer that share one name look fine on the
Python side — the catalog advertises both, keyed by (direction, format) — but the
extension registers COPY functions by name alone, so the writer is silently
dropped and ``COPY ... TO`` fails with "COPY TO is not supported for FORMAT".
Give each direction its own name.

Two rules the framework enforces, and this file demonstrates:

1. **The reader must emit ``expected_schema`` exactly.** DuckDB inserts no cast
   between the scan and the INSERT, so the target table's columns are the output
   schema — that is why ``on_bind`` is ``@final`` on a COPY-FROM function.
2. **The writer must finish inside ``close``.** There is no finalize phase, and
   ``write`` and ``close`` may run in different processes, so shards go through
   ``params.storage`` scoped by ``execution_id`` — never on ``self``.

The SQL ``FORMAT`` name is qualified by the catalog alias you attached as, NOT
the bare ``COPY_*_FORMAT`` string — the extension namespaces formats per attach,
so two workers may both call theirs ``tsvlite``:

    ATTACH 'tsv' (TYPE vgi, LOCATION 'uv run copy_format_worker.py');
    CREATE TABLE people (name VARCHAR, age BIGINT);
    COPY people FROM 'people.tsv' (FORMAT 'tsv.tsvlite');
    COPY (SELECT * FROM people) TO 'out.tsv' (FORMAT 'tsv.tsvlite_out', header true);
"""

from dataclasses import dataclass
from typing import Annotated

import pyarrow as pa

from vgi import Arg, Worker
from vgi.catalog import Catalog, Schema
from vgi.copy_from_function import CopyFromFunction
from vgi.copy_to_function import CopyToFunction
from vgi.table_buffering_function import TableBufferingParams
from vgi.table_function import OutputCollector, ProcessParams


@dataclass(slots=True, frozen=True, kw_only=True)
class ReadOptions:
    """Options accepted by ``COPY ... FROM ... (FORMAT tsvlite, ...)``.

    These are ordinary ``Arg``-annotated arguments — the source path is NOT one
    of them; it arrives on the bind. Each ``doc`` becomes the option's
    description in ``vgi_copy_formats()``.
    """

    skip_rows: Annotated[int, Arg("skip_rows", doc="Leading lines to discard", default=0)] = 0


@dataclass(slots=True, frozen=True, kw_only=True)
class WriteOptions:
    """Options accepted by ``COPY ... TO ... (FORMAT 'tsv.tsvlite_out', ...)``."""

    header: Annotated[bool, Arg("header", doc="Write a header line of column names", default=False)] = False


class ReadTsvLite(CopyFromFunction[ReadOptions]):
    """Read a tab-separated file into the COPY target table."""

    COPY_FROM_FORMAT = "tsvlite"
    COPY_FROM_COMMENT = "Tab-separated, one record per line, no quoting"

    @classmethod
    def read(
        cls,
        *,
        path: str,
        options: ReadOptions,
        expected_schema: pa.Schema,
        params: ProcessParams[ReadOptions],
        out: OutputCollector,
    ) -> None:
        """Parse ``path`` and emit batches matching ``expected_schema`` exactly."""
        with open(path, encoding="utf-8") as handle:
            lines = [line.rstrip("\n") for line in handle if line.strip()]
        lines = lines[options.skip_rows :]

        # Split into columns positionally, then let Arrow cast each column to the
        # type the target table declared. Emitting a mismatched type or arity is
        # rejected by the extension at COPY bind, not silently coerced.
        columns: list[list[str | None]] = [[] for _ in expected_schema]
        for line in lines:
            cells = line.split("\t")
            for index in range(len(expected_schema)):
                value = cells[index] if index < len(cells) else ""
                columns[index].append(value if value != "" else None)

        out.emit(
            pa.RecordBatch.from_arrays(
                [
                    pa.array(col, type=pa.string()).cast(field.type)
                    for col, field in zip(columns, expected_schema, strict=True)
                ],
                schema=expected_schema,
            )
        )
        out.finish()


class WriteTsvLite(CopyToFunction[WriteOptions]):
    """Write query results out as a tab-separated file."""

    COPY_TO_FORMAT = "tsvlite_out"
    COPY_TO_COMMENT = "Tab-separated, one record per line, no quoting"

    @classmethod
    def write(
        cls,
        *,
        batch: pa.RecordBatch,
        options: WriteOptions,
        file_path: str,
        params: TableBufferingParams[WriteOptions],
    ) -> None:
        """Sink: stash this batch as a shard. Runs per batch, possibly in parallel."""
        rows = ["\t".join("" if value is None else str(value) for value in row.values()) for row in batch.to_pylist()]
        # execution_id-scoped, because close() may run in a different process.
        params.storage.state_append(b"shards", b"", "\n".join(rows).encode("utf-8"))

    @classmethod
    def close(
        cls,
        *,
        options: WriteOptions,
        file_path: str,
        params: TableBufferingParams[WriteOptions],
    ) -> int:
        """Combine: read every shard back and write the file, once. Returns rows written."""
        shards = params.storage.state_log_scan(b"shards", b"")
        body = [chunk.decode("utf-8") for _key, chunk in shards if chunk]

        written = 0
        with open(file_path, "w", encoding="utf-8") as handle:
            if options.header:
                schema = params.init_call.bind_call.input_schema
                handle.write("\t".join(schema.names) + "\n")
            for chunk in body:
                handle.write(chunk + "\n")
                written += len(chunk.split("\n"))
        # Called even for an empty COPY — the file is created either way.
        return written


class TsvWorker(Worker):
    """A worker whose catalog advertises both COPY directions.

    The declarative ``Catalog`` introspects its function list for
    ``CopyFromFunction`` / ``CopyToFunction`` subclasses and advertises what it
    finds, so registering them here is all the wiring there is.
    """

    catalog = Catalog(
        name="tsv",
        schemas=[Schema(name="main", functions=[ReadTsvLite, WriteTsvLite])],
    )


if __name__ == "__main__":
    TsvWorker().run()

ATTACH 'tsv' (TYPE vgi, LOCATION 'uv run copy_format_worker.py');
CREATE TABLE people (name VARCHAR, age BIGINT);
COPY people FROM 'people.tsv' (FORMAT 'tsv.tsvlite', skip_rows 1);
The FORMAT name is qualified by the attach alias

COPY_FROM_FORMAT = "tsvlite" is not what users type. The extension namespaces formats per attach, so the SQL name is '<attach-alias>.<format>' — here 'tsv.tsvlite', because the catalog was attached as tsv. The bare name raises Catalog Error: Copy Function with name tsvlite does not exist!, helpfully suggesting the qualified one. Two workers may therefore both call their format tsvlite without colliding.

Emit expected_schema exactly

DuckDB inserts no cast between the scan and the INSERT. Your batches must match expected_schema exactly in both type and arity — the extension rejects a mismatch at COPY bind. The target table’s schema is what defines it; you don’t get to choose the output schema, which is why on_bind is @final here.

The file path and target schema arrive through CopyFromContext on the bind (params.bind_call.copy_from, or params.init_call.bind_call.copy_from after init).

Writing: COPY … TO

Section titled “Writing: COPY … TO”

Set COPY_TO_FORMAT and implement two methods:

  • write — called once per input batch, fanned out across DuckDB’s sink threads and per-thread workers.
  • close — called exactly once, on the coordinator worker, driven by DuckDB’s once-only copy_to_finalize. This is where the terminal write, flush and close happen.
from vgi.copy_to_function import CopyToFunction

class WriteTsvLite(CopyToFunction[WriteOptions]):
  COPY_TO_FORMAT = "tsvlite_out"   # NOT "tsvlite" — see the warning below

  @classmethod
  def write(cls, *, batch, options, file_path, params) -> None:
      # Per batch, possibly parallel, possibly on another process.
      params.storage.state_append(b"shards", b"", serialize(batch))

  @classmethod
  def close(cls, *, options, file_path, params) -> int:
      # Once, on the coordinator. Read the shards back and finish the destination.
      with open(file_path, "w") as dest:
          for _key, shard in params.storage.state_log_scan(b"shards", b""):
              dest.write(deserialize(shard))
      return rows_written        # informational; DuckDB reports its own count
COPY (SELECT * FROM people) TO 'out.tsv' (FORMAT 'tsv.tsvlite_out', header true);
Give each direction its own format name

A reader and a writer that share one COPY_*_FORMAT string look fine from Python — the catalog advertises both, keyed by (direction, format), and the source even notes that sharing should be allowed. The extension registers COPY functions by name alone, so the writer is silently dropped: vgi_copy_formats() shows only the from row, and COPY … TO fails with “Not implemented Error: COPY TO is not supported for FORMAT …”. There is no warning at attach time.

Name them apart — tsvlite and tsvlite_out, as the worker above does, matching the upstream fixtures’ example_lines / example_lines_out.

close() returns the row count and is called even for an empty COPY, so a header-only or zero-byte file still gets created rather than silently skipped.

There is no finalize — and write/close may be different processes

The destination must be fully written and closed inside close(); a writer that forgets leaves a silent partial file. And because write() and close() can run in different worker processes (pool rotation, HTTP), any shard state close() needs must live in cross-process storage scoped by params.execution_id — params.storage is the canonical choice — or go to a destination that tolerates concurrent writers (object-store multipart, an append API). Buffering on self or a module global breaks silently under rotation, exactly as it does for a buffering function.

Reach the destination path with copy_to_path(params) at write/close time. The source schema is the input schema (params.init_call.bind_call.input_schema).

Credentials for cloud destinations

Section titled “Credentials for cloud destinations”

A writer usually can’t reach s3:// without credentials, but on_bind is @final on a CopyToFunction — there’s no output schema for you to compute. The seam is on_secrets, called from on_bind for exactly this purpose:

@classmethod
def on_secrets(cls, params) -> None:
  # Scope by the destination path so DuckDB resolves the longest-prefix-matching secret.
  params.secrets.get("s3", scope=params.bind_call.copy_to.file_path, required=True)

The framework issues a two-phase bind retry to resolve every requested secret from the caller’s secret store, then surfaces the resolved values on params.secrets (a ResolvedSecrets) at write / close time. A requested secret that doesn’t exist resolves to “not found” rather than an error — pass required=True to make a missing secret fail the bind instead.

The path is available at on_secrets time

on_secrets runs during bind, where there is no init_call yet — so read the destination from params.bind_call.copy_to.file_path, not copy_to_path() (which is a write/close-time helper). The default on_secrets requests nothing, so writers that never touched credentials are unaffected.

Discovery

Section titled “Discovery”

vgi_copy_formats() lists everything the attached catalogs advertise — both directions, each tagged with its direction — with the option schema, types, defaults and doc descriptions taken from your Arg annotations. COPY_FROM_COMMENT / COPY_TO_COMMENT add free text alongside.

Next steps

Section titled “Next steps”