Enable Projection Pushdown Optimization for Recursive CTEs (#16696)

* Add column pruning support for RecursiveQuery operator in optimizer

- Extend optimize_projections to handle LogicalPlan::RecursiveQuery by applying
  projection pushdown to its inputs, improving query performance.
- Add integration test `recursive_query_column_pruning` verifying plan shows
  correct projection pruning for recursive CTEs.
- Implement `create_cte_work_table` in test context provider to support CTE tests.
- Add .github/copilot-instructions.md and AGENTS.md docs with Rust coding,
  linting, formatting, and contribution guidelines to maintain quality and
  consistency in generated and user code.

* Amend cte results

* Add recursive_cte.rs integration test and improve recursive CTE projection pushdown in cte.slt

- Add a new async test `recursive_cte_alias_instability` covering a complex recursive CTE query to the DataFusion core SQL tests, validating recursive CTE alias handling and query stability.
- Enhance existing sqllogictest file `cte.slt` by adding column projections to recursive CTE TableScans, improving plan efficiency and consistency.
- Fix projection pushdown for recursive CTEs in logical and physical plans for nodes, balances, recursive_cte, and numbers test cases.
- This addresses alias instability and projection inefficiencies previously observed in recursive CTE handling and improves test coverage for recursive SQL features.

* Remove redundant handling of RecursiveQuery in optimize_projections function

* main cte.slt

* Add tests for recursive CTE projection pushdown scenarios

* Remove unused create_cte_work_table function from MyContextProvider implementation

* Refactor TableScan projections in recursive query tests for clarity

* consolidate recursive cte tests

* Remove test that is included in slt

* Enhance optimization for recursive queries by adding checks for problematic structures

* Refactor recursive CTE tests to improve projection pushdown validation and add new test for alias instability

* refactor: rename function to clarify purpose of subquery alias detection in recursive queries

* test: add comments to clarify purpose of subquery alias handling in recursive CTE tests

* Refactor `plan_contains_subquery_alias` to count subquery aliases instead of returning a bool

- Changed `plan_contains_subquery_alias` to count occurrences of `SubqueryAlias` nodes in the plan.
- Added helper function `count_subquery_aliases` to recursively update the count.
- Updated logic to return true if there are two or more subquery aliases, preserving original behavior in `recursive_cte_alias_instability` test.

* Optimize subquery alias counting by early termination

- Updated `plan_contains_subquery_alias` to short-circuit counting once threshold (2) is reached.
- Modified `count_subquery_aliases` to accept a threshold and return early when count meets or exceeds it.
- Improves performance by avoiding unnecessary traversal of the entire plan.
- Added integration test snapshot for recursive query column pruning reflecting these changes.

* Refactor count_subquery_aliases to return count instead of using mutable reference

- Changed count_subquery_aliases to take count usize and return updated count, removing mutable reference.
- Added early-exit when count reaches threshold to avoid unnecessary traversal.
- Updated is_projection_unnecessary to use new count_subquery_aliases signature.
- Removed obsolete integration test snapshot file.
- Fixed recursive CTE plan and physical plan in optimizer integration and sqllogictest to include projection pruning on TableScans inside recursive queries.

This improves clarity, readability, and efficiency of subquery alias counting in logical plans, and fixes recursive query projections pruning as per related issue.

* docs(tests): add note referencing similar SQL in cte.slt for recursive CTE alias test

Add a comment in the recursive_cte_alias_instability test highlighting the similarity of the SQL query to one in datafusion/sqllogictest/test_files/cte.slt. This improves test clarity and cross-reference for maintainers.

* fix: clarify comment regarding alias ambiguity in recursive CTEs
feat: add initial example for explain_memory

* refactor: remove recursive_cte test module

* fix: add comment for clarity on recursive query handling

This commit adds a comment to the `optimize_projections` function in the `mod.rs` file to clarify the handling of recursive queries. The comment references a discussion on GitHub related to the implementation.

* remove stray file

* Revert "refactor: remove recursive_cte test module"

This reverts commit d8f6dd1.

* remove RecursiveQuery bypass

* refactor: improve handling of non-CTE subqueries in recursive queries

* refactor: remove unused recursive_cte module from SQL tests

* refactor: enhance projection optimization by handling non-CTE subqueries in recursive queries

* refactor: simplify TableScan projections in recursive query tests

* refactor: reorder filter and projection in recursive query logical plan

* refactor: restrict projection pushdown in recursive queries to only allow CTE references

* refactor: optimize TableScan projection in recursive query logical plan

* fix: correct dynamic filter predicate order in hash join test

* Update datafusion/optimizer/src/optimize_projections/mod.rs

Co-authored-by: Jeffrey Vo <[email protected]>

* Update datafusion/optimizer/src/optimize_projections/mod.rs

Co-authored-by: Jeffrey Vo <[email protected]>

* Enhance test description for recursive CTE projection pushdown

* Rename plan_contains_non_cte_subquery to plan_contains_other_subqueries for clarity

* Add test for recursive CTE with nested subquery to validate projection pushdown behavior

* Remove recursive_query_column_pruning

---------

Co-authored-by: Jeffrey Vo <[email protected]>

Author: kosiew

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -356,12 +356,35 @@ fn optimize_projections(
                 .collect::<Result<Vec<_>>>()?
         }
         LogicalPlan::EmptyRelation(_)
-        | LogicalPlan::RecursiveQuery(_)
         | LogicalPlan::Values(_)
         | LogicalPlan::DescribeTable(_) => {
             // These operators have no inputs, so stop the optimization process.
             return Ok(Transformed::no(plan));
         }
+        LogicalPlan::RecursiveQuery(recursive) => {
+            // Only allow subqueries that reference the current CTE; nested subqueries are not yet
+            // supported for projection pushdown for simplicity.
+            // TODO: be able to do projection pushdown on recursive CTEs with subqueries
+            if plan_contains_other_subqueries(
+                recursive.static_term.as_ref(),
+                &recursive.name,
+            ) || plan_contains_other_subqueries(
+                recursive.recursive_term.as_ref(),
+                &recursive.name,
+            ) {
+                return Ok(Transformed::no(plan));
+            }
+
+            plan.inputs()
+                .into_iter()
+                .map(|input| {
+                    indices
+                        .clone()
+                        .with_projection_beneficial()
+                        .with_plan_exprs(&plan, input.schema())
+                })
+                .collect::<Result<Vec<_>>>()?
+        }
         LogicalPlan::Join(join) => {
             let left_len = join.left.schema().fields().len();
             let (left_req_indices, right_req_indices) =
@@ -850,6 +873,46 @@ pub fn is_projection_unnecessary(
     ))
 }
 
+/// Returns true if the plan subtree contains any subqueries that are not the
+/// CTE reference itself. This treats any non-CTE [`LogicalPlan::SubqueryAlias`]
+/// node (including aliased relations) as a blocker, along with expression-level
+/// subqueries like scalar, EXISTS, or IN. These cases prevent projection
+/// pushdown for now because we cannot safely reason about their column usage.
+fn plan_contains_other_subqueries(plan: &LogicalPlan, cte_name: &str) -> bool {
+    if let LogicalPlan::SubqueryAlias(alias) = plan {
+        if alias.alias.table() != cte_name {
+            return true;
+        }
+    }
+
+    let mut found = false;
+    plan.apply_expressions(|expr| {
+        if expr_contains_subquery(expr) {
+            found = true;
+            Ok(TreeNodeRecursion::Stop)
+        } else {
+            Ok(TreeNodeRecursion::Continue)
+        }
+    })
+    .expect("expression traversal never fails");
+    if found {
+        return true;
+    }
+
+    plan.inputs()
+        .into_iter()
+        .any(|child| plan_contains_other_subqueries(child, cte_name))
+}
+
+fn expr_contains_subquery(expr: &Expr) -> bool {
+    expr.exists(|e| match e {
+        Expr::ScalarSubquery(_) | Expr::Exists(_) | Expr::InSubquery(_) => Ok(true),
+        _ => Ok(false),
+    })
+    // Safe unwrap since we are doing a simple boolean check
+    .unwrap()
+}
+
 #[cfg(test)]
 mod tests {
     use std::cmp::Ordering;

datafusion/optimizer/tests/optimizer_integration.rs

@@ -46,6 +46,48 @@ fn init() {
     let _ = env_logger::try_init();
 }
 
+#[test]
+fn recursive_cte_with_nested_subquery() -> Result<()> {
+    // Covers bailout path in `plan_contains_other_subqueries`, ensuring nested subqueries
+    // within recursive CTE branches prevent projection pushdown.
+    let sql = r#"
+        WITH RECURSIVE numbers(id, level) AS (
+            SELECT sub.id, sub.level FROM (
+                SELECT col_int32 AS id, 1 AS level FROM test
+            ) sub
+            UNION ALL
+            SELECT t.col_int32, numbers.level + 1
+            FROM test t
+            JOIN numbers ON t.col_int32 = numbers.id + 1
+        )
+        SELECT id, level FROM numbers
+    "#;
+
+    let plan = test_sql(sql)?;
+
+    assert_snapshot!(
+        format!("{plan}"),
+        @r#"
+        SubqueryAlias: numbers
+          Projection: sub.id AS id, sub.level AS level
+            RecursiveQuery: is_distinct=false
+              Projection: sub.id, sub.level
+                SubqueryAlias: sub
+                  Projection: test.col_int32 AS id, Int64(1) AS level
+                    TableScan: test
+              Projection: t.col_int32, numbers.level + Int64(1)
+                Inner Join: CAST(t.col_int32 AS Int64) = CAST(numbers.id AS Int64) + Int64(1)
+                  SubqueryAlias: t
+                    Filter: CAST(test.col_int32 AS Int64) IS NOT NULL
+                      TableScan: test
+                  Filter: CAST(numbers.id AS Int64) + Int64(1) IS NOT NULL
+                    TableScan: numbers
+        "#
+    );
+
+    Ok(())
+}
+
 #[test]
 fn case_when() -> Result<()> {
     let sql = "SELECT CASE WHEN col_int32 > 0 THEN 1 ELSE 0 END FROM test";
@@ -478,6 +520,94 @@ fn select_correlated_predicate_subquery_with_uppercase_ident() {
     );
 }
 
+#[test]
+fn recursive_cte_projection_pushdown() -> Result<()> {
+    // Test that projection pushdown works with recursive CTEs by ensuring
+    // only the required columns are projected from the base table, even when
+    // the CTE definition includes unused columns
+    let sql = "WITH RECURSIVE nodes AS (\
+        SELECT col_int32 AS id, col_utf8 AS name, col_uint32 AS extra FROM test \
+        UNION ALL \
+        SELECT id + 1, name, extra FROM nodes WHERE id < 3\
+    ) SELECT id FROM nodes";
+    let plan = test_sql(sql)?;
+
+    // The optimizer successfully performs projection pushdown by only selecting the needed
+    // columns from the base table and recursive table, eliminating unused columns
+    assert_snapshot!(
+        format!("{plan}"),
+        @r#"SubqueryAlias: nodes
+  RecursiveQuery: is_distinct=false
+    Projection: test.col_int32 AS id
+      TableScan: test projection=[col_int32]
+    Projection: CAST(CAST(nodes.id AS Int64) + Int64(1) AS Int32)
+      Filter: nodes.id < Int32(3)
+        TableScan: nodes projection=[id]
+"#
+    );
+    Ok(())
+}
+
+#[test]
+fn recursive_cte_with_unused_columns() -> Result<()> {
+    // Test projection pushdown with a recursive CTE where the base case
+    // includes columns that are never used in the recursive part or final result
+    let sql = "WITH RECURSIVE series AS (\
+        SELECT 1 AS n, col_utf8, col_uint32, col_date32 FROM test WHERE col_int32 = 1 \
+        UNION ALL \
+        SELECT n + 1, col_utf8, col_uint32, col_date32 FROM series WHERE n < 3\
+    ) SELECT n FROM series";
+    let plan = test_sql(sql)?;
+
+    // The optimizer successfully performs projection pushdown by eliminating unused columns
+    // even when they're defined in the CTE but not actually needed
+    assert_snapshot!(
+        format!("{plan}"),
+        @r#"SubqueryAlias: series
+  RecursiveQuery: is_distinct=false
+    Projection: Int64(1) AS n
+      Filter: test.col_int32 = Int32(1)
+        TableScan: test projection=[col_int32]
+    Projection: series.n + Int64(1)
+      Filter: series.n < Int64(3)
+        TableScan: series projection=[n]
+"#
+    );
+    Ok(())
+}
+
+#[test]
+/// Asserts the minimal plan shape once projection pushdown succeeds for a recursive CTE.
+/// Unlike the previous two tests that retain extra columns in either the base or recursive
+/// branches, this baseline shows the optimizer trimming everything down to the single
+/// column required by the final projection.
+fn recursive_cte_projection_pushdown_baseline() -> Result<()> {
+    // Test case that truly demonstrates projection pushdown working:
+    // The base case only selects needed columns
+    let sql = "WITH RECURSIVE countdown AS (\
+        SELECT col_int32 AS n FROM test WHERE col_int32 = 5 \
+        UNION ALL \
+        SELECT n - 1 FROM countdown WHERE n > 1\
+    ) SELECT n FROM countdown";
+    let plan = test_sql(sql)?;
+
+    // This demonstrates optimal projection pushdown where only col_int32 is projected from the base table,
+    // and only the needed column is selected from the recursive table
+    assert_snapshot!(
+        format!("{plan}"),
+        @r#"SubqueryAlias: countdown
+  RecursiveQuery: is_distinct=false
+    Projection: test.col_int32 AS n
+      Filter: test.col_int32 = Int32(5)
+        TableScan: test projection=[col_int32]
+    Projection: CAST(CAST(countdown.n AS Int64) - Int64(1) AS Int32)
+      Filter: countdown.n > Int32(1)
+        TableScan: countdown projection=[n]
+"#
+    );
+    Ok(())
+}
+
 fn test_sql(sql: &str) -> Result<LogicalPlan> {
     // parse the SQL
     let dialect = GenericDialect {}; // or AnsiDialect, or your own dialect ...
@@ -587,6 +717,14 @@ impl ContextProvider for MyContextProvider {
         None
     }
 
+    fn create_cte_work_table(
+        &self,
+        _name: &str,
+        schema: SchemaRef,
+    ) -> Result<Arc<dyn TableSource>> {
+        Ok(Arc::new(MyTableSource { schema }))
+    }
+
     fn options(&self) -> &ConfigOptions {
         &self.options
     }

datafusion/sqllogictest/test_files/cte.slt

@@ -110,7 +110,7 @@ logical_plan
 04)------EmptyRelation: rows=1
 05)----Projection: nodes.id + Int64(1) AS id
 06)------Filter: nodes.id < Int64(10)
-07)--------TableScan: nodes
+07)--------TableScan: nodes projection=[id]
 physical_plan
 01)RecursiveQueryExec: name=nodes, is_distinct=false
 02)--ProjectionExec: expr=[1 as id]
@@ -152,11 +152,10 @@ logical_plan
 01)Sort: balances.time ASC NULLS LAST, balances.name ASC NULLS LAST, balances.account_balance ASC NULLS LAST
 02)--SubqueryAlias: balances
 03)----RecursiveQuery: is_distinct=false
-04)------Projection: balance.time, balance.name, balance.account_balance
-05)--------TableScan: balance
-06)------Projection: balances.time + Int64(1) AS time, balances.name, balances.account_balance + Int64(10) AS account_balance
-07)--------Filter: balances.time < Int64(10)
-08)----------TableScan: balances
+04)------TableScan: balance projection=[time, name, account_balance]
+05)------Projection: balances.time + Int64(1) AS time, balances.name, balances.account_balance + Int64(10) AS account_balance
+06)--------Filter: balances.time < Int64(10)
+07)----------TableScan: balances projection=[time, name, account_balance]
 physical_plan
 01)SortExec: expr=[time@0 ASC NULLS LAST, name@1 ASC NULLS LAST, account_balance@2 ASC NULLS LAST], preserve_partitioning=[false]
 02)--RecursiveQueryExec: name=balances, is_distinct=false
@@ -958,7 +957,7 @@ logical_plan
 04)------EmptyRelation: rows=1
 05)----Projection: numbers.n + Int64(1)
 06)------Filter: numbers.n < Int64(10)
-07)--------TableScan: numbers
+07)--------TableScan: numbers projection=[n]
 physical_plan
 01)RecursiveQueryExec: name=numbers, is_distinct=false
 02)--ProjectionExec: expr=[1 as n]
@@ -984,7 +983,7 @@ logical_plan
 04)------EmptyRelation: rows=1
 05)----Projection: numbers.n + Int64(1)
 06)------Filter: numbers.n < Int64(10)
-07)--------TableScan: numbers
+07)--------TableScan: numbers projection=[n]
 physical_plan
 01)RecursiveQueryExec: name=numbers, is_distinct=false
 02)--ProjectionExec: expr=[1 as n]
@@ -1041,8 +1040,7 @@ logical_plan
 04)------Projection: Int64(0) AS k, Int64(0) AS v
 05)--------EmptyRelation: rows=1
 06)------Sort: r.v ASC NULLS LAST, fetch=1
-07)--------Projection: r.k, r.v
-08)----------TableScan: r
+07)--------TableScan: r projection=[k, v]
 physical_plan
 01)GlobalLimitExec: skip=0, fetch=5
 02)--RecursiveQueryExec: name=r, is_distinct=false