Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 78 additions & 1 deletion docs/source/library-user-guide/extending-operators.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one

# Extending Operators

DataFusion supports extending operators by transforming [`LogicalPlan`] and [`ExecutionPlan`] through customized [optimizer rules](https://docs.rs/datafusion/latest/datafusion/optimizer/trait.OptimizerRule.html). This section will use the µWheel project to illustrate such capabilities.
DataFusion supports extending operators with custom [`LogicalPlan`] and [`ExecutionPlan`] nodes and [optimizer rules](https://docs.rs/datafusion/latest/datafusion/optimizer/trait.OptimizerRule.html). The examples below show two ways to extend query planning.

[`logicalplan`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.LogicalPlan.html
[`executionplan`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/trait.ExecutionPlan.html
Expand Down Expand Up @@ -47,3 +47,80 @@ fn agg_to_table_scan(result: f64, schema: SchemaRef) -> Result<LogicalPlan> {
```

To get a deeper dive into the usage of the µWheel project, visit the [blog post](https://uwheel.rs/post/datafusion_uwheel/) by Max Meldrum.

## A custom TopK operator

The [TopK example] shows how to add a custom operator that finds the rows with
the highest values. It uses a query like this one to find the three customers
with the highest revenue:

```sql
SELECT customer_id, revenue FROM sales ORDER BY revenue DESC LIMIT 3;
```

DataFusion already has optimized support for TopK queries. This example is for
learning how to connect a custom operator to DataFusion, not for replacing the
built-in implementation.

### Describe the operation

The example defines `TopKPlanNode`, a custom logical plan node. A logical plan
describes the operations in a query. The node implements
`UserDefinedLogicalNodeCore` so DataFusion can inspect and update it. Its main
fields are:

```rust,ignore
struct TopKPlanNode {
k: usize,
input: LogicalPlan,
expr: SortExpr,
}
```

- `input` describes where the rows come from, such as the `sales` table.
- `expr` describes how to rank the rows, such as by revenue in descending order.
- `k` is the number of rows to keep. In this query, `k` is 3.

### Add the node to the plan

`TopKOptimizerRule` looks for a sort that has a row limit. The limit can be
part of the sort, or it can be a separate limit directly above the sort. The
sort must have exactly one sort expression. The rule replaces that part of the
plan with a `TopKPlanNode`:

```rust,ignore
LogicalPlan::Extension(Extension {
node: Arc::new(TopKPlanNode {
k: *fetch,
input: input.as_ref().clone(),
expr: expr[0].clone(),
}),
})
```

The rule takes `k` from the row limit. It copies the input and the sort
expression from the sort. `LogicalPlan::Extension` wraps the custom node so
that it can be part of a logical plan.

For the query above, the new plan includes:

```text
TopK: k=3
TableScan: sales
```

The rule is registered with the session so DataFusion can apply it while
optimizing the query.

### Execute the operation

`TopKPlanner` converts the logical `TopKPlanNode` into `TopKExec`, a physical
operator. `TopKExec` reads its input and uses `TopKReader` to keep the highest
values. Once it has read the input, the reader returns the retained rows.

The example is specialized for its test data. It expects particular columns
and types, and its map stores only one customer per revenue value. It does not
handle all sort options or other edge cases, so it should not be used as a
general-purpose TopK operator.

[topk example]: https://github.com/apache/datafusion/blob/main/datafusion/core/tests/user_defined/user_defined_plan.rs