# Batch UDTF Reference

**Status:** Feature-complete (beta) — 2026-03-06

This document is the reference for the implemented batch UDTF system. For
unimplemented future work, see [designs/udtf_future.md](designs/udtf_future.md).
For per-row 1:N chunking UDTFs (not yet implemented), see
[designs/udtf_chunking_design.md](designs/udtf_chunking_design.md).

---

## Problem

Geneva's UDF model operates **row-at-a-time**, producing exactly one output value
per input row. This 1:1 constraint makes it impossible to express many common
data preparation workflows:

| Workflow | Input | Output | Cardinality |
|----------|-------|--------|-------------|
| Deduplication | N rows | M rows (M <= N) | N:M |
| Clustering/grouping | N rows | K cluster rows | N:K |
| Aggregation | N rows | 1 summary row | N:1 |
| Cross-row join/merge | N rows | M rows | N:M |

A **User-Defined Table Function (UDTF)** lifts this restriction. A UDTF reads
from a source table and **yields** output batches with an arbitrary output schema
and row count.

---

## UDTF Taxonomy

| Type | Scope | Refresh | Parallelism | Status |
|------|-------|---------|-------------|--------|
| **Batch** (this doc) | Whole-table | Full | Partition-parallel (via `partition_by` or `partition_by_indexed_column`) | Implemented |
| Scalar/Chunking | Per-row | Incremental | Fragment-parallel | Not implemented ([design](designs/udtf_chunking_design.md)) |

**Batch UDTFs** see the entire source table (or a filtered subset) at once. They
always do a full refresh — re-read entire source, re-produce all output.

**Scalar/Chunking UDTFs** process one row at a time (1:N expansion). They can
refresh incrementally. See `designs/udtf_chunking_design.md`.

---

## API

### Decorator

```python
import geneva
import pyarrow as pa
from collections.abc import Iterator

@geneva.udtf(
    output_schema=pa.schema([
        pa.field("row_id", pa.int64()),
        pa.field("cluster_id", pa.int64()),
        pa.field("duplicate_row_ids", pa.list_(pa.int64())),
    ]),
    input_columns=["row_id", "phash"],   # optional: restrict source columns
    partition_by_indexed_column="phash",       # or partition_by="col"
    num_cpus=2.0,                         # Ray resource request
    num_gpus=0.0,
    memory=4 * 1024**3,                   # 4 GiB
)
class PHashDedupe:
    def __init__(self, threshold: int = 4):
        self.threshold = threshold

    def __call__(self, source: geneva.GenevaQueryBuilder) -> Iterator[pa.RecordBatch]:
        data: pa.Table = source.to_arrow().select(["row_id", "phash"])
        clusters = self._cluster(data)
        yield pa.RecordBatch.from_pydict({...})
```

Both class-based and function-based UDTFs are supported. Both must return an
`Iterator[pa.RecordBatch]`.

```python
# Function-based
@geneva.udtf(output_schema=pa.schema([...]))
def compute_scores(source: geneva.GenevaQueryBuilder, model_name: str = "default") -> Iterator[pa.RecordBatch]:
    for batch in source.to_batches():
        yield process(batch, model_name)
```

### UDTF Data Model

```python
@attrs.define
class UDTF:
    func: Callable              # The user's class or function
    output_schema: pa.Schema    # Required: output table schema
    name: str                   # Auto-derived from func if empty
    input_columns: list[str] | None  # Restrict source columns (None = all)
    num_cpus: float             # Ray resource: CPUs
    num_gpus: float             # Ray resource: GPUs
    memory: int | None          # Ray resource: bytes
    version: str                # Cache key (auto-derived from func hash)
    partition_by_indexed_column: str | None  # Index-based partitioning (mutually exclusive with partition_by)
    error_handling: ErrorHandlingConfig | None
```

### Creating a UDTF View

```python
conn = geneva.connect("/data/mydb")
images = conn.open_table("images")

deduped = conn.create_udtf_view(
    "deduped_images",
    source=images.search(None).select(["row_id", "phash"]),
    udtf=PHashDedupe(threshold=4),
)

# First refresh populates the view
deduped.refresh()
```

### Cluster and Context

UDTFs use Geneva's existing cluster infrastructure — no separate cluster config:

```python
# KubeRay cluster
with conn.context(cluster="my-cluster", manifest="my-manifest"):
    mv.refresh()

# Local Ray
with conn.local_ray_context():
    mv.refresh()
```

---

## Execution Model

### Single-worker mode (no `partition_by`)

The UDTF runs as a single Ray task. Use for global operations needing cross-row
visibility across the entire dataset.

### Partition-parallel mode (`partition_by="column"`)

The framework groups source data by the partition key and dispatches each
partition as an independent Ray task via `ActorPool.map_unordered()`. Results
are concatenated.

```python
@geneva.udtf(
    output_schema=...,
    partition_by="partition_id",
)
class PartitionedUDTF:
    def __call__(self, source_partition) -> Iterator[pa.RecordBatch]:
        data = source_partition.to_arrow()
        yield process(data)
```

### Index-based partition mode (`partition_by_indexed_column`)

Instead of partitioning by a materialized column, the framework reads partition
assignments directly from an existing IVF vector index (IVF_FLAT, IVF_PQ,
IVF_HNSW_SQ, etc.) at refresh time. This avoids materializing a `partition_id`
column and keeps partitions synchronized with the index.

```python
@geneva.udtf(
    output_schema=...,
    partition_by_indexed_column="phash",
)
class IndexPartitionedUDTF:
    def __call__(self, source_partition) -> Iterator[pa.RecordBatch]:
        data = source_partition.to_arrow()
        yield process(data)
```

**How it works:**

1. At refresh time, `_build_index_partition_work_items()` opens a
   `VectorIndexReader` on the IVF index and enumerates non-empty
   partitions.
2. Each partition becomes a work item carrying an `_IndexPartitionInfo`
   (partition ordinal, index name, column) instead of a SQL filter string.
3. On the remote actor, `_build_index_source()` re-reads the partition's row
   IDs from the index and uses `take_row_ids()` for direct row access.

**Key differences from `partition_by`:**

| | `partition_by` | `partition_by_indexed_column` |
|---|---|---|
| Partition source | Column values (SQL filter) | IVF index partitions (row ID take) |
| Requires column | Yes — column must exist | No — reads from index metadata |
| Partition count | Number of distinct values | Number of non-empty index partitions |
| Sync with index | Manual (`assign_partitions_from_index`) | Automatic — always reads latest index |

`partition_by` and `partition_by_indexed_column` are **mutually exclusive** — setting
both raises `ValueError`.

### Yielding Batches

The UDTF yields one or more `pa.RecordBatch` or `pa.Table` objects. Each batch
must conform to `output_schema`. The framework validates, writes, and optionally
checkpoints each batch.

```python
# Streaming style — yield per source batch (memory-efficient)
def __call__(self, source: GenevaQueryBuilder) -> Iterator[pa.RecordBatch]:
    for batch in source.to_batches(batch_size=1024):
        yield transform(batch)

# Bulk style — load all, compute, yield once
def __call__(self, source: GenevaQueryBuilder) -> Iterator[pa.RecordBatch]:
    all_data = source.to_arrow()
    result = expensive_cross_row_computation(all_data)
    yield result.to_batches()[0]
```

---

## Production Infrastructure

### Architecture

```
_refresh_udtf_matview (revised)
├── validate_admission(udtf_spec)          # admission control
├── JobTracker.remote(job_id, ...)         # metrics actor
├── CheckpointStore.from_uri(...)          # resume support
├── work_items = build_work_items(partition_by | partition_by_indexed_column | none)
├── ActorPool(UDTFProcessorActor.factory, num_actors=concurrency)
│   └── map_unordered(process_partition, partition_keys)
│       └── yields WriteBatch results as partitions complete
├── Fragment-per-partition writes          # streaming staged writes
├── JobTracker.mark_done("refresh")
└── cleanup checkpoint store
```

### UDTFProcessorActor

Ray actor that plugs into `ActorPool` as a worker. Each actor processes one
partition at a time:

```
UDTFProcessorActor
├── __init__(udtf_pickle, source_table_ref, query_spec, checkpoint_store)
│   └── Deserializes UDTF, opens source table once
├── process_partition(partition_key, partition_filter) → List[WriteBatch]
│   ├── Builds filtered query from partition_filter
│   ├── Calls udtf.execute(query_builder)
│   ├── For each yielded batch:
│   │   ├── Writes batch to checkpoint store (key: partition/batch_idx)
│   │   ├── Reports progress to JobTracker
│   │   └── Yields WriteBatch(data, partition_key, batch_idx)
│   └── Returns collected WriteBatch list
└── on_error handling per on_error config
```

### Capability Summary

| Capability | Implementation |
|---|---|
| **Task submission** | `ActorPool.map_unordered()` with `UDTFProcessorActor`, fault-tolerant actor restart |
| **Parallelism** | Partition-parallel via `partition_by` or `partition_by_indexed_column`; single task if neither set |
| **Writing** | Fragment-per-partition — each partition is a separate Lance fragment, O(largest partition) memory |
| **Checkpointing** | `CheckpointStore` — per-batch skip on resume, per-partition `__done__` markers |
| **Admission control** | `validate_udtf_admission()` — CPU/GPU/memory check before dispatch |
| **Metrics** | `JobTracker` actor — partition progress tracking |
| **Failure handling** | Partition-level fail/retry/skip; `ErrorStore` integration |
| **Version-aware refresh** | O(1) source version check, early exit if unchanged |
| **Namespace support** | `storage_options_provider` threaded through all writes |

### Error Handling

The `@udtf` decorator accepts `on_error` (stored as `UDTF.error_handling`).
Error handling operates at **partition granularity** — the unit of work is the
entire `execute()` call for a partition.

| Mode | Behavior | Use case |
|---|---|---|
| **Fail partition** (default) | Exception kills the partition, refresh fails | Correctness-critical UDTFs |
| **Retry partition** | Retry entire `execute()` call with configurable backoff | Transient failures (network, OOM) |
| **Skip partition** | Log error to `ErrorStore`, continue with remaining partitions | Best-effort / tolerant workloads |

Key differences from MV pipeline error handling:
- No `SkipRowsStrategy` — UDTFs yield whole batches, not row-by-row
- Retry granularity is the full partition `execute()` call
- `ErrorStore` uses `partition_key` as location identifier (not `fragment_id`)

### Checkpointing

Checkpoint keys are scoped by `udtf_name:udtf_version/src-{source_version}/{partition_col}={partition_value}`.

- Each yielded batch is written to `CheckpointStore` before reporting completion
- On resume, completed batches are skipped; `__done__` markers skip entire partitions
- Version-aware invalidation: checkpoint keys include source version, so stale
  checkpoints are ignored when the source changes

**Requirement:** Checkpointed UDTFs must be deterministic — the same input must
yield the same batch sequence for resume to work correctly.

### Version-Aware Refresh

On refresh start, compare the source table's version against the version stored
in destination metadata (`MATVIEW_META_BASE_VERSION`). If unchanged, skip refresh
entirely (O(1) check).

---

## Key Design Decisions

### Partition-based concurrency, not fragment-based

The MV pipeline parallelizes across *source fragments*. UDTFs parallelize across
*partition keys* specified by the user via `partition_by` or `partition_by_indexed_column`.
`ActorPool` dispatches partition work items, not fragment IDs.

### Full overwrite, not incremental

UDTFs are black-box transformations. There is no way to determine which output
rows correspond to which input rows. Every refresh overwrites the entire output
table.

### Per-batch checkpoints, not per-fragment

A single partition may yield many batches over a long execution. Checkpointing
per-yielded-batch allows resume at batch granularity within a partition.

### MV-only (no standalone `apply_udtf()`)

Following industry practice (Snowflake, Databricks, Pixeltable), UDTFs are
always part of a persistent definition executed via `create_udtf_view()` +
`refresh()`.

### Dedicated factory method (`create_udtf_view`)

UDTF-backed views differ from UDF-backed MVs in every dimension (cardinality,
refresh semantics, execution model, output schema). A dedicated factory keeps
the concepts clean.

---

## Files

| File | Description |
|---|---|
| `src/geneva/table.py` | `UDTFProcessorActor`, `_build_column_partition_work_items`, `_build_index_partition_work_items`, `_refresh_udtf_matview` with ActorPool, checkpointing, error handling |
| `src/geneva/partitioning.py` | `create_ivf_flat_index()`, `assign_partitions_from_index()` |
| `src/geneva/runners/ray/admission.py` | `validate_udtf_admission()`, `calculate_udtf_job_resources()` |
| `src/geneva/checkpoint.py` | `format_udtf_checkpoint_prefix()`, `format_udtf_partition_prefix()` |
| `tests/test_udtf.py` | Decorator, execution, serialization, partitioning (column + index work items), checkpointing, config validation |
| `tests/test_udtf_actorpool.py` | 11 tests: ActorPool, admission, on_error, ErrorStore, fragment-per-partition |
| `tests/test_udtf_checkpoint.py` | Checkpoint key helpers and read/write tests |

---

## Worked Examples

### Image Deduplication (N:M)

```python
@geneva.udtf(
    output_schema=pa.schema([
        pa.field("row_id", pa.int64()),
        pa.field("cluster_id", pa.int64()),
        pa.field("is_duplicate", pa.bool_()),
    ]),
    input_columns=["row_id", "phash"],
    num_cpus=4,
    memory=8 * 1024**3,
)
class PHashDedupe:
    def __init__(self, threshold: int = 4):
        self.threshold = threshold

    def __call__(self, source: geneva.GenevaQueryBuilder) -> Iterator[pa.RecordBatch]:
        data = source.to_arrow().select(["row_id", "phash"])
        clusters = self._cluster(data)
        yield pa.RecordBatch.from_pydict({...})

deduped = conn.create_udtf_view(
    "deduped_images",
    source=images.search(None).select(["row_id", "phash"]),
    udtf=PHashDedupe(threshold=4),
)
deduped.refresh()
```

### Aggregation (N:1)

```python
@geneva.udtf(
    output_schema=pa.schema([
        pa.field("label", pa.string()),
        pa.field("count", pa.int64()),
        pa.field("mean_score", pa.float64()),
    ]),
)
class GroupStats:
    def __init__(self, group_by: str = "label"):
        self.group_by = group_by

    def __call__(self, source: GenevaQueryBuilder) -> Iterator[pa.RecordBatch]:
        df = source.to_pandas()
        agg = df.groupby(self.group_by).agg(
            count=("label", "size"),
            mean_score=("score", "mean"),
        ).reset_index()
        yield pa.RecordBatch.from_pandas(agg)
```

### K-Means Clustering (N:K)

```python
@geneva.udtf(
    output_schema=pa.schema([
        pa.field("row_id", pa.int64()),
        pa.field("cluster_id", pa.int64()),
        pa.field("distance_to_centroid", pa.float64()),
    ]),
    num_cpus=4,
    memory=16 * 1024**3,
)
class KMeansClustering:
    def __init__(self, k: int = 100):
        self.k = k

    def __call__(self, source: GenevaQueryBuilder) -> Iterator[pa.RecordBatch]:
        import numpy as np
        embeddings = source.to_arrow().select(["row_id", "embedding"])
        row_ids = embeddings.column("row_id").to_pylist()
        vectors = np.stack(embeddings.column("embedding").to_pylist())
        centroids, assignments, distances = self._fit(vectors)

        chunk_size = 10_000
        for start in range(0, len(row_ids), chunk_size):
            end = min(start + chunk_size, len(row_ids))
            yield pa.RecordBatch.from_pydict({
                "row_id": row_ids[start:end],
                "cluster_id": assignments[start:end].tolist(),
                "distance_to_centroid": distances[start:end].tolist(),
            })
```

### Partitioned Edge Detection

```python
@geneva.udtf(
    output_schema=pa.schema([
        ("row_id_a", pa.string()),
        ("row_id_b", pa.string()),
        ("hamming_dist", pa.int32()),
    ]),
    input_columns=["row_id", "phash"],
    partition_by="partition_id",
    num_cpus=2,
)
class EdgeDetection:
    def __init__(self, threshold: int = 4):
        self.threshold = threshold

    def __call__(self, source: GenevaQueryBuilder) -> Iterator[pa.RecordBatch]:
        data = source.to_arrow()
        row_ids = data["row_id"].to_pylist()
        phashes = data["phash"].to_pylist()
        edges = self._pairwise_compare(row_ids, phashes, self.threshold)
        if edges:
            yield pa.RecordBatch.from_pydict({
                "row_id_a": [e[0] for e in edges],
                "row_id_b": [e[1] for e in edges],
                "hamming_dist": [e[2] for e in edges],
            })
```

### Partitioned Edge Detection (index-based)

Uses `partition_by_indexed_column` to partition directly from an existing IVF vector
index, avoiding the need to materialize a `partition_id` column:

```python
from geneva.partitioning import create_ivf_flat_index

# 1. Build an IVF index (IVF_FLAT shown here; IVF_PQ, IVF_HNSW_SQ, etc. also work)
create_ivf_flat_index(tbl, "phash", k=16)

# 2. Define the UDTF with partition_by_indexed_column
@geneva.udtf(
    output_schema=pa.schema([
        ("row_id_a", pa.string()),
        ("row_id_b", pa.string()),
        ("hamming_dist", pa.int32()),
    ]),
    input_columns=["row_id", "phash"],
    partition_by_indexed_column="phash",
    num_cpus=2,
)
class EdgeDetection:
    def __init__(self, threshold: int = 4):
        self.threshold = threshold

    def __call__(self, source: GenevaQueryBuilder) -> Iterator[pa.RecordBatch]:
        data = source.to_arrow()
        row_ids = data["row_id"].to_pylist()
        phashes = data["phash"].to_pylist()
        edges = self._pairwise_compare(row_ids, phashes, self.threshold)
        if edges:
            yield pa.RecordBatch.from_pydict({
                "row_id_a": [e[0] for e in edges],
                "row_id_b": [e[1] for e in edges],
                "hamming_dist": [e[2] for e in edges],
            })
```
