Avoid recompute CTEs (common table expressions) / share input plans

[alamb]

Is your feature request related to a problem or challenge?

The core usecase is:

with x as (
    <expensive full join of two large tables producing small-ish result>
)

select * from x where ...
union all
select * from x where ...
union all
select * from x where ...

DataFusion will effectively run the subquery x three times (it will basically copy the LogicalPlan for x wherever it is used.

                                       ┌─────────────────────┐                                     
                                       │      UNION ALL      │                                     
                                       │                     │                                     
                                       └─────────────────────┘                                     
                                             ▲    ▲   ▲                                            
                                             │    │   │                                            
               ┌─────────────────────────────┘    │   └─────────────────────────────┐              
               │                                  │                                 │              
               │                                  │                                 │              
      ┌────────────────┐                 ┌────────────────┐                ┌────────────────┐      
      │    Filter 1    │                 │    Filter 2    │                │    Filter 3    │      
      └────────────────┘                 └────────────────┘                └────────────────┘      
               ▲                                  ▲                                 ▲              
               │                                  │                                 │              
               │                                  │                                 │              
 ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐      ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐     ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐
      ┌────────────────┐                 ┌────────────────┐                ┌────────────────┐      
 │    │ Expensive Join │     │      │    │ Expensive Join │     │     │    │ Expensive Join │     │
      └────────────────┘                 └────────────────┘                └────────────────┘      
 │             ▲             │      │             ▲             │     │             ▲             │
        ┌──────┴──────┐                    ┌──────┴──────┐                   ┌──────┴──────┐       
 │      │             │      │      │      │             │      │     │      │             │      │
    .───────.     .───────.            .───────.     .───────.           .───────.     .───────.   
 │ ╱         ╲   ╱         ╲ │      │ ╱         ╲   ╱         ╲ │     │ ╱         ╲   ╱         ╲ │
  (  Input 1  ) (  Input 2  )        (  Input 1  ) (  Input 2  )       (  Input 1  ) (  Input 2  ) 
 │ `.       ,'   `.       ,' │      │ `.       ,'   `.       ,' │     │ `.       ,'   `.       ,' │
     `─────'       `─────'              `─────'       `─────'             `─────'       `─────'    
 │"x"                        │      │"x"                        │     │"x"                        │
  ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─        ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─       ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ 

This design has certain benefits:

1. It is straightforward to implement (probably why DataFusion is like this)
2. If the different UNION ALL arms have different predicates, they could potentially be pushed down in one branch but not the others.

Describe the solution you'd like

However, in many cases it would likely be better to do to the expensive join only once and reuse the results like this:

                     ┌─────────────────────┐                     
                     │      UNION ALL      │                     
                     │                     │                     
                     └─────────────────────┘                     
                           ▲    ▲   ▲                            
                           │    │   │                            
         ┌─────────────────┘    │   └───────────────────┐        
         │                      │                       │        
         │                      │                       │        
         │                      │                       │        
┌────────────────┐     ┌────────────────┐      ┌────────────────┐
│    Filter 1    │     │    Filter 2    │      │    Filter 3    │
└────────────────┘     └────────────────┘      └────────────────┘
         ▲                      ▲                       ▲        
         │                      │                       │        
         │                      │                       │        
         └────────────────────┐ │ ┌─────────────────────┘        
                              │ │ │                              
                              │ │ │                              
                              │ │ │                              
                  ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐                  
                       ┌────────────────┐                        
                  │    │ Expensive Join │     │                  
                       └────────────────┘                        
                  │             ▲             │                  
                         ┌──────┴──────┐                         
                  │      │             │      │                  
                     .───────.     .───────.                     
                  │ ╱         ╲   ╱         ╲ │                  
                   (  Input 1  ) (  Input 2  )                   
                  │ `.       ,'   `.       ,' │                  
                      `─────'       `─────'                      
                  │"x"                        │                  
                   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─                   

Describe alternatives you've considered

I think there are several considerations for this design, the biggest is that it is a 'diamond' plan where the same stream can be consumed at different rates potentially needing to buffer the entire intermediate result or else the plan will deadlock

For example

                 ┌─────────────────────┐                     
                 │      Hash Join      │                     
                 │                     │                     
                 └─────────────────────┘                     
                       ▲          ▲                          
                       │          │                          
         ┌─────────────┘          └────────────┐             
         │                                     │             
         │                                     │             
         │                                     │             
    Build Side                            Probe Side         
 (read completely                (not read at all until Build
before probe side)                 Side is completely read)  
         ▲                                     ▲             
         │                                     │             
         │                                     │             
         └────────────────────┬────────────────┘             
                              │                              
                              │                              
                              │                              
                  ┌ ─ ─ ─ ─ ─ ┴ ─ ─ ─ ─ ─ ─ ─ ┐              
                       ┌────────────────┐                    
                  │    │ Expensive Join │     │              
                       └────────────────┘                    
                  │             ▲             │              
                         ┌──────┴──────┐                     
                  │      │             │      │              
                     .───────.     .───────.                 
                  │ ╱         ╲   ╱         ╲ │              
                   (  Input 1  ) (  Input 2  )               
                  │ `.       ,'   `.       ,' │              
                      `─────'       `─────'                  
                  │"x"                        │              
                   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─               

Additional context

This came from a discord thread from @sergiimk

[alamb]

This also came up as potentially part of a solution for #8582

[Dandandan]

This is an interesting topic / improvement!

Two suggestions from my side

• It would be very nice if this could be generalized to equivalent subqueries (i.e. already expanded CTEs or subqueries that have the same result) instead of CTEs.
• PostgreSQL has the hint [NOT] MATERIALIZED to force / disable CTE materialization. Often it's much faster to disable materialization, as otherwise pushdown / CBO optimizations will not apply (as you suggested). It might be interesting to consider supporting this in DataFusion as well, as it is hard from an optimizer standpoint to decide what to do.
  https://www.postgresql.org/docs/current/queries-with.html#QUERIES-WITH-CTE-MATERIALIZATION

[crepererum]

where the same stream can be consumed at different rates potentially needing to buffer the entire intermediate result or else the plan will deadlock

I think this is very similar to our repartitioning code and the trade-offs and problems we see there. The reason is that a repartition is basically also as single input with mulitiple consumers. Just think of it like the same data but with a column "bool: belongs to this output" added.

[suibianwanwank]

Hi, I’d like to take a try at this task.

My plan is to first support CTE with the WITH ... AS MATERIALIZED syntax. After that, we can explore broader optimizations in the optimizer phase.

As for the implementation of CTE in the ExecutionPlan, I currently have two ideas:

1. Add a top-level CTEQueryExec node to ensure the CTE query is fully executed (blocking), and cache its result in a CTEWorkTable so that subsequent CTEScan nodes can read from it. (It seems DuckDB takes a similar approach.)
   However, I'm not sure if this is feasible in DataFusion, as it might require blocking the ExecutionPlan and collecting all RecordBatches.

2. Materialize the CTE before execution by creating a temporary table in the current schema and writing the required data into it. Then, replace all following scan with queries against this temporary table. (I noticed that Databend uses a similar approach.)

I’d like to gather some feedback first to make sure the design is general enough and correct. Thanks!

[findepi]

Maintaining streaming processing is very useful in the number of circumstances: a query with LIMIT, a query with TopN over sorted input data, an interactive query.
I don't know how CTE reuse execution v1 should look like, but we should work towards the goal that all plan nodes dependent on the CTE can incrementally consume data passed to the CTE. Effectively, each should maintain an independent iterator over data being produced. Additionally, CTE should be aware of number of upstream consumers, which allows it to free data no longer needed by any consumer.

Implementation wise, it could be done with a separate blocking deque (VecDeque?) for each consumer.
Incoming data is added to all queues (shared Arcs, no data copying).
Consumers pulls from the queue. When it's empty, it waits for new data or end of processing signal.

[crepererum]

Implementation wise, it could be done with a separate blocking deque (VecDeque?) for each consumer.
Incoming data is added to all queues (shared Arcs, no data copying).
Consumers pulls from the queue. When it's empty, it waits for new data or end of processing signal.

I think this problem is actually very similar to RepatitionExec and the distributor channels in there, except that RepartitionExec sends 1 row to 1 output, the construct required here would just send them to ALL outputs. And we'll likely want a similar buffer limiter, so that the CTE output isn't buffered unlimited when no consumer polls the data (at the same time, multiple consumer may poll at the same time). So I would propose to use the same underlying "distributor channel" primitive (or whatever we replace that with).