Skip to content

Commit a2db3f0

Browse files
Jiashu-Hualamb
andauthored
Documentation: Plan custom expressions (#15353)
* Added Custom Expression Planning Section * improve syntax * Update docs/source/library-user-guide/adding-udfs.md Co-authored-by: Andrew Lamb <[email protected]> * Update docs/source/library-user-guide/adding-udfs.md Co-authored-by: Andrew Lamb <[email protected]> * Update docs/source/library-user-guide/adding-udfs.md Co-authored-by: Andrew Lamb <[email protected]> * Update docs/source/library-user-guide/adding-udfs.md Co-authored-by: Andrew Lamb <[email protected]> * Update docs/source/library-user-guide/adding-udfs.md Co-authored-by: Andrew Lamb <[email protected]> * Update adding-udfs.md added assert for CI process * formatted language request by prettier --------- Co-authored-by: Andrew Lamb <[email protected]>
1 parent 0ff8984 commit a2db3f0

File tree

1 file changed

+85
-0
lines changed

1 file changed

+85
-0
lines changed

docs/source/library-user-guide/adding-udfs.md

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1160,6 +1160,91 @@ async fn main() -> Result<()> {
11601160
// +---+
11611161
```
11621162

1163+
## Custom Expression Planning
1164+
1165+
DataFusion provides native support for common SQL operators by default such as `+`, `-`, `||`. However it does not provide support for other operators such as `@>`. To override DataFusion's default handling or support unsupported operators, developers can extend DataFusion by implementing custom expression planning, a core feature of DataFusion
1166+
1167+
### Implementing Custom Expression Planning
1168+
1169+
To extend DataFusion with support for custom operators not natively available, you need to:
1170+
1171+
1. Implement the `ExprPlanner` trait: This allows you to define custom logic for planning expressions that DataFusion doesn't natively recognize. The trait provides the necessary interface to translate SQL AST nodes into logical `Expr`.
1172+
1173+
For detailed documentation please see: [Trait ExprPlanner](https://docs.rs/datafusion/latest/datafusion/logical_expr/planner/trait.ExprPlanner.html)
1174+
1175+
2. Register your custom planner: Integrate your implementation with DataFusion's `SessionContext` to ensure your custom planning logic is invoked during the query optimization and execution planning phase.
1176+
1177+
For a detailed documentation see: [fn register_expr_planner](https://docs.rs/datafusion/latest/datafusion/execution/trait.FunctionRegistry.html#method.register_expr_planner)
1178+
1179+
See example below:
1180+
1181+
```rust
1182+
# use arrow::array::RecordBatch;
1183+
# use std::sync::Arc;
1184+
1185+
# use datafusion::common::{assert_batches_eq, DFSchema};
1186+
# use datafusion::error::Result;
1187+
# use datafusion::execution::FunctionRegistry;
1188+
# use datafusion::logical_expr::Operator;
1189+
# use datafusion::prelude::*;
1190+
# use datafusion::sql::sqlparser::ast::BinaryOperator;
1191+
# use datafusion_common::ScalarValue;
1192+
# use datafusion_expr::expr::Alias;
1193+
# use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr};
1194+
# use datafusion_expr::BinaryExpr;
1195+
1196+
# #[derive(Debug)]
1197+
# // Define the custom planner
1198+
# struct MyCustomPlanner;
1199+
1200+
// Implement ExprPlanner to add support for the `->` custom operator
1201+
impl ExprPlanner for MyCustomPlanner {
1202+
fn plan_binary_op(
1203+
&self,
1204+
expr: RawBinaryExpr,
1205+
_schema: &DFSchema,
1206+
) -> Result<PlannerResult<RawBinaryExpr>> {
1207+
match &expr.op {
1208+
// Map `->` to string concatenation
1209+
BinaryOperator::Arrow => {
1210+
// Rewrite `->` as a string concatenation operation
1211+
// - `left` and `right` are the operands (e.g., 'hello' and 'world')
1212+
// - `Operator::StringConcat` tells DataFusion to concatenate them
1213+
Ok(PlannerResult::Planned(Expr::BinaryExpr(BinaryExpr {
1214+
left: Box::new(expr.left.clone()),
1215+
right: Box::new(expr.right.clone()),
1216+
op: Operator::StringConcat,
1217+
})))
1218+
}
1219+
_ => Ok(PlannerResult::Original(expr)),
1220+
}
1221+
}
1222+
}
1223+
1224+
use datafusion::execution::context::SessionContext;
1225+
use datafusion::arrow::util::pretty;
1226+
1227+
#[tokio::main]
1228+
async fn main() -> Result<()> {
1229+
let config = SessionConfig::new().set_str("datafusion.sql_parser.dialect", "postgres");
1230+
let mut ctx = SessionContext::new_with_config(config);
1231+
ctx.register_expr_planner(Arc::new(MyCustomPlanner))?;
1232+
let results = ctx.sql("select 'foo'->'bar';").await?.collect().await?;
1233+
1234+
let expected = [
1235+
"+----------------------------+",
1236+
"| Utf8(\"foo\") || Utf8(\"bar\") |",
1237+
"+----------------------------+",
1238+
"| foobar |",
1239+
"+----------------------------+",
1240+
];
1241+
assert_batches_eq!(&expected, &results);
1242+
1243+
pretty::print_batches(&results)?;
1244+
Ok(())
1245+
}
1246+
```
1247+
11631248
[1]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/simple_udf.rs
11641249
[2]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/simple_udwf.rs
11651250
[3]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/simple_udaf.rs

0 commit comments

Comments
 (0)