[adriangb]

Which issue does this PR close?

• Closes Prune unused columns from the body of a materialized CTE #26240.

Stacked on #26269. Only the last commit is new.

Rationale for this change

The body of a materialized CTE kept all its columns, also when no reference read them.

CREATE TABLE wide(a INT, b INT, c INT, d INT, e INT) AS VALUES (1, 10, 100, 1000, 10000);

EXPLAIN WITH w AS MATERIALIZED (SELECT * FROM wide)
SELECT a FROM w
UNION ALL
SELECT b FROM w;

Before, all five columns were read and buffered:

MaterializedCte: name=w
  SubqueryAlias: w
    TableScan: wide projection=[a, b, c, d, e]
  Union
    Projection: w.a
      MaterializedCteScan: name=w
    Projection: w.b AS a
      MaterializedCteScan: name=w

After, only a and b are read and buffered:

MaterializedCte: name=w
  SubqueryAlias: w
    TableScan: wide projection=[a, b]
  Union
    MaterializedCteScan: name=w projection=[a]
    Projection: w.b AS a
      MaterializedCteScan: name=w projection=[b]

What changes are included in this PR?

• MaterializedCteScan has a new field projection: Option<Vec<usize>>: the indexes of the columns of the CTE body that the scan outputs.
• OptimizeProjections narrows each MaterializedCteScan to the columns that its parent needs, as it does for a TableScan.
• At a MaterializedCte, OptimizeProjections optimizes the continuation first. It then computes the union of the columns of all the scans of the CTE, projects the body to this union, and gives each scan the new positions of its columns. Scans in subquery expressions and in the bodies of other CTEs are counted.
• MaterializedCteScanExec applies the projection to the shared buffered batches.
• The physical planner checks that each scan agrees with the physical schema of its body when it binds the scan, and returns an internal error if it does not.

The projection in the scan makes the intermediate state (a narrow scan and a wide body) a valid plan, so no separate rule is necessary.

Notes:

• A reference that reads no column (SELECT count(*) FROM c) gets projection=[]. If no reference reads a column, the body keeps its first column.
• The projection is put above the body. A volatile column can be removed, but the number of rows does not change. Columns below a DISTINCT or a recursive query are not removed.
• A scan in a correlated subquery that is not decorrelated is not narrowed. The CTE then keeps all its columns.
• The pgjson display does not show the projection.

Comparison with PostgreSQL and DuckDB. This is what DuckDB does: RemoveUnusedColumns prunes the body to the columns that its references read (remove_unused_columns.cpp). PostgreSQL plans a CTE for full retrieval and does not prune (subselect.c).

Are these changes tested?

Yes. cte_materialized.slt: EXPLAIN and result checks for the example above, a column that only a filter reads, count(*), a CTE that another CTE reads, references in subqueries, a volatile body, a DISTINCT body, aliased columns, a self join, a recursive body and a CTE that nothing reads. Unit tests for the scan projection, nested CTEs, idempotence and the zero-column case.

Are there any user-facing changes?

EXPLAIN shows projection=[...] on MaterializedCteScan and MaterializedCteScanExec when the scan does not read all the columns. MaterializedCteScan has the new public field projection. The upgrade guide for 56.0.0 says this.

[github-actions]

Thank you for opening this pull request!

Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).

Details

     Cloning af752e5d5620b13e759e1713c6c0b72c43b4fb9b
    Building datafusion v55.1.0 (current)
       Built [  68.014s] (current)
     Parsing datafusion v55.1.0 (current)
      Parsed [   0.037s] (current)
    Building datafusion v55.1.0 (baseline)
       Built [  18.771s] (baseline)
     Parsing datafusion v55.1.0 (baseline)
      Parsed [   0.037s] (baseline)
    Checking datafusion v55.1.0 -> v55.1.0 (no change; assume patch)
     Checked [   0.656s] 229 checks: 229 pass, 31 skip
     Summary no semver update required
    Finished [  89.088s] datafusion
    Building datafusion-expr v55.1.0 (current)
       Built [  23.344s] (current)
     Parsing datafusion-expr v55.1.0 (current)
      Parsed [   0.110s] (current)
    Building datafusion-expr v55.1.0 (baseline)
       Built [   7.734s] (baseline)
     Parsing datafusion-expr v55.1.0 (baseline)
      Parsed [   0.116s] (baseline)
    Checking datafusion-expr v55.1.0 -> v55.1.0 (no change; assume patch)
     Checked [   1.669s] 229 checks: 228 pass, 1 fail, 0 warn, 31 skip

--- failure constructible_struct_adds_field: struct exhaustively constructible through public API adds field ---

Description:
A pub struct that could be exhaustively constructed with a literal using only public API has a new pub field, breaking existing exhaustive literals.
        ref: https://doc.rust-lang.org/reference/expressions/struct-expr.html
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.51.0/src/lints/constructible_struct_adds_field.ron

Failed in:
  field MaterializedCteScan.projection in /home/runner/work/datafusion/datafusion/datafusion/expr/src/logical_plan/materialized_cte.rs:204
  field MaterializedCteScan.projection in /home/runner/work/datafusion/datafusion/datafusion/expr/src/logical_plan/materialized_cte.rs:204
  field MaterializedCteScan.projection in /home/runner/work/datafusion/datafusion/datafusion/expr/src/logical_plan/materialized_cte.rs:204
  field MaterializedCteScan.projection in /home/runner/work/datafusion/datafusion/datafusion/expr/src/logical_plan/materialized_cte.rs:204

     Summary semver requires new major version: 1 major and 0 minor checks failed
    Finished [  34.118s] datafusion-expr
    Building datafusion-optimizer v55.1.0 (current)
       Built [  17.779s] (current)
     Parsing datafusion-optimizer v55.1.0 (current)
      Parsed [   0.036s] (current)
    Building datafusion-optimizer v55.1.0 (baseline)
       Built [   7.749s] (baseline)
     Parsing datafusion-optimizer v55.1.0 (baseline)
      Parsed [   0.037s] (baseline)
    Checking datafusion-optimizer v55.1.0 -> v55.1.0 (no change; assume patch)
     Checked [   0.172s] 229 checks: 229 pass, 31 skip
     Summary no semver update required
    Finished [  26.510s] datafusion-optimizer
    Building datafusion-physical-plan v55.1.0 (current)
       Built [  17.543s] (current)
     Parsing datafusion-physical-plan v55.1.0 (current)
      Parsed [   0.238s] (current)
    Building datafusion-physical-plan v55.1.0 (baseline)
       Built [  13.777s] (baseline)
     Parsing datafusion-physical-plan v55.1.0 (baseline)
      Parsed [   0.240s] (baseline)
    Checking datafusion-physical-plan v55.1.0 -> v55.1.0 (no change; assume patch)
     Checked [   0.684s] 229 checks: 229 pass, 31 skip
     Summary no semver update required
    Finished [  33.551s] datafusion-physical-plan

[codecov-commenter]

Codecov Report

❌ Patch coverage is 81.05850% with 136 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.88%. Comparing base (484183d) to head (108276b).
⚠️ Report is 161 commits behind head on main.

Files with missing lines	Patch %	Lines
...tafusion/expr/src/logical_plan/materialized_cte.rs	83.50%	20 Missing and 28 partials ⚠️
...tafusion/optimizer/src/optimize_projections/mod.rs	78.35%	6 Missing and 23 partials ⚠️
datafusion/physical-plan/src/materialized_cte.rs	85.41%	4 Missing and 10 partials ⚠️
datafusion/expr/src/logical_plan/plan.rs	65.78%	11 Missing and 2 partials ⚠️
datafusion/expr/src/logical_plan/display.rs	0.00%	8 Missing ⚠️
datafusion/optimizer/src/analyzer/type_coercion.rs	83.72%	4 Missing and 3 partials ⚠️
datafusion/core/src/physical_planner.rs	86.36%	2 Missing and 4 partials ⚠️
...ion/substrait/src/logical_plan/producer/rel/mod.rs	0.00%	4 Missing ⚠️
datafusion/optimizer/src/decorrelate.rs	76.92%	1 Missing and 2 partials ⚠️
datafusion/optimizer/src/optimizer.rs	60.00%	0 Missing and 2 partials ⚠️
... and 2 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #26274      +/-   ##
==========================================
- Coverage   82.89%   82.88%   -0.02%     
==========================================
  Files        1154     1154              
  Lines      457630   458101     +471     
  Branches   457630   458101     +471     
==========================================
+ Hits       379370   379712     +342     
- Misses      55256    55310      +54     
- Partials    23004    23079      +75     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.