You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/22 06:29:22 UTC
(datafusion) branch gh-readonly-queue/main/pr-25954-457ccaa56be2af30b3a6c6edd76ba4d155c836dd created (now 1e2ba67cf9)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-25954-457ccaa56be2af30b3a6c6edd76ba4d155c836dd
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 1e2ba67cf9 docs: add examples for walking and rewriting logical plans (#25954)
This branch includes the following new commits:
new 1e2ba67cf9 docs: add examples for walking and rewriting logical plans (#25954)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: docs: add examples for walking and rewriting logical plans (#25954)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25954-457ccaa56be2af30b3a6c6edd76ba4d155c836dd
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 1e2ba67cf99b3f5cfb6989385f52612a4c0977c6
Author: mattsu <35...@users.noreply.github.com>
AuthorDate: Thu Oct 22 06:28:59 2026 +0000
docs: add examples for walking and rewriting logical plans (#25954)
## Which issue does this PR close?
- Closes #10628.
## Rationale for this change
Users need concrete examples of inspecting and rewriting `LogicalPlan`
trees with the `TreeNode` APIs, including subqueries embedded in
expressions.
## What changes are included in this PR?
- Extend the `LogicalPlan` rustdoc example to collect table names with
`TreeNode::apply`.
- Add a tested `TreeNode::rewrite` example that removes Filter nodes,
explaining traversal order and the change in query semantics.
- Add the `query_planning -- plan_walk` example demonstrating plan
inspection, stateful rewriting, subquery traversal, and collecting
predicates from both Filter nodes and TableScan filters.
- Register the example in the query planning runner and update the
generated examples README.
## What is the testing strategy for this PR?
The rustdoc examples and assertions in the executable example verify the
collected table names, rewritten plan, subquery traversal, and filter
predicates.
- `cargo test --profile ci -p datafusion-expr --doc
logical_plan::plan::LogicalPlan` (8 passed)
- `cargo run --profile ci -p datafusion-examples --example
query_planning -- plan_walk`
- `cargo fmt --all`
- `cargo clippy --all-targets --all-features -- -D warnings`
- `./ci/scripts/doc_prettier_check.sh --write --allow-dirty`
- `./dev/rust_lint.sh` (run via `uv run --package datafusion-dev --with
PyYAML`)
## Are there any user-facing changes?
The API documentation now includes additional plan walking and rewriting
examples and links to an executable walkthrough. There are no changes to
query execution or public APIs.
---------
Co-authored-by: kosiew <ko...@gmail.com>
---
datafusion-examples/README.md | 1 +
.../examples/query_planning/main.rs | 8 +-
.../examples/query_planning/plan_walk.rs | 216 +++++++++++++++++++++
datafusion/expr/src/logical_plan/plan.rs | 67 ++++++-
4 files changed, 288 insertions(+), 4 deletions(-)
diff --git a/datafusion-examples/README.md b/datafusion-examples/README.md
index f3c55001a6..c8bc380ecc 100644
--- a/datafusion-examples/README.md
+++ b/datafusion-examples/README.md
@@ -186,6 +186,7 @@ cargo run --example dataframe -- dataframe
| optimizer_rule | [`query_planning/optimizer_rule.rs`](examples/query_planning/optimizer_rule.rs) | Replace predicates via a custom OptimizerRule |
| parse_sql_expr | [`query_planning/parse_sql_expr.rs`](examples/query_planning/parse_sql_expr.rs) | Parse SQL into DataFusion Expr |
| plan_to_sql | [`query_planning/plan_to_sql.rs`](examples/query_planning/plan_to_sql.rs) | Generate SQL from expressions or plans |
+| plan_walk | [`query_planning/plan_walk.rs`](examples/query_planning/plan_walk.rs) | Walk and rewrite LogicalPlans with TreeNode APIs |
| planner_api | [`query_planning/planner_api.rs`](examples/query_planning/planner_api.rs) | APIs for logical and physical plan manipulation |
| pruning | [`query_planning/pruning.rs`](examples/query_planning/pruning.rs) | Use pruning to skip irrelevant files |
| thread_pools | [`query_planning/thread_pools.rs`](examples/query_planning/thread_pools.rs) | Configure custom thread pools for DataFusion execution |
diff --git a/datafusion-examples/examples/query_planning/main.rs b/datafusion-examples/examples/query_planning/main.rs
index 2e4310082c..e7e11c2b1a 100644
--- a/datafusion-examples/examples/query_planning/main.rs
+++ b/datafusion-examples/examples/query_planning/main.rs
@@ -21,7 +21,7 @@
//!
//! ## Usage
//! ```bash
-//! cargo run --example query_planning -- [all|analyzer_rule|expr_api|optimizer_rule|parse_sql_expr|plan_to_sql|planner_api|pruning|thread_pools]
+//! cargo run --example query_planning -- [all|analyzer_rule|expr_api|optimizer_rule|parse_sql_expr|plan_to_sql|plan_walk|planner_api|pruning|thread_pools]
//! ```
//!
//! Each subcommand runs a corresponding example:
@@ -42,6 +42,9 @@
//! - `plan_to_sql`
//! (file: plan_to_sql.rs, desc: Generate SQL from expressions or plans)
//!
+//! - `plan_walk`
+//! (file: plan_walk.rs, desc: Walk and rewrite LogicalPlans with TreeNode APIs)
+//!
//! - `planner_api`
//! (file: planner_api.rs, desc: APIs for logical and physical plan manipulation)
//!
@@ -56,6 +59,7 @@ mod expr_api;
mod optimizer_rule;
mod parse_sql_expr;
mod plan_to_sql;
+mod plan_walk;
mod planner_api;
mod pruning;
mod thread_pools;
@@ -73,6 +77,7 @@ enum ExampleKind {
OptimizerRule,
ParseSqlExpr,
PlanToSql,
+ PlanWalk,
PlannerApi,
Pruning,
ThreadPools,
@@ -98,6 +103,7 @@ impl ExampleKind {
ExampleKind::OptimizerRule => optimizer_rule::optimizer_rule().await?,
ExampleKind::ParseSqlExpr => parse_sql_expr::parse_sql_expr().await?,
ExampleKind::PlanToSql => plan_to_sql::plan_to_sql_examples().await?,
+ ExampleKind::PlanWalk => plan_walk::plan_walk()?,
ExampleKind::PlannerApi => planner_api::planner_api().await?,
ExampleKind::Pruning => pruning::pruning()?,
ExampleKind::ThreadPools => thread_pools::thread_pools().await?,
diff --git a/datafusion-examples/examples/query_planning/plan_walk.rs b/datafusion-examples/examples/query_planning/plan_walk.rs
new file mode 100644
index 0000000000..95458f3357
--- /dev/null
+++ b/datafusion-examples/examples/query_planning/plan_walk.rs
@@ -0,0 +1,216 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Walk and rewrite logical plans, including subqueries embedded in expressions.
+//! Run with `cargo run --example query_planning -- plan_walk`.
+
+use std::collections::BTreeSet;
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::common::Result;
+use datafusion::common::tree_node::{
+ Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter,
+};
+use datafusion::logical_expr::logical_plan::builder::table_scan_with_filters;
+use datafusion::logical_expr::{Expr, LogicalPlan, col, in_subquery, lit, table_scan};
+
+pub fn plan_walk() -> Result<()> {
+ let schema = Schema::new(vec![
+ Field::new("name", DataType::Utf8, false),
+ Field::new("salary", DataType::Int32, false),
+ ]);
+
+ // table_scan returns a LogicalPlanBuilder. Build a tree of operators:
+ // Projection(name, salary)
+ // Filter(salary > 1000)
+ // TableScan(employee)
+ // Only a schema is needed to inspect these plans; no data is read.
+ let plan = table_scan(Some("employee"), &schema, None)?
+ .filter(col("salary").gt(lit(1000)))?
+ .project(vec![col("name"), col("salary")])?
+ .build()?;
+ println!("Original plan:\n{}", plan.display_indent());
+
+ // apply borrows the plan and visits each node before visiting its inputs.
+ // Continue visits the inputs, Jump skips them, and Stop ends the walk.
+ let mut tables = BTreeSet::new();
+ plan.apply(|node| {
+ if let LogicalPlan::TableScan(scan) = node {
+ tables.insert(scan.table_name.to_string());
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ assert_eq!(tables, BTreeSet::from(["employee".to_string()]));
+ println!("Referenced tables: {tables:?}");
+
+ // rewrite consumes the plan. Clone it here to keep the original for comparison.
+ // This illustrative rewrite changes query semantics by removing filters;
+ // it is not an optimizer rule that preserves query results.
+ let mut rewriter = RemoveFilters::default();
+ let rewritten = plan.clone().rewrite(&mut rewriter)?;
+ assert!(rewritten.transformed);
+ assert_eq!(rewriter.removed, 1);
+ assert_eq!(rewritten.data.schema(), plan.schema());
+ assert_eq!(
+ rewritten.data.display_indent().to_string(),
+ "Projection: employee.name, employee.salary\n TableScan: employee"
+ );
+ println!("Without Filter nodes:\n{}", rewritten.data.display_indent());
+
+ // Consecutive Filters are removed bottom-up: f_up sees the outer Filter
+ // after its inner Filter has already been replaced with the TableScan.
+ let plan = table_scan(Some("employee"), &schema, None)?
+ .filter(col("salary").gt(lit(1000)))?
+ .filter(col("salary").lt(lit(2000)))?
+ .build()?;
+ assert_eq!(
+ plan.display_indent().to_string(),
+ "Filter: employee.salary < Int32(2000)\n Filter: employee.salary > Int32(1000)\n TableScan: employee"
+ );
+ let mut rewriter = RemoveFilters::default();
+ let rewritten = plan.rewrite(&mut rewriter)?;
+ assert!(rewritten.transformed);
+ assert_eq!(rewriter.removed, 2);
+ assert_eq!(
+ rewritten.data.display_indent().to_string(),
+ "TableScan: employee"
+ );
+
+ // A subquery in an expression is not an ordinary plan input, so apply
+ // and rewrite do not traverse it. Use their *_with_subqueries variants.
+ // Include a scan filter to illustrate predicates that have been pushed down
+ // into a table source. Here it is set explicitly rather than by an optimizer.
+ let scan_predicate = col("eligible.salary").gt(lit(1000));
+ let subquery_predicate = col("eligible.salary").lt(lit(2000));
+ let subquery = table_scan_with_filters(
+ Some("eligible"),
+ &schema,
+ None,
+ vec![scan_predicate.clone()],
+ )?
+ .filter(subquery_predicate.clone())?
+ .project(vec![col("salary")])?
+ .build()?;
+ let plan = table_scan(Some("employee"), &schema, None)?
+ .filter(col("salary").gt(lit(0)))?
+ .project(vec![
+ col("name"),
+ in_subquery(col("salary"), Arc::new(subquery)).alias("is_eligible"),
+ ])?
+ .build()?;
+
+ tables.clear();
+ plan.apply(|node| {
+ if let LogicalPlan::TableScan(scan) = node {
+ tables.insert(scan.table_name.to_string());
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ assert_eq!(tables, BTreeSet::from(["employee".to_string()]));
+
+ let summary = summarize(&plan)?;
+ assert_eq!(
+ summary.tables,
+ BTreeSet::from(["eligible".to_string(), "employee".to_string()])
+ );
+ // The outer Filter, the subquery's Filter, and the scan's pushed-down filter.
+ assert_eq!(summary.filters.len(), 3);
+ assert!(summary.filters.contains(&scan_predicate));
+ println!("Tables including subqueries: {:?}", summary.tables);
+ for predicate in &summary.filters {
+ println!("Filter predicate: {predicate}");
+ }
+
+ // Ordinary rewrite removes only the outer Filter. The subquery's Filter
+ // and pushed-down scan predicate remain because it does not visit subqueries.
+ let mut rewriter = RemoveFilters::default();
+ let rewritten = plan.clone().rewrite(&mut rewriter)?;
+ assert!(rewritten.transformed);
+ assert_eq!(rewriter.removed, 1);
+ let summary = summarize(&rewritten.data)?;
+ assert_eq!(summary.filters.len(), 2);
+ assert!(summary.filters.contains(&subquery_predicate));
+ assert!(summary.filters.contains(&scan_predicate));
+
+ let mut rewriter = RemoveFilters::default();
+ let rewritten = plan.rewrite_with_subqueries(&mut rewriter)?;
+ assert!(rewritten.transformed);
+ assert_eq!(rewriter.removed, 2);
+ // Both Filter nodes were removed, but the subquery in the Projection and
+ // its pushed-down scan filter remain. Rewriting Filter nodes does not
+ // modify TableScan.filters.
+ let summary = summarize(&rewritten.data)?;
+ assert_eq!(
+ summary.tables,
+ BTreeSet::from(["eligible".to_string(), "employee".to_string()])
+ );
+ assert_eq!(summary.filters, vec![scan_predicate]);
+
+ Ok(())
+}
+
+#[derive(Debug, Default)]
+struct PlanSummary {
+ tables: BTreeSet<String>,
+ filters: Vec<Expr>,
+}
+
+fn summarize(plan: &LogicalPlan) -> Result<PlanSummary> {
+ let mut summary = PlanSummary::default();
+ plan.apply_with_subqueries(|node| {
+ match node {
+ LogicalPlan::TableScan(scan) => {
+ summary.tables.insert(scan.table_name.to_string());
+ summary.filters.extend(scan.filters.iter().cloned());
+ }
+ LogicalPlan::Filter(filter) => {
+ summary.filters.push(filter.predicate.clone());
+ }
+ _ => {}
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ // Filter pushdown support determines whether predicates appear in Filter
+ // nodes, TableScan.filters, or both. This list may therefore have duplicates.
+ // Other predicates, such as join conditions, are not collected here.
+ Ok(summary)
+}
+
+// A rewriter can keep state across nodes, such as the number of removed filters.
+#[derive(Default)]
+struct RemoveFilters {
+ removed: usize,
+}
+
+impl TreeNodeRewriter for RemoveFilters {
+ type Node = LogicalPlan;
+
+ // f_up runs after the inputs (and, with rewrite_with_subqueries, embedded
+ // subqueries) have been rewritten. Returning the rewritten input also works
+ // when several Filter nodes are nested.
+ // For a stateless bottom-up rewrite, transform_up accepts a closure instead.
+ fn f_up(&mut self, node: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+ match node {
+ LogicalPlan::Filter(filter) => {
+ self.removed += 1;
+ Ok(Transformed::yes(filter.input.as_ref().clone()))
+ }
+ _ => Ok(Transformed::no(node)),
+ }
+ }
+}
diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
index 10cc92d6b5..600f1331c6 100644
--- a/datafusion/expr/src/logical_plan/plan.rs
+++ b/datafusion/expr/src/logical_plan/plan.rs
@@ -111,7 +111,8 @@ pub use datafusion_common::{JoinConstraint, JoinType};
/// # Visiting and Rewriting `LogicalPlan`s
///
/// Using the [`tree_node`] API, you can recursively walk all nodes in a
-/// `LogicalPlan`. For example, to find all column references in a plan:
+/// `LogicalPlan`. [`TreeNode::apply`] visits a node before its inputs. For
+/// example, to find table names and expressions in a plan:
///
/// ```
/// # use std::collections::HashSet;
@@ -134,17 +135,22 @@ pub use datafusion_common::{JoinConstraint, JoinType};
/// .project(vec![col("name"), col("salary")])?
/// .build()?;
///
-/// // use apply to walk the plan and collect all expressions
+/// // use apply to walk the plan and collect table names and expressions
+/// let mut table_names = HashSet::new();
/// let mut expressions = HashSet::new();
/// plan.apply(|node| {
+/// if let LogicalPlan::TableScan(scan) = node {
+/// table_names.insert(scan.table_name.to_string());
+/// }
/// // collect all expressions in the plan
/// node.apply_expressions(|expr| {
/// expressions.insert(expr.clone());
/// Ok(TreeNodeRecursion::Continue) // control walk of expressions
/// })?;
/// Ok(TreeNodeRecursion::Continue) // control walk of plan nodes
-/// }).unwrap();
+/// })?;
///
+/// assert_eq!(table_names, HashSet::from(["employee".to_string()]));
/// // we found the expression in projection and filter
/// assert_eq!(expressions.len(), 3);
/// println!("Found expressions: {:?}", expressions);
@@ -160,6 +166,11 @@ pub use datafusion_common::{JoinConstraint, JoinType};
/// # }
/// ```
///
+/// [`TreeNode::apply`] and [`TreeNode::rewrite`] traverse plan inputs, but do
+/// not visit subqueries embedded in expressions. Use
+/// [`LogicalPlan::apply_with_subqueries`] and
+/// [`LogicalPlan::rewrite_with_subqueries`] to include those subqueries.
+///
/// You can also rewrite plans using the [`tree_node`] API. For example, to
/// replace the filter predicate in a plan:
///
@@ -211,6 +222,56 @@ pub use datafusion_common::{JoinConstraint, JoinType};
/// # Ok(())
/// # }
/// ```
+///
+/// For more control over traversal, implement a [`TreeNodeRewriter`] and use
+/// [`TreeNode::rewrite`]. Its `f_down` method runs before rewriting inputs and
+/// `f_up` runs after them. The following example removes Filter nodes in
+/// `f_up` by replacing each with its already rewritten input. This changes
+/// query semantics and is only an illustration of the API.
+///
+/// ```
+/// # use arrow::datatypes::{DataType, Field, Schema};
+/// # use datafusion_expr::{col, lit, LogicalPlan, table_scan};
+/// # use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
+/// # use datafusion_common::Result;
+/// # fn main() -> Result<()> {
+/// # let schema = Schema::new(vec![
+/// # Field::new("name", DataType::Utf8, false),
+/// # Field::new("salary", DataType::Int32, false),
+/// # ]);
+/// let plan = table_scan(Some("employee"), &schema, None)?
+/// .filter(col("salary").gt(lit(1000)))?
+/// .project(vec![col("name"), col("salary")])?
+/// .build()?;
+///
+/// struct RemoveFilters;
+///
+/// impl TreeNodeRewriter for RemoveFilters {
+/// type Node = LogicalPlan;
+///
+/// fn f_up(&mut self, node: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+/// match node {
+/// LogicalPlan::Filter(filter) => {
+/// Ok(Transformed::yes(filter.input.as_ref().clone()))
+/// }
+/// _ => Ok(Transformed::no(node)),
+/// }
+/// }
+/// }
+///
+/// let rewritten = plan.rewrite(&mut RemoveFilters)?;
+/// assert!(rewritten.transformed);
+/// assert_eq!(rewritten.data.display_indent().to_string(),
+/// "Projection: employee.name, employee.salary\n TableScan: employee");
+/// # Ok(())
+/// # }
+/// ```
+///
+/// See the [plan walking example] for a complete walkthrough, including
+/// collecting tables and filter predicates from subqueries.
+///
+/// [`TreeNodeRewriter`]: datafusion_common::tree_node::TreeNodeRewriter
+/// [plan walking example]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/query_planning/plan_walk.rs
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
pub enum LogicalPlan {
/// Evaluates an arbitrary list of expressions (essentially a
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org