make dynamic filters in `ExecutionPlan` nodes discoverable

[jayshrivastava]

Is your feature request related to a problem or challenge?

Disclaimer: I'm writing this from the perspective of datafusion-distributed, but I imagine that this problem is common to other projects (ballista, comet) that distribute datafusion queries. I am not familiar with comet or ballista, so I would appreciate any feedback from folks who work on those projects.

More context in the (WIP) Dynamic Filtering in DataFusion-Distributed RFC

Background

Datafusion natively preserves dynamic filter producer<->consumer relationships across network boundaries only when you serialize the producer (HashJoinExec) and the consumer (DataSourceExec) in serialization run. There's tests which assert that this works.

                   ┌───────────────────────────────────────────────┐                                                
                   │          ┌───────────────────────┐            │                                                
                   │          │    RepartitionExec    │            │                                                
                   │          │                       │            │                                                
                   │          └───────────────────────┘            │                                                
                   │          ┌───────────────────────┐            │                                                
                   │          │     HashJoinExec      │            │                                                
 ┌───────────────┐ │          │  mode = Partitioned   │            │                                                
 │   Machine 1   │ │          └───────────────────────┘            │                                                
 └───────────────┘ │ ┌───────────────────┐  ┌───────────────────┐  │                                                
                   │ │  DataSourceExec   │  │  ProjectionExec   │  │                                                
                   │ └───────────────────┘  └───────────────────┘  │                                                
                   │                        ┌───────────────────┐  │                                                
                   │                        │  DataSourceExec   │  │                                                
                   │                        └───────────────────┘  │                                                
                   └───────────────────────────────────────────────┘                                                
                                           │                                                                        
                                           │    machine 1 serializes the plan                                       
                                           │    machine 2 deserializes the plan and preserves dynamic filter        
                                           │    relationships. dynamic filtering works                                                       
                                           ▼                                                                        
                   ┌───────────────────────────────────────────────┐                                                
                   │           ┌───────────────────────┐           │                                                
                   │           │    RepartitionExec    │           │                                                
                   │           │                       │           │                                                
                   │           └───────────────────────┘           │                                                
                   │           ┌───────────────────────┐           │                                                
                   │           │     HashJoinExec      │           │                                                
                   │           │  mode = Partitioned   │           │                                                
┌───────────────┐  │           └───────────────────────┘           │                                                
│   Machine 2   │  │  ┌───────────────────┐  ┌───────────────────┐ │                                                
└───────────────┘  │  │  DataSourceExec   │  │  ProjectionExec   │ │                                                
                   │  └───────────────────┘  └───────────────────┘ │                                                
                   │                         ┌───────────────────┐ │                                                
                   │                         │  DataSourceExec   │ │                                                
                   │                         └───────────────────┘ │                                                
                   └───────────────────────────────────────────────┘                                                

The problem is that datafusion doesn't support the case where you split up the consumer from the producer:

                   ┌───────────────────────────────────────────────┐                       
                   │          ┌───────────────────────┐            │                       
                   │          │    RepartitionExec    │            │                       
                   │          │                       │            │                       
                   │          └───────────────────────┘            │                       
                   │          ┌───────────────────────┐            │                       
                   │          │     HashJoinExec      │            │                       
 ┌───────────────┐ │          │  mode = Partitioned   │            │                       
 │   Machine 1   │ │          └───────────────────────┘            │                       
 └───────────────┘ │ ┌───────────────────┐  ┌───────────────────┐  │                       
                   │ │  DataSourceExec   │  │  ProjectionExec   │  │                       
                   │ └───────────────────┘  └───────────────────┘  │                       
                   │                        ┌───────────────────┐  │                       
                   │                        │  DataSourceExec   │  │                       
                   │                        └───────────────────┘  │                       
                   └───────────────────────────────────────────────┘                       
                                           │                                               
                                           │                                               
                                           │    machine 1 sends the probe side to machine 2
                                           │    how will the hash join send updates to the data source?                                           
                                           ▼                                               
                   ┌───────────────────────────────────────────────┐                       
                   │                                               │                       
                   │                                               │                       
                   │                                               │                       
                   │                                               │                       
                   │                                               │                       
                   │                                               │                       
                   │                                               │                       
┌───────────────┐  │                                               │                       
│   Machine 2   │  │                         ┌───────────────────┐ │                       
└───────────────┘  │                         │  ProjectionExec   │ │                       
                   │                         └───────────────────┘ │                       
                   │                         ┌───────────────────┐ │                       
                   │                         │  DataSourceExec   │ │                       
                   │                         └───────────────────┘ │                       
                   └───────────────────────────────────────────────┘                                                                                                                                                                                                                             

Describe the solution you'd like

Proposal:

1. Distributed projects (ballista, comet, datafusion-distributed) should be responsible for implementing dynamic filter update routing and propagation during query execution
2. Vanilla datafusion should provide a way to find and extract dynamic filters from ExecutionPlan to allow a project like datafusion-distributed to implement (1). The minimal required features are
   (a) identify which dynamic filters are producers and consumers in an ExecutionPlan tree
   (b) get access to all &DynamicFilterPhysicalExpr in an ExecutionPlan (as long as one is holding a reference to a DynamicFilterPhysicalExpr, they can read and write updates to it)
   (c) allow easy extension for custom ExecutionPlan implementations

Option 1 - New ExecutionPlan API

trait ExecutionPlan {
    fn dynamic_filter_exprs(&self) -> Vec<DynamicFilterNodeBehavior> 
}

enum DynamicFilterNodeBehavior {
    Producer(Arc<dyn PhysicalExpr>),
    Consumer(Arc<dyn PhysicalExpr>),
}

Notes

• Previously, a similar API was reverted in Revert "Add ExecutionPlan::apply_expressions() (#20337)" #22437 because implementing it was complex and there was no use case in vanilla datafusion to assert correctness
  • The proposed dynamic_filter_exprs() API is simpler to implement, making the breaking change less annoying
  • Still no use case in vanilla datafusion

Pros:

• Does exactly what we want

Cons:

• Breaking change. It must be implemented by every ExecutionPlan. We cannot default to vec![] because ExecutionPlan nodes may fail to declare their dynamic filters.
• Very specific. We usually treat dynamic filters as any other filter expression. However, this method cuts through the abstraction and brings them to the forefront.

Option 2 - Store Dynamic Filters in the Session (SessionState)

We can try updating the SessionState to store which plan nodes are producers and consumers of dynamic filters. This can be populated during the FilterPushdown optimization rule using the existing API:

trait PhysicalOptimizerRule {
    fn optimize_with_context(
        &self,
        plan: Arc<dyn ExecutionPlan>,
        context: &dyn PhysicalOptimizerContext,
    ) -> Result<Arc<dyn ExecutionPlan>>
}

The PhysicalOptimizerContext only stores a statistics registry today. We would need to extend it to store dynamic filters:

trait PhysicalOptimizerContext {
      fn runtime_filter_registry(
          &self,
      )
  }

Here's a rough sketch of what we would store in the Session.

struct FilterPushdownReport {
    bindings: HashMap<(u64, u64), Vec<RuntimeFilterBinding>>, // (expression_id, query_id) -> bindings
}

enum RuntimeFilterBinding {
    Producer {
        node: Arc<dyn ExecutionPlan>,
        expression: Arc<dyn PhysicalExpr>,
    },
    Consumer {
        node: Arc<dyn ExecutionPlan>,
        expression: Arc<dyn PhysicalExpr>,
        usage: FilterUsage, // Pruning or FullyEvaluated
    },
}

Pros

• Not a breaking change

Cons

• We have to deal with a state which spans multiple queries
• Arc<dyn PhysicalExpr> references may grow stale. datafusion-distributed certainly modifies the ExecutionPlan tree and could invalidate these pointers.
  • FilterPushdownPhase::Post says it should be run at the end of the optimization process since any changes to the plan may break the dynamic filter’s references.
  • We could make it the distributed project's responsibility to mutate / update the FilterPushdownReport if it changes any ExecutionPlans, but that would still rely on pointer comparisons.

Open questions

• Is it acceptable to add features in vanilla datafusion which aren't used by vanilla datafusion? Who is responsible for testing these? (Related: Add an example of an in-process model of distributed execution #23282).
  • It would be nice if vanilla datafusion committed to a minimal feature (ie. "DynamicFilterPhysicalExpr can be discovered from ExecutionPlan") with minimal tests support use cases like datafusion-distributed, ballista, and comet.

Describe alternatives you've considered

We could downcast ExecutionPlan to concrete types and try to discover dynamic filters like that, however this method is less extensible (how would we support custom ExecutionPlan implementations?) and it isn't guaranteed to work because not every concrete type will make their dynamic filters pub.

As of writing, these public methods on some execution plan nodes exist, byt they were only added to support serialization and will likely be removed by #23494.

HashJoinExec::dynamic_filter_expr()
AggregateExec::dynamic_filter_expr()
SortExec::dynamic_filter_expr()

Additional context

No response

[jayshrivastava]

cc @alamb @timsaucer @andygrove @gene-bordegaray @gabotechs @adriangb Regarding our discussion during the community meeting today.

[andygrove]

also cc @phillipleblanc @milenkovicm @avantgardnerio @martin-g from Ballista community

[andygrove]

Thanks for creating this issue @jayshrivastava, and thanks for linking to #23282.

My opinion (which I know may not be shared by others) is that it would be beneficial for "core" DataFusion to offer a choice of execution modes:

1. Eager execution of the whole query tree, streaming results immediately (current model)
2. A distributed-style execution, but without the distributed (networking, multi-process) part. Essentially this would model dfd/ballista style execution and allow for development of related capabilities in core. I would be fine with this being an optional crate. This was the idea I was trying to demonstrate in Add an example of an in-process model of distributed execution #23282

Expanding further on this, enabling a distributed-style execution path with optional persistence of data between stages could allow single process DF to scale more on certain workloads (even though performance may not be amazing).

[adriangb]

Breaking change. It must be implemented by every ExecutionPlan. We cannot default to vec![] because ExecutionPlan nodes may fail to declare their dynamic filters.

I think what made the previous implementation tricky was that it required declaring all filter expressions and the API was so generic that a reasonable consumer could rely on it in ways that would break if it was not implemented properly. This more narrow version is at least safer if not safe in that sense: dynamic filters are generally optional performance improvements. Their absence or mishandling doesn't result in incorrect results or errors, just in performance degradation. So maybe defaulting to vec![] is a viable option.

For DynamicFilterNodeBehavior and the Producer/Consumer split: is it necessary to know the relationship between nodes? I.e. if I have a SortExec and a HashJoinExec they will both declare themselves as producers, then there will be 1 or maybe 2 DataSourceExec that are consumers. But if the plan is being split up and distributed amongst nodes I imagine it would be necessary or at last helpful to know which producers correspond to which consumers? If not you'd have to broadcast updates from all producers to all consumers.

[jayshrivastava]

But if the plan is being split up and distributed amongst nodes I imagine it would be necessary or at last helpful to know which producers correspond to which consumers? If not you'd have to broadcast updates from all producers to all consumers.

Agreed. I was looking to use expression_id for this. The expression_id helps identify pairs that are linked together, but it doesn't make it clear in which direction the updates need to go.

[jayshrivastava]

We could infer from context (ex. HashJoinExec is always a producer), but then there's the question of custom ExecutionPlan implementations. I figured we'd want them to declare themselves to be producers or consumers.

[andygrove]

We could infer from context (ex. HashJoinExec is always a producer), but then there's the question of custom ExecutionPlan implementations. I figured we'd want them to declare themselves to be producers or consumers.

Exactly. There is a draft PR in Ballista to add a custom hash join operator (with spilling). We would want DataFusion to treat it as a HashJoinExec even though it is a custom implementation.

[avantgardnerio]

I am not able to give it my full attention yet, and I apologize for that. But I think Impala does this:

  PLAN-ROOT SINK                                                                                                                                                                                                                                           
  |                                                                                                                                                                                                                                                        
  04:EXCHANGE [UNPARTITIONED]                                                                                                                                                                                                                              
  |                                                                                                                                                                                                                                                        
  02:HASH JOIN [LEFT SEMI JOIN, BROADCAST]
  |  hash predicates: year = year                                                                                                                                                                                                                          
  |  runtime filters: RF000 <- year                    <-- producer                                                                                                                                                                                        
  |                                                                                                                                                                                                                                                        
  |--03:EXCHANGE [BROADCAST]
  |  |
  |  01:SCAN HDFS [default.yy2]
  |     partitions=1/1 files=1 size=620B
  |
  00:SCAN HDFS [default.yy]
     partitions=5/5 files=5 size=1.71KB
     runtime filters: RF000 -> year                    <-- consumer

I think the only change we would need to track things is:

1. add 2 methods to ExecutionPlanNode: get_sender_ids() -> Vec<usize> and get_receiver_ids()
2. have the rules that connect 2 operators to increment and assign those counters to ops

Then any further optimizer / planner could simply inspect those values to find connections between operators that (may) span stage boundaries, and have an option to handle them as it sees fit (reset the plan, skip the rule, fail, etc)

https://impala.apache.org/docs/build/html/topics/impala_runtime_filtering.html

[alamb]

1. Distributed projects (ballista, comet, datafusion-distributed) should be responsible for implementing dynamic filter update routing and propagation during query execution

Agree (I also agree with your other points)

I think you have three usecases:

1. A way to discover all dynamic filters in a aplan
2. A way to discover, for each ExecutionPlan, which dynamic filter it produces
3. A way to discover, for each ExecutionPlan, which dynamic filter it consumes

I think a "consumer" is a consumer simply by nature of evaluating an expression tree that happens to have DynamicFilterPhysicalExpr somewhere.

There is no explicit way to find out what nodes are "producer"s today from what I can tell -- the only thing that makes an ExecutionPlan a producer is if it pushes dynamic filter expressions down in FilterPushdownPhase::post

Note that DynamicFilterPhysicalExpr already implements PhysicalExpr::expression_id so there is a unique way to identify each expression

To satisfy the usescase then I think the best set of APIs are:

1. A generic API to walk and modify expressions on ExecutionPlan similar to LogicalPlan::map_expressions and LogicalPlan::apply_expressions (basically what @LiaCastaneda added in Add ExecutionPlan::apply_expressions() #20337 )
2. Add an API to ExecutionPlan to return what dynamic filters it produced (basically promote existing methods like SortExec::dynamic_filters and HashJoinExec::dynamic_filters to the ExecutionPlan trait)

trait ExecutionPlan
fn dynamic_filters(&self) -> Vec<Arc<DynamicFilterPhysicalExpr>> {
  Vec::new()
}

I think adding ExecutionPlan::dynamic_filters is better than an enum with producer/consumer as it is clear it is only about what filters are produced (not consumed)

[alamb]

I think returning DynamicFilterNodeBehavior:: Consumer will require the equivalent of apply_expressions anyways (to visit all expressions in a node).

[alamb]

Add an API to ExecutionPlan to return what dynamic filters it produced (basically promote existing methods like SortExec::dynamic_filters and HashJoinExec::dynamic_filters to the ExecutionPlan trait)

FWIW this would be almost free today (as the method already exists on several of the nodes) and gives you the "producer" side immediately with minimal downstream impact