datafusion_distributed/distributed_planner/
distributed_config.rs

1use crate::config_extension_ext::set_distributed_option_extension;
2use datafusion::common::{DataFusionError, extensions_options, plan_err};
3use datafusion::config::{ConfigExtension, ConfigOptions};
4use datafusion::execution::TaskContext;
5use datafusion::prelude::SessionConfig;
6use std::sync::Arc;
7
8extensions_options! {
9    /// Configuration for the distributed planner.
10    pub struct DistributedConfig {
11        /// Sets the number of bytes each partitions is expected to scan from parquet files. If
12        /// more partitions than the ones available in one machine would be needed, several machines
13        /// are used, and the scan is distributed.
14        /// Lowering this number will increase parallelism.
15        pub file_scan_config_bytes_per_partition: usize, default = 16 * 1024 * 1024
16        /// Task multiplying factor for when a node declares that it changes the cardinality
17        /// of the data:
18        /// - If a node is increasing the cardinality of the data, this factor will increase.
19        /// - If a node reduces the cardinality of the data, this factor will decrease.
20        /// - In any other situation, this factor is left intact.
21        pub cardinality_task_count_factor: f64, default = cardinality_task_count_factor_default()
22        /// When encountering a UNION operation, isolate its children depending on the task context.
23        /// For example, on a UNION operation with 3 children running in 3 distributed tasks,
24        /// instead of executing the 3 children in each 3 tasks with a DistributedTaskContext of
25        /// 1/3, 2/3, and 3/3 respectively, Execute:
26        /// - The first child in the first task with a DistributedTaskContext of 1/1
27        /// - The second child in the second task with a DistributedTaskContext of 1/1
28        /// - The third child in the third task with a DistributedTaskContext of 1/1
29        pub children_isolator_unions: bool, default = true
30        /// Propagate collected metrics from all nodes in the plan across network boundaries
31        /// so that they can be reconstructed on the head node of the plan.
32        pub collect_metrics: bool, default = true
33        /// Enable broadcast joins for CollectLeft hash joins. When enabled, the build side of
34        /// a CollectLeft join is broadcast to all consumer tasks.
35        pub broadcast_joins: bool, default = true
36        /// The compression used for sending data over the network between workers.
37        /// It can be set to either `zstd`, `lz4` or `none`.
38        pub compression: String, default = "lz4".to_string()
39        /// Overrides `datafusion.execution.batch_size` for worker-executed stages. Because
40        /// `RepartitionExec` reads `session_config().batch_size()` at execute time to size its
41        /// output batches (via its internal `LimitedBatchCoalescer`), this knob lets users tune
42        /// shuffle batch sizes independently of the global `datafusion.execution.batch_size`.
43        ///
44        /// Set to 0 (the default) to apply no override and inherit `datafusion.execution.batch_size`.
45        pub shuffle_batch_size: usize, default = 0
46        /// Maximum tasks that will be assigned per stage during distributed planning.
47        /// If set to 0, this value is the number of workers returned by the provided `WorkerResolver`.
48        /// It defaults to 0.
49        pub max_tasks_per_stage: usize, default = 0
50        /// Enable the PartialReduce optimization, which inserts an extra aggregation pass
51        /// above hash RepartitionExec before network shuffles to reduce shuffle data size.
52        /// Disabled by default because its effectiveness is workload-dependent: it helps when
53        /// aggregation significantly reduces cardinality, but adds overhead when it does not.
54        pub partial_reduce: bool, default = false
55        /// Soft byte budget that each per-worker connection will buffer in memory before pausing
56        /// the gRPC pull from that worker. Per-partition channels are unbounded (to avoid
57        /// head-of-line blocking between sibling partitions), so backpressure is enforced
58        /// globally per [WorkerConnection] using this budget. A single message larger than this
59        /// budget will still be admitted (otherwise we would livelock), so the actual peak per
60        /// connection is `worker_connection_buffer_budget_bytes + max_message_size`.
61        pub worker_connection_buffer_budget_bytes: usize, default = 64 * 1024 * 1024
62        /// Calculates the task count of the different stages at execution time, based on runtime
63        /// information collected by sampling at the head of the stages.
64        ///
65        /// With this option enabled, the shape of the distributed plan is only known after fully
66        /// executing it, as it's dynamically created on the fly during execution.
67        pub dynamic_task_count: bool, default = false
68        /// If `dynamic_task_count` is enabled, this value is the amount of bytes each
69        /// partition is expected to handle. Lower values will result in greater parallelism.
70        pub dynamic_bytes_per_partition: usize, default = 16 * 1024 * 1024
71    }
72}
73
74fn cardinality_task_count_factor_default() -> f64 {
75    if cfg!(test) || cfg!(feature = "integration") {
76        1.5
77    } else {
78        1.0
79    }
80}
81
82impl DistributedConfig {
83    /// Gets the [DistributedConfig] from the [ConfigOptions]'s extensions.
84    pub fn from_config_options(cfg: &ConfigOptions) -> Result<&Self, DataFusionError> {
85        let Some(distributed_cfg) = cfg.extensions.get::<DistributedConfig>() else {
86            return plan_err!("DistributedConfig is not in ConfigOptions.extensions");
87        };
88        Ok(distributed_cfg)
89    }
90    /// Gets the [DistributedConfig] from the [ConfigOptions]'s extensions.
91    pub fn from_config_options_mut(cfg: &mut ConfigOptions) -> Result<&mut Self, DataFusionError> {
92        let Some(distributed_cfg) = cfg.extensions.get_mut::<DistributedConfig>() else {
93            return plan_err!("DistributedConfig is not in ConfigOptions.extensions");
94        };
95        Ok(distributed_cfg)
96    }
97
98    /// Gets the [DistributedConfig] from the [ConfigOptions]'s in the provided [SessionConfig].
99    pub fn from_session_config(session_cfg: &SessionConfig) -> Result<&Self, DataFusionError> {
100        Self::from_config_options(session_cfg.options())
101    }
102
103    /// Gets the [DistributedConfig] from the [ConfigOptions]'s in the provided [TaskContext].
104    pub fn from_task_context(ctx: &Arc<TaskContext>) -> Result<&Self, DataFusionError> {
105        Self::from_session_config(ctx.session_config())
106    }
107
108    /// Ensures that the [DistributedConfig] is present in the [SessionConfig]'s [ConfigOptions].
109    /// If not, it will insert a default [DistributedConfig] into the [SessionConfig]'s [ConfigOptions].
110    pub(crate) fn ensure_in_config(cfg: &mut SessionConfig) {
111        if cfg
112            .options()
113            .extensions
114            .get::<DistributedConfig>()
115            .is_none()
116        {
117            set_distributed_option_extension(cfg, DistributedConfig::default())
118        }
119    }
120}
121
122impl ConfigExtension for DistributedConfig {
123    const PREFIX: &'static str = "distributed";
124}
datafusion_distributed/distributed_planner/distributed_config.rs

datafusion_distributed/distributed_planner/
distributed_config.rs
