Skip to content

Commit 4ad354c

Browse files
jayshrivastavajayshrivastava
committed
checkpoint
1 parent a4a9541 commit 4ad354c

11 files changed

Lines changed: 170 additions & 100 deletions

File tree

‎src/distributed_ext.rs‎

Lines changed: 33 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use crate::config_extension_ext::{
33
set_distributed_option_extension, set_distributed_option_extension_from_headers,
44
};
55
use crate::events::{
6-
DesiredTaskCountHandler, DesiredTaskCountHandlers, RouteTasksHandler, RouteTasksHandlers,
6+
Event, EventHandlerChainGeneric, Handler, RouteTasksHandler, RouteTasksHandlers,
77
ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers,
88
};
99
use crate::passthrough_headers::set_passthrough_headers;
@@ -575,6 +575,7 @@ pub trait DistributedExt: Sized {
575575
/// in-place mutation.
576576
fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext);
577577

578+
/// TODO:
578579
/// Registers a handler that supplies desired or maximum task-count hints for plan nodes.
579580
/// The distributed planner reconciles these hints when choosing each stage's task count.
580581
///
@@ -591,25 +592,14 @@ pub trait DistributedExt: Sized {
591592
/// }
592593
///
593594
/// SessionStateBuilder::new()
594-
/// .with_distributed_desired_task_count_handler(handle_custom_desired_task_count);
595+
/// .with_distributed_event_handler(handle_custom_desired_task_count);
595596
/// ```
596597
///
597598
/// ```text
598599
/// ┌────────────────┐
599600
/// │CustomDataSource│──────────▶ 3 desired tasks
600601
/// └────────────────┘
601602
/// ```
602-
fn with_distributed_desired_task_count_handler<T: DesiredTaskCountHandler>(
603-
self,
604-
handler: T,
605-
) -> Self;
606-
607-
/// Same as [DistributedExt::with_distributed_desired_task_count_handler] but with an
608-
/// in-place mutation.
609-
fn set_distributed_desired_task_count_handler<T: DesiredTaskCountHandler>(
610-
&mut self,
611-
handler: T,
612-
);
613603
614604
/// Registers a handler that can replace leaf nodes with distributed variants once the
615605
/// distributed planner has decided a final task count.
@@ -690,6 +680,12 @@ pub trait DistributedExt: Sized {
690680
/// Same as [DistributedExt::with_distributed_route_tasks_handler] but with an in-place
691681
/// mutation.
692682
fn set_distributed_route_tasks_handler<T: RouteTasksHandler>(&mut self, handler: T);
683+
684+
// Sets a handler for the event type E
685+
//
686+
687+
fn with_distributed_event_handler<E: Event, H: Handler<E>>(self, handler: H) -> Self;
688+
fn set_distributed_event_handler<E: Event, H: Handler<E>>(&mut self, handler: H);
693689
}
694690

695691
/// Trait to have a unified interface for getting structs & properties from SessionConfig that are used in distributed context.
@@ -854,10 +850,6 @@ impl DistributedExt for SessionConfig {
854850
self.set_extension(Arc::new(local_worker_context));
855851
}
856852

857-
fn set_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(&mut self, h: H) {
858-
DesiredTaskCountHandlers::push_custom(self, Arc::new(h))
859-
}
860-
861853
fn set_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(&mut self, h: H) {
862854
ScaleUpLeafNodeHandlers::push_custom(self, Arc::new(h));
863855
}
@@ -866,6 +858,10 @@ impl DistributedExt for SessionConfig {
866858
RouteTasksHandlers::push_custom(self, Arc::new(h));
867859
}
868860

861+
fn set_distributed_event_handler<E: Event, H: Handler<E>>(&mut self, h: H) {
862+
EventHandlerChainGeneric::<E>::push_custom(self, Arc::new(h));
863+
}
864+
869865
delegate! {
870866
to self {
871867
#[call(set_distributed_option_extension)]
@@ -958,20 +954,21 @@ impl DistributedExt for SessionConfig {
958954
#[expr($;self)]
959955
fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self;
960956

961-
#[call(set_distributed_desired_task_count_handler)]
962-
#[expr($;self)]
963-
fn with_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(mut self, h: H) -> Self;
964-
965957
#[call(set_distributed_scale_up_leaf_node_handler)]
966958
#[expr($;self)]
967959
fn with_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(mut self, h: H) -> Self;
968960

969961
#[call(set_distributed_route_tasks_handler)]
970962
#[expr($;self)]
971963
fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(mut self, h: H) -> Self;
964+
965+
#[call(set_distributed_event_handler)]
966+
#[expr($;self)]
967+
fn with_distributed_event_handler<E: Event, H: Handler<E>>(mut self, h: H) -> Self;
972968
}
973969
}
974970
}
971+
975972
impl DistributedGetterExt for SessionConfig {
976973
fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError> {
977974
get_distributed_worker_resolver(self)
@@ -1098,11 +1095,6 @@ impl DistributedExt for SessionStateBuilder {
10981095
#[expr($;self)]
10991096
fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self;
11001097

1101-
fn set_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(&mut self, h: H);
1102-
#[call(set_distributed_desired_task_count_handler)]
1103-
#[expr($;self)]
1104-
fn with_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(mut self, h: H) -> Self;
1105-
11061098
fn set_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(&mut self, h: H);
11071099
#[call(set_distributed_scale_up_leaf_node_handler)]
11081100
#[expr($;self)]
@@ -1112,6 +1104,11 @@ impl DistributedExt for SessionStateBuilder {
11121104
#[call(set_distributed_route_tasks_handler)]
11131105
#[expr($;self)]
11141106
fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(mut self, h: H) -> Self;
1107+
1108+
fn set_distributed_event_handler<E: Event, H: Handler<E>>(&mut self, h: H);
1109+
#[call(set_distributed_event_handler)]
1110+
#[expr($;self)]
1111+
fn with_distributed_event_handler<E: Event, H: Handler<E>>(mut self, h: H) -> Self;
11151112
}
11161113
}
11171114
}
@@ -1243,11 +1240,6 @@ impl DistributedExt for SessionState {
12431240
#[expr($;self)]
12441241
fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self;
12451242

1246-
fn set_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(&mut self, h: H);
1247-
#[call(set_distributed_desired_task_count_handler)]
1248-
#[expr($;self)]
1249-
fn with_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(mut self, h: H) -> Self;
1250-
12511243
fn set_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(&mut self, h: H);
12521244
#[call(set_distributed_scale_up_leaf_node_handler)]
12531245
#[expr($;self)]
@@ -1257,6 +1249,11 @@ impl DistributedExt for SessionState {
12571249
#[call(set_distributed_route_tasks_handler)]
12581250
#[expr($;self)]
12591251
fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(mut self, h: H) -> Self;
1252+
1253+
fn set_distributed_event_handler<E: Event, H: Handler<E>>(&mut self, h: H);
1254+
#[call(set_distributed_event_handler)]
1255+
#[expr($;self)]
1256+
fn with_distributed_event_handler<E: Event, H: Handler<E>>(mut self, h: H) -> Self;
12601257
}
12611258
}
12621259
}
@@ -1381,11 +1378,6 @@ impl DistributedExt for SessionContext {
13811378
#[expr($;self)]
13821379
fn with_distributed_local_worker_context(self, local_worker_context: LocalWorkerContext) -> Self;
13831380

1384-
fn set_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(&mut self, h: H);
1385-
#[call(set_distributed_desired_task_count_handler)]
1386-
#[expr($;self)]
1387-
fn with_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(self, h: H) -> Self;
1388-
13891381
fn set_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(&mut self, h: H);
13901382
#[call(set_distributed_scale_up_leaf_node_handler)]
13911383
#[expr($;self)]
@@ -1395,6 +1387,11 @@ impl DistributedExt for SessionContext {
13951387
#[call(set_distributed_route_tasks_handler)]
13961388
#[expr($;self)]
13971389
fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(self, h: H) -> Self;
1390+
1391+
fn set_distributed_event_handler<E: Event, H: Handler<E>>(&mut self, h: H);
1392+
#[call(set_distributed_event_handler)]
1393+
#[expr($;self)]
1394+
fn with_distributed_event_handler<E: Event, H: Handler<E>>(self, h: H) -> Self;
13981395
}
13991396
}
14001397
}

‎src/distributed_planner/inject_network_boundaries.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::events::TaskCountAnnotation::{Desired, Maximum};
22
use crate::events::{
3-
DesiredTaskCountEvent, DesiredTaskCountHandlers, ScaleUpLeafNodeEvent, ScaleUpLeafNodeHandlers,
3+
DesiredTaskCount, DesiredTaskCountEvent, ScaleUpLeafNodeEvent, ScaleUpLeafNodeHandlers,
44
TaskCountAnnotation,
55
};
66
use crate::execution_plans::{ChildWeight, ChildrenIsolatorUnionExec};
@@ -238,7 +238,7 @@ async fn _inject_network_boundaries(
238238
plan: &plan,
239239
session_config: nb_ctx.cfg,
240240
};
241-
return if let Some(estimate) = DesiredTaskCountHandlers::handle(ev) {
241+
return if let Some(estimate) = DesiredTaskCount::handle(ev) {
242242
Ok(nb_ctx.plan_with_task_count(plan, estimate.task_count.limit(nb_ctx.max_tasks()?)))
243243
} else {
244244
// We could not determine how many tasks this leaf node should run on, so
@@ -262,7 +262,7 @@ async fn _inject_network_boundaries(
262262
plan: &plan,
263263
session_config: nb_ctx.cfg,
264264
};
265-
let mut task_count = DesiredTaskCountHandlers::handle(ev).map_or(Desired(1), |v| v.task_count);
265+
let mut task_count = DesiredTaskCount::handle(ev).map_or(Desired(1), |v| v.task_count);
266266
if nb_ctx.d_cfg.children_isolator_unions && plan.is::<UnionExec>() {
267267
// Unions have the chance to decide how many tasks they should run on. If there's a union
268268
// with a bunch of children, the user might want to increase parallelism and increase the

‎src/distributed_planner/session_state_builder_ext.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use crate::distributed_planner::DistributedConfig;
22
use crate::distributed_planner::distributed_query_planner::DistributedQueryPlanner;
33
use crate::events::{
4-
DesiredTaskCountHandlers, ScaleUpLeafNodeHandlers, file_scan_config_desired_task_count,
4+
EventHandlerChainGeneric, ScaleUpLeafNodeHandlers, file_scan_config_desired_task_count,
55
file_scan_config_scale_up_leaf_node,
66
};
77
use datafusion::execution::SessionStateBuilder;
@@ -27,7 +27,7 @@ impl SessionStateBuilderExt for SessionStateBuilder {
2727
.optimizer
2828
.enable_physical_uncorrelated_scalar_subquery = false;
2929

30-
DesiredTaskCountHandlers::push_builtin(cfg, Arc::new(file_scan_config_desired_task_count));
30+
EventHandlerChainGeneric::push_builtin(cfg, Arc::new(file_scan_config_desired_task_count));
3131
ScaleUpLeafNodeHandlers::push_builtin(cfg, Arc::new(file_scan_config_scale_up_leaf_node));
3232
let prev = std::mem::take(self.query_planner());
3333
self.with_query_planner(Arc::new(DistributedQueryPlanner { prev }))

‎src/events/common.rs‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,69 @@
11
use datafusion::execution::config::SessionConfig;
22
use std::sync::Arc;
33

4+
pub trait Handler<E: Event>: Send + Sync + 'static {
5+
fn handle(&self, ev: E::Data<'_>) -> E::Response;
6+
}
7+
8+
impl<E, H> Handler<E> for Arc<H>
9+
where
10+
E: Event,
11+
H: Handler<E> + ?Sized,
12+
{
13+
fn handle(&self, ev: E::Data<'_>) -> E::Response {
14+
self.as_ref().handle(ev)
15+
}
16+
}
17+
18+
pub trait Event: Send + Sync + 'static {
19+
/// Data contained in the event.
20+
type Data<'a>;
21+
type Response;
22+
}
23+
24+
pub(crate) struct EventHandlerChainGeneric<E: Event> {
25+
pub(crate) builtin: Vec<Arc<dyn Handler<E>>>,
26+
pub(crate) custom: Vec<Arc<dyn Handler<E>>>,
27+
}
28+
29+
impl<E: Event> Default for EventHandlerChainGeneric<E> {
30+
fn default() -> Self {
31+
Self {
32+
builtin: Vec::new(),
33+
custom: Vec::new(),
34+
}
35+
}
36+
}
37+
38+
impl<E: Event> Clone for EventHandlerChainGeneric<E> {
39+
fn clone(&self) -> Self {
40+
Self {
41+
builtin: self.builtin.clone(),
42+
custom: self.custom.clone(),
43+
}
44+
}
45+
}
46+
47+
impl<E: Event> EventHandlerChainGeneric<E> {
48+
pub(crate) fn push_builtin(cfg: &mut SessionConfig, handler: Arc<dyn Handler<E>>) {
49+
let mut handlers = cfg
50+
.get_extension::<Self>()
51+
.map(|v| v.as_ref().clone())
52+
.unwrap_or_default();
53+
handlers.builtin.push(handler);
54+
cfg.set_extension(Arc::new(handlers));
55+
}
56+
57+
pub(crate) fn push_custom(cfg: &mut SessionConfig, handler: Arc<dyn Handler<E>>) {
58+
let mut handlers = cfg
59+
.get_extension::<Self>()
60+
.map(|v| v.as_ref().clone())
61+
.unwrap_or_default();
62+
handlers.custom.push(handler);
63+
cfg.set_extension(Arc::new(handlers));
64+
}
65+
}
66+
467
pub(crate) struct EventHandlerChain<H: ?Sized> {
568
builtin: Vec<Arc<H>>,
669
custom: Vec<Arc<H>>,

‎src/events/defaults/file_scan_config.rs‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
use crate::DistributedConfig;
2+
#[cfg(test)]
3+
use crate::events::DesiredTaskCount;
24
use crate::events::{
35
DesiredTaskCountEvent, DesiredTaskCountEventResponse, ScaleUpLeafNodeEvent,
46
ScaleUpLeafNodeEventResponse,
@@ -98,7 +100,6 @@ fn rebalance_round_robin<T>(items: Vec<T>, target_groups: usize) -> Vec<Vec<T>>
98100
mod tests {
99101
use super::*;
100102
use crate::DistributedExt;
101-
use crate::events::DesiredTaskCountHandlers;
102103
use crate::test_utils::parquet::register_parquet_tables;
103104
use datafusion::error::DataFusionError;
104105
use datafusion::physical_plan::ExecutionPlan;
@@ -107,11 +108,11 @@ mod tests {
107108
#[tokio::test]
108109
async fn test_first_desired_task_count_handler_wins() -> Result<(), DataFusionError> {
109110
let cfg = SessionConfig::new()
110-
.with_distributed_desired_task_count_handler(desired_ten)
111-
.with_distributed_desired_task_count_handler(desired_twenty);
111+
.with_distributed_event_handler(desired_ten)
112+
.with_distributed_event_handler(desired_twenty);
112113

113114
let plan = make_data_source_exec().await?;
114-
let response = DesiredTaskCountHandlers::handle(DesiredTaskCountEvent {
115+
let response = DesiredTaskCount::handle(DesiredTaskCountEvent {
115116
plan: &plan,
116117
session_config: &cfg,
117118
})
@@ -123,11 +124,11 @@ mod tests {
123124
#[tokio::test]
124125
async fn test_desired_task_count_handlers_continue_until_some() -> Result<(), DataFusionError> {
125126
let cfg = SessionConfig::new()
126-
.with_distributed_desired_task_count_handler(no_desired_task_count)
127-
.with_distributed_desired_task_count_handler(desired_thirty);
127+
.with_distributed_event_handler(no_desired_task_count)
128+
.with_distributed_event_handler(desired_thirty);
128129

129130
let plan = make_data_source_exec().await?;
130-
let response = DesiredTaskCountHandlers::handle(DesiredTaskCountEvent {
131+
let response = DesiredTaskCount::handle(DesiredTaskCountEvent {
131132
plan: &plan,
132133
session_config: &cfg,
133134
})

0 commit comments

Comments
 (0)