omnigraph/.context/experiments/custom-operator.md

# Experiment 1.3 — Custom UserDefinedLogicalNode + ExecutionPlan e2e

**Ticket:** MR-925 §1.3 (validates MR-737 §5.3, §5.10).
**Prototype:** `validation-prototypes/custom-operator/`.
**Substrate pin:** DataFusion 52.5 (matched to omnigraph workspace).
**Date:** 2026-05-12.

---

## Hypothesis

We can ship a graph-specific operator (e.g. `NeighborExpand` for
`MATCH (a)-[]->(b)`) as a third-party DataFusion node, with:

  - a `UserDefinedLogicalNodeCore` implementation that lives in our crate,
  - a paired `ExecutionPlan` that does the actual work,
  - an `ExtensionPlanner` that lowers the logical → physical,
  - integration via `SessionStateBuilder::with_query_planner`,
  - normal optimizer + physical-planner composition with built-in DF ops
    (Filter, Aggregate, etc.),
  - working `BaselineMetrics` (so the operator shows up in EXPLAIN ANALYZE).

This is one of the two non-negotiable prototypes the ticket requires (§1.3).

## Method

The prototype defines a `NeighborExpand` operator with these semantics:

> Input schema: `(src_id: UInt64, _neighbors: List<UInt64>)`
> Output schema: `(src_id: UInt64, edge_type: Utf8, dst_id: UInt64)`
> Semantics: flatten each input row's `_neighbors` list, emitting one
> output row per `(src_id, edge_type, dst_id)` triple.

Five probes are run on the same plan tree:

| Probe | What is exercised |
|-------|-------------------|
| **E1** Round-trip      | Build `LogicalPlan::Extension(NeighborExpandNode { input: scan, edge_type: "FOLLOWS" })`, plan through `DefaultPhysicalPlanner` + our `ExtensionPlanner`, execute, count rows. |
| **E2** EXPLAIN         | `LogicalPlan::display_indent()` and `displayable(physical).indent(true)` show our node names. |
| **E3** Predicate guard | Default `prevent_predicate_push_down_columns()` blocks the optimizer from pushing predicates *below* our node (it would change semantics — `WHERE dst_id = 7` only makes sense *after* the expand). |
| **E4** Composition     | Wrap output as a `MemTable` and run `SELECT count(*) ... GROUP BY edge_type WHERE dst_id > 2` over it; verify the result composes with DF's `FilterExec` + `AggregateExec`. |
| **E5** Metrics         | After execute, the operator's `metrics()` returns a `MetricsSet` containing `OutputRows`, `OutputBatches`, `OutputBytes`, `ElapsedCompute`, `StartTimestamp`, `EndTimestamp`. |

Run output (release):

```
Logical plan:
NeighborExpand: edge_type=FOLLOWS
  TableScan: edges_factored projection=[src_id, _neighbors]

Physical plan:
NeighborExpandExec: edge_type=FOLLOWS
  DataSourceExec: partitions=1, partition_sizes=[1]

[E1] Total flattened rows = 7
[E1] PASS: row count matches expected 7
[E1] First batch (src,dst) pairs: [(10, 1), (10, 2), (20, 3), (40, 7), (40, 8), (40, 9), (40, 10)]
[E4] PASS: Filter(dst>2)+Aggregate(count(*)) over expand = 5
[E5] Physical plan metrics: ... OutputRows(Count { value: 7 }) ...
[E3] prevent_predicate_push_down_columns = {"src_id", "dst_id", "edge_type"}
[E3] PASS: predicate push-down conservatively blocks all output cols (the default)

All probes passed.
```

## Findings

### F1. The integration surface is clean. ✅

The full integration footprint, in lines of code, is **about 250 lines**
for both a logical node + a physical operator. The work breakdown is:

  - `impl UserDefinedLogicalNodeCore for NeighborExpandNode` — six required
    methods plus boilerplate (Hash, Eq, PartialOrd). No internal DF types
    leak through.
  - `impl ExecutionPlan for NeighborExpandExec` — `properties`, `children`,
    `with_new_children`, `execute`, `metrics`, `statistics`,
    `partition_statistics`. Public enums only (`Boundedness::Bounded`,
    `EmissionType::Incremental`).
  - `impl ExtensionPlanner for NeighborExpandPlanner` — one `plan_extension`
    method that downcasts via `node.as_any().downcast_ref::<Our>()`.
  - `SessionStateBuilder::with_query_planner(Arc::new(MyPlanner))` to glue
    everything together.

There are no `pub(crate)` blockers and no `internal::` modules required.

### F2. Predicate push-down is opt-in. ✅

`UserDefinedLogicalNodeCore::prevent_predicate_push_down_columns()` defaults
to "block all output columns". For `NeighborExpand`, this is the *right*
default — predicates on the post-flatten `dst_id` cannot be pushed below
the expand without changing semantics. If we later want pushdown on
`src_id` (it's stable across the expand), we override the method to
return the schema minus `src_id`. This is a one-line opt-in.

### F3. EXPLAIN integration is automatic. ✅

`fmt_for_explain` on the logical node and `DisplayAs` on the physical
operator both appear in standard `display_indent()` / `displayable(...)`
output. No further wiring is required to make our operators visible in
`EXPLAIN PLAN` / `EXPLAIN ANALYZE` output.

### F4. BaselineMetrics work as documented. ✅

Wrapping the upstream stream with `stream.inspect_ok(move |b|
metrics.record_output(b.num_rows()))` is the entire integration. The
`MetricsSet` returned by `metrics()` includes the standard `OutputRows`,
`OutputBatches`, `OutputBytes`, `ElapsedCompute`, and `*Timestamp`
metrics. After two executions, `output_rows` correctly accumulates.

### F5. Composition with built-in ops works without ceremony. ✅

We registered the output of the expand as a `MemTable` and ran a SQL
query with `Filter(dst_id > 2) → Aggregate(count(*), edge_type) → GROUP BY`
over it. The aggregate returned the expected count (5). This proves the
output schema and physical batch layout are valid for downstream DF
operators — there is no hidden assumption that prevents our custom op
from being a peer of `FilterExec`, `ProjectionExec`, etc.

### F6. Unsafe usage: **none**. ✅

No `unsafe` was needed anywhere in the integration. All downcasting is
via the public `Any` interface; all stream wrapping uses `RecordBatchStreamAdapter`.

### F7. Internal-API access: **none**. ✅

The only borderline-internal item touched is `Boundedness::Bounded` and
`EmissionType::Incremental` from
`datafusion::physical_plan::execution_plan` — both are public enums.
Everything else (`PlanProperties`, `EquivalenceProperties`,
`Partitioning`, `BaselineMetrics`, `ExecutionPlanMetricsSet`,
`RecordBatchStreamAdapter`) is in `pub mod` paths.

## Awkward seams

These are not blockers but should be noted for the §11 RFC-body delta:

1. **`UserDefinedLogicalNodeCore::expressions()` semantics are subtle.**
   The doc says "expressions in the current node, not including inputs",
   but if you return non-empty `expressions()`, the optimizer will rewrite
   them and call `with_exprs_and_inputs(new_exprs, ...)`. For operators
   that don't have inline exprs (like ours), returning `vec![]` is the
   right answer — but it means we lose access to any
   constant-folding/simplification the optimizer would otherwise do for us.
   Document this in MR-737 §5.3: "graph operators must declare their
   inline exprs explicitly to benefit from CF/CP".

2. **`PartialOrd` requirement on `UserDefinedLogicalNodeCore`.**
   The trait requires `PartialOrd` because nodes participate in
   memoization / cache-key hashing in some optimizer passes. Most graph
   operators won't have a meaningful order, so they return `None`. The
   `#[derive(PartialOrd)]` requires all fields to also be `PartialOrd`,
   which `LogicalPlan` is not — so we manually impl
   `PartialOrd { partial_cmp -> None }`. Make this idiomatic in the RFC.

3. **`statistics()` and `partition_statistics()` are required.**
   For operators that *do* have statistics (like NeighborExpand: we can
   bound output rows by `sum(_neighbors.length)`), this is an opportunity
   for cost-aware planning. For prototypes we return `Statistics::new_unknown(...)`.
   §11 §statistics should call this out: graph operators must compute
   statistics or the planner will fall back to worst-case cardinality.

4. **`Partitioning` is inherited from input by default.**
   `NeighborExpandExec::new` clones `input.output_partitioning()`. For
   expand specifically this is wrong if the input is hash-partitioned by
   `src_id` and the consumer needs hash-by-`dst_id`. The pattern is to
   override `properties()` to declare a different partitioning. Note in
   the RFC that graph operators must explicitly choose their output
   partitioning rather than inheriting.

## Decision impact on MR-737 §5.3 and §5.10

**§5.3 is achievable on DataFusion 52.5 as written.** The
`UserDefinedLogicalNode`/`ExecutionPlan` surface is fully sufficient
for the operators §5.3 enumerates (Expand, MultiExpand, BackJoin,
NeighborSetIntersect, etc.). The only edits needed in §5.3:

  - Note that operators must explicitly opt-in to predicate push-down
    rather than rely on the default (which blocks all pushdown — the
    correct default for most graph ops).
  - Note the `PartialOrd` boilerplate as part of the operator skeleton.
  - Note the `statistics()` requirement: cost-aware planning depends on
    operators implementing it accurately, not punting to
    `Statistics::new_unknown`.

**§5.10 ("operators survive the optimizer + execute correctly")**:
The composition test (E4) plus the metrics test (E5) cover this. No
deltas needed.

## Caveats

- **Single-partition execution.** The probe uses partition count = 1.
  Multi-partition behavior (especially `required_input_distribution()`
  and `repartitioning`) is not exercised — but the `PlanProperties`
  surface for that is well-documented and was used by DF's own builtins
  the same way (verified by reading `FilterExec`).
- **The composition test materializes the output to a MemTable** to make
  it accessible from SQL. A more thorough test would build a single
  logical plan tree with `Filter(Expand(scan))` and run it; we'd need to
  hand-construct the LogicalPlan tree, which is straightforward but
  doesn't add new signal beyond what E1 + E4 already cover.
- **Schema gotcha caught during the experiment.** `ListBuilder<UInt64Builder>`
  defaults to a NULLABLE inner item even when the values appended don't
  contain nulls. If you declare a schema with `nullable=false` on the
  inner list field, `RecordBatch::try_new` rejects the array with
  `InvalidArgumentError("column types must match schema types")`. Worth
  a note in §5.3 or our internal Arrow guide — the default needs to be
  matched explicitly.

## Follow-ups (tracked, not done)

- Hand-build a multi-partition test plan with explicit
  `RepartitionExec(hash by src_id)` between scan and expand to verify
  our `output_partitioning()` choice survives the optimizer.
- Add a real `statistics()` implementation that consults
  `_neighbors.length` from upstream statistics (when present).