You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "mattsu2020 (via GitHub)" <gi...@apache.org> on 2026/10/01 22:10:19 UTC
[PR] docs: add examples for walking and rewriting logical plans [datafusion]
mattsu2020 opened a new pull request, #25954:
URL: https://github.com/apache/datafusion/pull/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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] docs: add examples for walking and rewriting logical plans [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on PR #25954:
URL: https://github.com/apache/datafusion/pull/25954#issuecomment-6334332136
🚀
@mattsu2020
Thank you for your contribution.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] docs: add examples for walking and rewriting logical plans [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew merged PR #25954:
URL: https://github.com/apache/datafusion/pull/25954
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] docs: add examples for walking and rewriting logical plans [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #25954:
URL: https://github.com/apache/datafusion/pull/25954#discussion_r4335173221
##########
datafusion-examples/examples/query_planning/plan_walk.rs:
##########
@@ -0,0 +1,185 @@
+// 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());
+
+ // 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 = table_scan_with_filters(
+ Some("eligible"),
+ &schema,
+ None,
+ vec![scan_predicate.clone()],
+ )?
+ .filter(col("salary").lt(lit(2000)))?
+ .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}");
+ }
+
+ 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
Review Comment:
Optional: A small `Filter -> Filter -> TableScan` example asserting `removed == 2` and leaving only the scan would make the nested-Filter behavior clearer. The current Filters are on separate branches, so they don't demonstrate consecutive Filter removal.
##########
datafusion-examples/examples/query_planning/plan_walk.rs:
##########
@@ -0,0 +1,185 @@
+// 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());
+
+ // 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 = table_scan_with_filters(
+ Some("eligible"),
+ &schema,
+ None,
+ vec![scan_predicate.clone()],
+ )?
+ .filter(col("salary").lt(lit(2000)))?
+ .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}");
+ }
+
+ let mut rewriter = RemoveFilters::default();
Review Comment:
Optional: Could we compare this with ordinary `rewrite` using a fresh `RemoveFilters`? Asserting `removed == 1` and checking that `summarize` still finds the subquery Filter and pushed-down scan predicate would help illustrate why `rewrite_with_subqueries` is needed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org