datafusion_distributed/
distributed_ext.rs

1use crate::codec::{set_distributed_user_codec, set_distributed_user_codec_arc};
2use crate::config_extension_ext::{
3    set_distributed_option_extension, set_distributed_option_extension_from_headers,
4};
5use crate::events::{
6    DesiredTaskCountHandler, DesiredTaskCountHandlers, RouteTasksHandler, RouteTasksHandlers,
7    ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers, WorkerPlanRewriteHandler,
8    WorkerPlanRewriteHandlers,
9};
10use crate::passthrough_headers::set_passthrough_headers;
11use crate::protocol::set_distributed_channel_resolver;
12use crate::work_unit_feed::set_distributed_work_unit_feed;
13use crate::worker_resolver::set_distributed_worker_resolver;
14use crate::{
15    ChannelResolver, DistributedConfig, LocalWorkerContext, WorkUnitFeed, WorkUnitFeedProvider,
16    WorkerResolver, get_distributed_worker_resolver,
17};
18use datafusion::common::DataFusionError;
19use datafusion::config::ConfigExtension;
20use datafusion::execution::{SessionState, SessionStateBuilder};
21use datafusion::physical_plan::ExecutionPlan;
22use datafusion::prelude::{SessionConfig, SessionContext};
23use datafusion_proto::physical_plan::PhysicalExtensionCodec;
24use delegate::delegate;
25use http::HeaderMap;
26use std::sync::Arc;
27
28/// Extends DataFusion with distributed capabilities.
29pub trait DistributedExt: Sized {
30    /// Adds the provided [ConfigExtension] to the distributed context. The [ConfigExtension] will
31    /// be serialized using gRPC metadata and sent across tasks. Users are expected to call this
32    /// method with their own extensions to be able to access them in any place in the
33    /// plan.
34    ///
35    /// This method also adds the provided [ConfigExtension] to the current session option
36    /// extensions, the same as calling [SessionConfig::with_option_extension].
37    ///
38    /// Example:
39    ///
40    /// ```rust
41    /// # use async_trait::async_trait;
42    /// # use datafusion::common::{extensions_options, DataFusionError};
43    /// # use datafusion::config::ConfigExtension;
44    /// # use datafusion::execution::{SessionState, SessionStateBuilder};
45    /// # use datafusion::prelude::SessionConfig;
46    /// # use datafusion_distributed::{DistributedExt, WorkerSessionBuilder, WorkerQueryContext};
47    ///
48    /// extensions_options! {
49    ///     pub struct CustomExtension {
50    ///         pub foo: String, default = "".to_string()
51    ///         pub bar: usize, default = 0
52    ///         pub baz: bool, default = false
53    ///     }
54    /// }
55    ///
56    /// impl ConfigExtension for CustomExtension {
57    ///     const PREFIX: &'static str = "custom";
58    /// }
59    ///
60    /// let mut my_custom_extension = CustomExtension::default();
61    /// // Now, the CustomExtension will be able to cross network boundaries. Upon making an Arrow
62    /// // Flight request, it will be sent through gRPC metadata.
63    /// let state = SessionStateBuilder::new()
64    ///     .with_distributed_option_extension(my_custom_extension)
65    ///     .build();
66    ///
67    /// async fn build_state(ctx: WorkerQueryContext) -> Result<SessionState, DataFusionError> {
68    ///     // This function can be provided to a Worker to tell it how to
69    ///     // build sessions that retrieve the CustomExtension from gRPC metadata.
70    ///     Ok(ctx
71    ///         .builder
72    ///         .with_distributed_option_extension_from_headers::<CustomExtension>(&ctx.headers)?
73    ///         .build())
74    /// }
75    /// ```
76    fn with_distributed_option_extension<T: ConfigExtension + Default>(self, t: T) -> Self;
77
78    /// Same as [DistributedExt::with_distributed_option_extension] but with an in-place mutation
79    fn set_distributed_option_extension<T: ConfigExtension + Default>(&mut self, t: T);
80
81    /// Adds the provided [ConfigExtension] to the distributed context. The [ConfigExtension] will
82    /// be serialized using gRPC metadata and sent across tasks. Users are expected to call this
83    /// method with their own extensions to be able to access them in any place in the
84    /// plan.
85    ///
86    /// - If there was a [ConfigExtension] of the same type already present, it's updated with an
87    ///   in-place mutation based on the headers that came over the wire.
88    /// - If there was no [ConfigExtension] set before, it will get added, as if
89    ///   [SessionConfig::with_option_extension] was being called.
90    ///
91    /// Example:
92    ///
93    /// ```rust
94    /// # use async_trait::async_trait;
95    /// # use datafusion::common::{extensions_options, DataFusionError};
96    /// # use datafusion::config::ConfigExtension;
97    /// # use datafusion::execution::{SessionState, SessionStateBuilder};
98    /// # use datafusion::prelude::SessionConfig;
99    /// # use datafusion_distributed::{DistributedExt, WorkerSessionBuilder, WorkerQueryContext};
100    ///
101    /// extensions_options! {
102    ///     pub struct CustomExtension {
103    ///         pub foo: String, default = "".to_string()
104    ///         pub bar: usize, default = 0
105    ///         pub baz: bool, default = false
106    ///     }
107    /// }
108    ///
109    /// impl ConfigExtension for CustomExtension {
110    ///     const PREFIX: &'static str = "custom";
111    /// }
112    ///
113    /// let mut my_custom_extension = CustomExtension::default();
114    /// // Now, the CustomExtension will be able to cross network boundaries. Upon making an Arrow
115    /// // Flight request, it will be sent through gRPC metadata.
116    /// let state = SessionStateBuilder::new()
117    ///     .with_distributed_option_extension(my_custom_extension)
118    ///     .build();
119    ///
120    /// async fn build_state(ctx: WorkerQueryContext) -> Result<SessionState, DataFusionError> {
121    ///     // This function can be provided to a Worker to tell it how to
122    ///     // build sessions that retrieve the CustomExtension from gRPC metadata.
123    ///     Ok(ctx
124    ///         .builder
125    ///         .with_distributed_option_extension_from_headers::<CustomExtension>(&ctx.headers)?
126    ///         .build())
127    /// }
128    /// ```
129    fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(
130        self,
131        headers: &HeaderMap,
132    ) -> Result<Self, DataFusionError>;
133
134    /// Same as [DistributedExt::with_distributed_option_extension_from_headers] but with an in-place mutation
135    fn set_distributed_option_extension_from_headers<T: ConfigExtension + Default>(
136        &mut self,
137        headers: &HeaderMap,
138    ) -> Result<(), DataFusionError>;
139
140    /// Injects a user-defined [PhysicalExtensionCodec] that is capable of encoding/decoding
141    /// custom execution nodes. Multiple user-defined [PhysicalExtensionCodec] can be added
142    /// by calling this method several times.
143    ///
144    /// Example:
145    ///
146    /// ```
147    /// # use std::sync::Arc;
148    /// # use datafusion::common::DataFusionError;
149    /// # use datafusion::execution::{SessionState, FunctionRegistry, SessionStateBuilder, TaskContext};
150    /// # use datafusion::physical_plan::ExecutionPlan;
151    /// # use datafusion::prelude::SessionConfig;
152    /// # use datafusion_proto::physical_plan::{PhysicalExtensionCodec, PhysicalProtoConverterExtension};
153    /// # use datafusion_distributed::{DistributedExt, WorkerQueryContext};
154    ///
155    /// #[derive(Debug)]
156    /// struct CustomExecCodec;
157    ///
158    /// impl PhysicalExtensionCodec for CustomExecCodec {
159    ///     fn try_decode(&self, buf: &[u8], inputs: &[Arc<dyn ExecutionPlan>], ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
160    ///         todo!()
161    ///     }
162    ///
163    ///     fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>, proto_converter: &dyn PhysicalProtoConverterExtension) -> datafusion::common::Result<()> {
164    ///         todo!()
165    ///     }
166    /// }
167    ///
168    /// let state = SessionStateBuilder::new()
169    ///     .with_distributed_user_codec(CustomExecCodec)
170    ///     .build();
171    ///
172    /// async fn build_state(ctx: WorkerQueryContext) -> Result<SessionState, DataFusionError> {
173    ///     // This function can be provided to a Worker to tell it how to
174    ///     // encode/decode CustomExec nodes.
175    ///     Ok(SessionStateBuilder::new()
176    ///         .with_distributed_user_codec(CustomExecCodec)
177    ///         .build())
178    /// }
179    /// ```
180    fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(self, codec: T) -> Self;
181
182    /// Same as [DistributedExt::with_distributed_user_codec] but with an in-place mutation
183    fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
184
185    /// Same as [DistributedExt::with_distributed_user_codec] but with a dynamic argument.
186    fn with_distributed_user_codec_arc(self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
187
188    /// Same as [DistributedExt::set_distributed_user_codec] but with a dynamic argument.
189    fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>);
190
191    /// This is what tells Distributed DataFusion the URLs of the workers available for serving queries.
192    ///
193    /// It injects a [WorkerResolver] implementation for Distributed DataFusion to resolve worker
194    /// nodes in the cluster. When running in distributed mode, setting a [WorkerResolver] is required.
195    ///
196    /// Even if this is required to be present in the [SessionContext] that first initiates and
197    /// plans the query, it's not necessary to be present in a Worker's session state builder,
198    /// as no planning happens there.
199    ///
200    /// Example:
201    ///
202    /// ```
203    /// # use async_trait::async_trait;
204    /// # use datafusion::common::DataFusionError;
205    /// # use datafusion::execution::{SessionState, SessionStateBuilder};
206    /// # use datafusion::prelude::SessionConfig;
207    /// # use url::Url;
208    /// # use std::sync::Arc;
209    /// # use datafusion_distributed::{WorkerResolver, DistributedExt, SessionStateBuilderExt, WorkerQueryContext};
210    ///
211    /// struct CustomWorkerResolver;
212    ///
213    /// #[async_trait]
214    /// impl WorkerResolver for CustomWorkerResolver {
215    ///     fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
216    ///         todo!()
217    ///     }
218    /// }
219    ///
220    /// // This tweaks the SessionState so that it can plan for distributed queries and execute them.
221    /// let state = SessionStateBuilder::new()
222    ///     .with_distributed_worker_resolver(CustomWorkerResolver)
223    ///     .with_distributed_planner()
224    ///     .build();
225    /// ```
226    fn with_distributed_worker_resolver<T: WorkerResolver + 'static>(self, resolver: T) -> Self;
227
228    /// Same as [DistributedExt::with_distributed_channel_resolver] but with an in-place mutation.
229    fn set_distributed_worker_resolver<T: WorkerResolver + 'static>(&mut self, resolver: T);
230
231    /// This is what tells Distributed DataFusion how to build a Worker gRPC client out of a worker URL.
232    ///
233    /// There's a default implementation that caches the Worker client instances so that there's
234    /// only one per URL, but users can decide to override that behavior in favor of their own solution.
235    ///
236    /// Example:
237    ///
238    /// ```
239    /// # use async_trait::async_trait;
240    /// # use datafusion::common::DataFusionError;
241    /// # use datafusion::execution::{SessionState, SessionStateBuilder};
242    /// # use datafusion::prelude::SessionConfig;
243    /// # use url::Url;
244    /// # use std::sync::Arc;
245    /// # use datafusion_distributed::{ChannelResolver, DistributedExt, SessionStateBuilderExt, WorkerChannel, WorkerQueryContext, grpc};
246    ///
247    /// struct CustomChannelResolver;
248    ///
249    /// #[async_trait]
250    /// impl ChannelResolver for CustomChannelResolver {
251    ///     async fn get_worker_client_for_url(&self, url: &Url) -> Result<Box<dyn WorkerChannel>, DataFusionError> {
252    ///         // Build a custom worker client wrapped with tower layers or something similar.
253    ///         todo!()
254    ///     }
255    /// }
256    ///
257    /// // This tweaks the SessionState so that it can plan for distributed queries and execute them.
258    /// let state = SessionStateBuilder::new()
259    ///     .with_distributed_channel_resolver(CustomChannelResolver)
260    ///     .with_distributed_planner()
261    ///     .build();
262    ///
263    /// // This function can be provided to a Worker so that, upon receiving a distributed
264    /// // part of a plan, it knows how to resolve gRPC channels from URLs for making network calls to other nodes.
265    /// async fn build_state(ctx: WorkerQueryContext) -> Result<SessionState, DataFusionError> {
266    ///     // If you have a custom channel resolver, it should also be passed in the
267    ///     // Worker session builder.
268    ///     Ok(ctx
269    ///         .builder
270    ///         .with_distributed_channel_resolver(CustomChannelResolver)
271    ///         .build())
272    /// }
273    /// ```
274    fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
275        self,
276        resolver: T,
277    ) -> Self;
278
279    /// Same as [DistributedExt::with_distributed_channel_resolver] but with an in-place mutation.
280    fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
281        &mut self,
282        resolver: T,
283    );
284
285    /// Sets the number of bytes each partition in a stage with a FileScanConfig node is
286    /// expected to scan. A task runs `target_partitions` partitions, so the task count is
287    /// roughly `total_scan_bytes / bytes_per_partition / target_partitions` (capped at the
288    /// number of available workers). Reducing this number increases the amount of tasks.
289    ///
290    /// ```text
291    ///     ┌───────────────────────┐
292    ///     │SortPreservingMergeExec│
293    ///     └───────────────────────┘
294    ///                 ▲
295    /// ┌ ─ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ─ ─ Stage 2
296    ///     ┌───────────┴───────────┐    │
297    /// │   │       SortExec        │
298    ///     └───────────────────────┘    │
299    /// │   ┌───────────────────────┐
300    ///     │     AggregateExec     │    │
301    /// │   └───────────────────────┘
302    ///  ─ ─ ─ ─ ─ ─ ─ ─▲─ ─ ─ ─ ─ ─ ─ ─ ┘
303    /// ┌ ─ ─ ─ ─ ─ ─ ─ ┴ ─ ─ ─ ─ ─ ─ ─ ─ Stage 1
304    ///     ┌───────────────────────┐    │
305    /// │   │      FilterExec       │
306    ///     └───────────────────────┘    │
307    /// │   ┌───────────────────────┐        Sets the bytes scanned per
308    ///     │    FileScanConfig     │◀───┼─   partition. Less
309    /// │   └───────────────────────┘        bytes_per_partition == more tasks
310    ///  ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘
311    ///```
312    fn with_distributed_file_scan_config_bytes_per_partition(
313        self,
314        bytes_per_partition: usize,
315    ) -> Result<Self, DataFusionError>;
316
317    /// Same as [DistributedExt::with_distributed_file_scan_config_bytes_per_partition] but with an in-place mutation.
318    fn set_distributed_file_scan_config_bytes_per_partition(
319        &mut self,
320        bytes_per_partition: usize,
321    ) -> Result<(), DataFusionError>;
322
323    /// The number of tasks in each stage is calculated in a bottom-to-top fashion.
324    ///
325    /// Bottom stages containing leaf nodes will provide an estimation of the amount of tasks
326    /// for those stages, but upper stages might see a reduction (or increment) in the amount
327    /// of tasks based on the cardinality effect bottom stages have in the data.
328    ///
329    /// For example: If there are two stages, and the leaf stage is estimated to use 10 tasks,
330    ///  the upper stage might use less (e.g. 5) if it sees that the leaf stage is returning
331    ///  less data because of filters or aggregations.
332    ///
333    /// This function sets the scale factor for when encountering these nodes that change the
334    /// cardinality of the data. For example, if a stage with 10 tasks contains an AggregateExec
335    /// node, and the scale factor is 2.0, the following stage will use  10 / 2.0 = 5 tasks.
336    ///
337    /// ```text
338    ///     ┌───────────────────────┐
339    ///     │SortPreservingMergeExec│
340    ///     └───────────────────────┘
341    ///                 ▲
342    /// ┌ ─ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ─ ─ Stage 2 (N/scale_factor tasks)
343    ///     ┌───────────┴───────────┐    │
344    /// │   │       SortExec        │
345    ///     └───────────────────────┘    │
346    /// │   ┌───────────────────────┐
347    ///     │     AggregateExec     │    │
348    /// │   └───────────────────────┘
349    ///  ─ ─ ─ ─ ─ ─ ─ ─▲─ ─ ─ ─ ─ ─ ─ ─ ┘
350    /// ┌ ─ ─ ─ ─ ─ ─ ─ ┴ ─ ─ ─ ─ ─ ─ ─ ─ Stage 1 (N tasks)
351    ///     ┌───────────────────────┐    │       A filter reduces cardinality,
352    /// │   │      FilterExec       │◀────────therefore the next stage will have
353    ///     └───────────────────────┘    │    less tasks according to this factor
354    /// │   ┌───────────────────────┐
355    ///     │    FileScanConfig     │    │
356    /// │   └───────────────────────┘
357    ///  ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘
358    /// ```
359    fn with_distributed_cardinality_effect_task_scale_factor(
360        self,
361        factor: f64,
362    ) -> Result<Self, DataFusionError>;
363
364    /// Same as [DistributedExt::with_distributed_cardinality_effect_task_scale_factor] but with
365    /// an in-place mutation.
366    fn set_distributed_cardinality_effect_task_scale_factor(
367        &mut self,
368        factor: f64,
369    ) -> Result<(), DataFusionError>;
370
371    /// Enables metrics collection across network boundaries so that all the metrics gather in
372    /// each node are accessible from the head stage that started running the query.
373    fn with_distributed_metrics_collection(self, enabled: bool) -> Result<Self, DataFusionError>;
374
375    /// Same as [DistributedExt::with_distributed_metrics_collection] but with an in-place mutation.
376    fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>;
377
378    /// Enables children isolator unions for distributing UNION operations across as many tasks as
379    /// the sum of all the tasks required for each child.
380    ///
381    /// For example, if there is a UNION with 3 children, requiring one task each, it will result
382    /// in a plan with 3 tasks where each task runs one child:
383    ///
384    /// ```text
385    /// ┌─────────────────────────────┐┌─────────────────────────────┐┌─────────────────────────────┐
386    /// │           Task 1            ││           Task 2            ││           Task 3            │
387    /// │┌───────────────────────────┐││┌───────────────────────────┐││┌───────────────────────────┐│
388    /// ││ ChildrenIsolatorUnionExec ││││ ChildrenIsolatorUnionExec ││││ ChildrenIsolatorUnionExec ││
389    /// │└───▲─────────▲─────────▲───┘││└───▲─────────▲─────────▲───┘││└───▲─────────▲─────────▲───┘│
390    /// │    │                        ││              │              ││                        │    │
391    /// │┌───┴───┐ ┌  ─│ ─   ┌  ─│ ─  ││┌  ─│ ─   ┌───┴───┐ ┌  ─│ ─  ││┌  ─│ ─   ┌  ─│ ─   ┌───┴───┐│
392    /// ││Child 1│  Child 2│  Child 3│││ Child 1│ │Child 2│  Child 3│││ Child 1│  Child 2│ │Child 3││
393    /// │└───────┘ └  ─  ─   └  ─  ─  ││└  ─  ─   └───────┘ └  ─  ─  ││└  ─  ─   └  ─  ─   └───────┘│
394    /// └─────────────────────────────┘└─────────────────────────────┘└─────────────────────────────┘
395    /// ```
396    fn with_distributed_children_isolator_unions(
397        self,
398        enabled: bool,
399    ) -> Result<Self, DataFusionError>;
400
401    /// Same as [DistributedExt::with_distributed_children_isolator_unions] but with an in-place mutation.
402    fn set_distributed_children_isolator_unions(
403        &mut self,
404        enabled: bool,
405    ) -> Result<(), DataFusionError>;
406
407    /// Enables broadcast joins for CollectLeft hash joins. When enabled, the build side of
408    /// a CollectLeft join is broadcast to all consumer tasks instead of being coalesced
409    /// into a single partition.
410    ///
411    /// Note: This option is disabled by default until the implementation is smarter about when to
412    /// broadcast.
413    fn with_distributed_broadcast_joins(self, enabled: bool) -> Result<Self, DataFusionError>;
414
415    /// Same as [DistributedExt::with_distributed_broadcast_joins_enabled] but with an in-place mutation.
416    fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError>;
417
418    #[cfg(feature = "grpc")]
419    /// The compression type to use for sending data over the wire.
420    ///
421    /// The default is [CompressionType::LZ4_FRAME].
422    fn with_distributed_compression(
423        self,
424        compression: Option<arrow_ipc::CompressionType>,
425    ) -> Result<Self, DataFusionError>;
426
427    #[cfg(feature = "grpc")]
428    /// Same as [DistributedExt::with_distributed_compression] but with an in-place mutation.
429    fn set_distributed_compression(
430        &mut self,
431        compression: Option<arrow_ipc::CompressionType>,
432    ) -> Result<(), DataFusionError>;
433
434    /// Overrides `datafusion.execution.batch_size` for worker-executed stages, letting users
435    /// tune shuffle batch sizes (specifically `RepartitionExec`'s output batching via its
436    /// internal `LimitedBatchCoalescer`) independently of the global batch size.
437    ///
438    /// Set to 0 (the default) to apply no override.
439    fn with_distributed_shuffle_batch_size(
440        self,
441        batch_size: usize,
442    ) -> Result<Self, DataFusionError>;
443
444    /// Same as [DistributedExt::with_distributed_shuffle_batch_size] but with an in-place mutation.
445    fn set_distributed_shuffle_batch_size(
446        &mut self,
447        batch_size: usize,
448    ) -> Result<(), DataFusionError>;
449
450    /// Sets arbitrary HTTP headers that will be forwarded unchanged to worker nodes.
451    /// These headers are included in outgoing Arrow Flight requests to workers.
452    ///
453    /// Returns an error if any header name starts with the reserved prefix
454    /// `x-datafusion-distributed-config-`, which is used internally.
455    ///
456    /// Example:
457    ///
458    /// ```rust
459    /// # use datafusion::execution::SessionStateBuilder;
460    /// # use datafusion_distributed::DistributedExt;
461    /// # use http::HeaderMap;
462    ///
463    /// let mut passthrough = HeaderMap::new();
464    /// passthrough.insert("x-custom-priority", "high".parse().unwrap());
465    ///
466    /// let state = SessionStateBuilder::new()
467    ///     .with_distributed_passthrough_headers(passthrough)
468    ///     .unwrap()
469    ///     .build();
470    /// ```
471    fn with_distributed_passthrough_headers(
472        self,
473        headers: HeaderMap,
474    ) -> Result<Self, DataFusionError>;
475
476    /// Same as [DistributedExt::with_distributed_passthrough_headers] but with an in-place mutation.
477    fn set_distributed_passthrough_headers(
478        &mut self,
479        headers: HeaderMap,
480    ) -> Result<(), DataFusionError>;
481
482    /// Sets the maximum tasks that will be assigned for each stage.
483    ///
484    /// If not specified, the number of workers returned by the provided [WorkerResolver] is taken.
485    fn with_distributed_max_tasks_per_stage(
486        self,
487        max_tasks_per_stage: usize,
488    ) -> Result<Self, DataFusionError>;
489
490    /// Same as [DistributedExt::with_distributed_max_tasks_per_stage] but with an in-place mutation.
491    fn set_distributed_max_tasks_per_stage(
492        &mut self,
493        max_tasks_per_stage: usize,
494    ) -> Result<(), DataFusionError>;
495
496    /// Enables or disables the PartialReduce optimization, which inserts an extra aggregation
497    /// pass above hash RepartitionExec before network shuffles to reduce shuffle data size.
498    /// Disabled by default because its effectiveness is workload-dependent: it helps when
499    /// aggregation significantly reduces cardinality, but adds overhead when it does not.
500    fn with_distributed_partial_reduce(self, enabled: bool) -> Result<Self, DataFusionError>;
501
502    /// Same as [DistributedExt::with_distributed_partial_reduce] but with an in-place mutation.
503    fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError>;
504
505    /// Sets the soft byte budget that each per-worker connection will buffer in memory before
506    /// pausing the gRPC pull from that worker. Per-partition channels are unbounded (to avoid
507    /// head-of-line blocking between sibling partitions), so backpressure is enforced globally
508    /// per worker connection using this budget.
509    fn with_distributed_worker_connection_buffer_budget_bytes(
510        self,
511        budget_bytes: usize,
512    ) -> Result<Self, DataFusionError>;
513
514    /// Same as [DistributedExt::with_distributed_worker_connection_buffer_budget_bytes] but with
515    /// an in-place mutation.
516    fn set_distributed_worker_connection_buffer_budget_bytes(
517        &mut self,
518        budget_bytes: usize,
519    ) -> Result<(), DataFusionError>;
520
521    /// Registers a [WorkUnitFeed] so that Distributed DataFusion can discover it while traversing
522    /// plans. For more info, refer to [WorkUnitFeed] docs.
523    ///
524    /// This method uses some type system trickery so that users can provide a callback like this:
525    ///
526    /// ```ignore
527    /// # use datafusion::execution::SessionStateBuilder;
528    ///
529    /// SessionStateBuilder::new()
530    ///     .with_distributed_work_unit_feed(|p: &MyCustomPlan| &p.my_work_unit_feed);
531    /// ```
532    fn with_distributed_work_unit_feed<T, P, F>(self, getter: F) -> Self
533    where
534        T: ExecutionPlan + 'static,
535        P: WorkUnitFeedProvider + 'static,
536        P::WorkUnit: 'static,
537        F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static;
538
539    /// Same as [DistributedExt::with_distributed_work_unit_feed] but with an in-place mutation.
540    fn set_distributed_work_unit_feed<T, P, F>(&mut self, getter: F)
541    where
542        T: ExecutionPlan + 'static,
543        P: WorkUnitFeedProvider + 'static,
544        P::WorkUnit: 'static,
545        F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static;
546
547    /// Dynamically allocates tasks to the different stages based on runtime statistics
548    /// collected during execution.
549    fn with_distributed_dynamic_task_count(self, enabled: bool) -> Result<Self, DataFusionError>;
550
551    /// Same as [DistributedExt::with_distributed_dynamic_task_count] but with an in-place mutation.
552    fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError>;
553
554    /// Target throughput in bytes per partition per second used by the dynamic task count
555    /// allocator to decide how many tasks to assign to each stage based on runtime statistics.
556    fn with_distributed_dynamic_bytes_per_partition(
557        self,
558        dynamic_bytes_per_partition: usize,
559    ) -> Result<Self, DataFusionError>;
560
561    /// Same as [DistributedExt::with_distributed_dynamic_bytes_per_partition] but with an
562    /// in-place mutation.
563    fn set_distributed_dynamic_bytes_per_partition(
564        &mut self,
565        dynamic_bytes_per_partition: usize,
566    ) -> Result<(), DataFusionError>;
567
568    /// Target throughput in bytes per partition per second used by the dynamic task count
569    /// allocator to decide how many tasks to assign to each stage based on runtime statistics.
570    fn with_distributed_local_worker_context(
571        self,
572        local_worker_context: LocalWorkerContext,
573    ) -> Self;
574
575    /// Same as [DistributedExt::with_distributed_local_worker_context] but with an
576    /// in-place mutation.
577    fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext);
578
579    /// Registers a handler that supplies desired or maximum task-count hints for plan nodes.
580    /// The distributed planner reconciles these hints when choosing each stage's task count.
581    ///
582    /// A function with the following signature can be provided as argument:
583    ///
584    /// ```rust
585    /// # use datafusion::execution::SessionStateBuilder;
586    /// # use datafusion::physical_plan::empty::EmptyExec;
587    /// # use datafusion::common::Result;
588    /// # use datafusion_distributed::{DistributedExt, DesiredTaskCountEvent, DesiredTaskCountEventResponse};
589    ///
590    /// fn handle_custom_desired_task_count(event: DesiredTaskCountEvent) -> Option<Result<DesiredTaskCountEventResponse>> {
591    ///     let _exec = event.plan.downcast_ref::<EmptyExec>()?;
592    ///     Some(Ok(DesiredTaskCountEventResponse::desired(3)))
593    /// }
594    ///
595    /// SessionStateBuilder::new()
596    ///     .with_distributed_desired_task_count_handler(handle_custom_desired_task_count);
597    /// ```
598    ///
599    /// ```text
600    /// ┌────────────────┐
601    /// │CustomDataSource│──────────▶ 3 desired tasks
602    /// └────────────────┘
603    /// ```
604    fn with_distributed_desired_task_count_handler<T: DesiredTaskCountHandler>(
605        self,
606        handler: T,
607    ) -> Self;
608
609    /// Same as [DistributedExt::with_distributed_desired_task_count_handler] but with an
610    /// in-place mutation.
611    fn set_distributed_desired_task_count_handler<T: DesiredTaskCountHandler>(
612        &mut self,
613        handler: T,
614    );
615
616    /// Registers a handler that can replace leaf nodes with distributed variants once the
617    /// distributed planner has decided a final task count.
618    ///
619    /// A function with the following signature can be provided as argument:
620    ///
621    /// ```rust
622    /// # use std::sync::Arc;
623    /// # use datafusion::error::Result;
624    /// # use datafusion::execution::SessionStateBuilder;
625    /// # use datafusion::physical_plan::ExecutionPlan;
626    /// # use datafusion::physical_plan::empty::EmptyExec;
627    /// # use datafusion_distributed::{DistributedExt, DistributedLeafExec, ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse};
628    ///
629    /// fn handle_custom_scale_up_leaf_node(event: ScaleUpLeafNodeEvent) -> Option<Result<ScaleUpLeafNodeEventResponse>> {
630    ///     let exec = event.plan.downcast_ref::<EmptyExec>()?;
631    ///     // Build per-task plan variants, each handling its own non-overlapping subset of data.
632    ///     let variants = (0..event.task_count)
633    ///         .map(|_| Arc::new(EmptyExec::new(exec.schema())) as _);
634    ///     Some(
635    ///         DistributedLeafExec::try_new(Arc::clone(event.plan), variants)
636    ///             .map(|exec| ScaleUpLeafNodeEventResponse::new(Arc::new(exec))),
637    ///     )
638    /// }
639    ///
640    /// SessionStateBuilder::new()
641    ///     .with_distributed_scale_up_leaf_node_handler(handle_custom_scale_up_leaf_node);
642    /// ```
643    ///
644    /// ```text
645    ///                             ┌────────────────────────────────────────────────────────────┐
646    ///                             │                    DistributedLeafExec                     │
647    /// ┌────────────────┐          │┌──────────────────┐┌──────────────────┐┌──────────────────┐│
648    /// │CustomDataSource│─3 tasks─▶││ CustomDataSource ││ CustomDataSource ││ CustomDataSource ││
649    /// └────────────────┘          ││      (1/3)       ││      (2/3)       ││      (3/3)       ││
650    ///                             │└──────────────────┘└──────────────────┘└──────────────────┘│
651    ///                             └────────────────────────────────────────────────────────────┘
652    /// ```
653    fn with_distributed_scale_up_leaf_node_handler<T: ScaleUpLeafNodeHandler>(
654        self,
655        handler: T,
656    ) -> Self;
657
658    /// Same as [DistributedExt::with_distributed_scale_up_leaf_node_handler] but with an
659    /// in-place mutation.
660    fn set_distributed_scale_up_leaf_node_handler<T: ScaleUpLeafNodeHandler>(&mut self, handler: T);
661
662    /// Registers a handler that maps a stage's task slots to worker URLs before execution.
663    /// A response must contain one URL per task, in task-index order.
664    ///
665    /// A function with the following signature can be provided as argument:
666    ///
667    /// ```rust
668    /// # use datafusion::error::Result;
669    /// # use datafusion::execution::SessionStateBuilder;
670    /// # use datafusion_distributed::{DistributedExt, DistributedGetterExt, RouteTasksEvent, RouteTasksEventResponse};
671    ///
672    /// fn handle_custom_route_tasks(event: RouteTasksEvent) -> Option<Result<RouteTasksEventResponse>> {
673    ///     let routing = event.task_ctx.session_config()
674    ///         .get_distributed_worker_resolver()
675    ///         .and_then(|resolver| resolver.get_urls())
676    ///         .map(|workers| RouteTasksEventResponse::new(
677    ///             workers.into_iter().cycle().take(event.task_count).collect()
678    ///         ));
679    ///     Some(routing)
680    /// }
681    ///
682    /// SessionStateBuilder::new()
683    ///     .with_distributed_route_tasks_handler(handle_custom_route_tasks);
684    /// ```
685    ///
686    /// ```text
687    /// task 0  ──► http://worker1
688    /// task 1  ──► http://worker2     RouteTasksHandler
689    /// task 2  ──► http://worker3
690    /// ```
691    fn with_distributed_route_tasks_handler<T: RouteTasksHandler>(self, handler: T) -> Self;
692
693    /// Same as [DistributedExt::with_distributed_route_tasks_handler] but with an in-place
694    /// mutation.
695    fn set_distributed_route_tasks_handler<T: RouteTasksHandler>(&mut self, handler: T);
696
697    /// Registers a handler that rewrites a decoded worker stage plan before it is executed.
698    ///
699    /// Handlers registered with this method need to meet the following criteria:
700    ///
701    /// - The topology of the plan must remain the same: no nodes or edges may be added or removed.
702    /// - Each output partition must produce the same rows, including their multiplicity.
703    /// - The output schema, partitioning, ordering, boundedness, and emission type of the head
704    ///   node must remain unchanged. Intermediate nodes may change their output schema.
705    ///
706    /// Handlers run in registration order, each receiving the plan returned by the preceding
707    /// handler. Returning an error aborts worker plan registration.
708    ///
709    /// Register this handler on the [`SessionStateBuilder`] used by each
710    /// [`crate::WorkerSessionBuilder`]. Registering it on the coordinating session has no effect:
711    /// the coordinator session context is not used by workers
712    ///
713    /// ```rust
714    /// # use datafusion::common::Result;
715    /// # use datafusion::execution::SessionState;
716    /// # use datafusion_distributed::{DistributedExt, Worker, WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse, WorkerQueryContext};
717    ///
718    /// async fn build_worker_session(ctx: WorkerQueryContext) -> Result<SessionState> {
719    ///     Ok(ctx
720    ///         .builder
721    ///         .with_distributed_worker_plan_rewrite_handler(
722    ///             |event: WorkerPlanRewriteEvent<'_>| {
723    ///                 Ok(WorkerPlanRewriteEventResponse::new(event.plan))
724    ///             },
725    ///         )
726    ///         .build())
727    /// }
728    ///
729    /// let _worker = Worker::from_session_builder(build_worker_session);
730    /// ```
731    fn with_distributed_worker_plan_rewrite_handler<T: WorkerPlanRewriteHandler>(
732        self,
733        handler: T,
734    ) -> Self;
735
736    /// Same as [DistributedExt::with_distributed_worker_plan_rewrite_handler] but with an
737    /// in-place mutation.
738    fn set_distributed_worker_plan_rewrite_handler<T: WorkerPlanRewriteHandler>(
739        &mut self,
740        handler: T,
741    );
742}
743
744/// Trait to have a unified interface for getting structs & properties from SessionConfig that are used in distributed context.
745pub trait DistributedGetterExt: Sized {
746    /// Gets the [WorkerResolver] from this session's config.
747    fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError>;
748}
749
750impl DistributedExt for SessionConfig {
751    fn set_distributed_option_extension<T: ConfigExtension + Default>(&mut self, t: T) {
752        set_distributed_option_extension(self, t)
753    }
754
755    fn set_distributed_option_extension_from_headers<T: ConfigExtension + Default>(
756        &mut self,
757        headers: &HeaderMap,
758    ) -> Result<(), DataFusionError> {
759        set_distributed_option_extension_from_headers::<T>(self, headers)?;
760        Ok(())
761    }
762
763    fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T) {
764        set_distributed_user_codec(self, codec)
765    }
766
767    fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>) {
768        set_distributed_user_codec_arc(self, codec)
769    }
770
771    fn set_distributed_worker_resolver<T: WorkerResolver + 'static>(&mut self, resolver: T) {
772        set_distributed_worker_resolver(self, resolver);
773    }
774
775    fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
776        &mut self,
777        resolver: T,
778    ) {
779        set_distributed_channel_resolver(self, resolver);
780    }
781
782    fn set_distributed_file_scan_config_bytes_per_partition(
783        &mut self,
784        bytes_per_partition: usize,
785    ) -> Result<(), DataFusionError> {
786        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
787        d_cfg.file_scan_config_bytes_per_partition = bytes_per_partition;
788        Ok(())
789    }
790
791    fn set_distributed_cardinality_effect_task_scale_factor(
792        &mut self,
793        factor: f64,
794    ) -> Result<(), DataFusionError> {
795        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
796        d_cfg.cardinality_task_count_factor = factor;
797        Ok(())
798    }
799
800    fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError> {
801        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
802        d_cfg.collect_metrics = enabled;
803        Ok(())
804    }
805
806    fn set_distributed_children_isolator_unions(
807        &mut self,
808        enabled: bool,
809    ) -> Result<(), DataFusionError> {
810        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
811        d_cfg.children_isolator_unions = enabled;
812        Ok(())
813    }
814
815    fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError> {
816        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
817        d_cfg.broadcast_joins = enabled;
818        Ok(())
819    }
820
821    #[cfg(feature = "grpc")]
822    fn set_distributed_compression(
823        &mut self,
824        compression: Option<arrow_ipc::CompressionType>,
825    ) -> Result<(), DataFusionError> {
826        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
827        d_cfg.compression = match compression {
828            Some(arrow_ipc::CompressionType::ZSTD) => "zstd".to_string(),
829            Some(arrow_ipc::CompressionType::LZ4_FRAME) => "lz4".to_string(),
830            _ => "none".to_string(),
831        };
832        Ok(())
833    }
834
835    fn set_distributed_shuffle_batch_size(
836        &mut self,
837        batch_size: usize,
838    ) -> Result<(), DataFusionError> {
839        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
840        d_cfg.shuffle_batch_size = batch_size;
841        Ok(())
842    }
843
844    fn set_distributed_passthrough_headers(
845        &mut self,
846        headers: HeaderMap,
847    ) -> Result<(), DataFusionError> {
848        set_passthrough_headers(self, headers)
849    }
850
851    fn set_distributed_max_tasks_per_stage(
852        &mut self,
853        max_tasks_per_stage: usize,
854    ) -> Result<(), DataFusionError> {
855        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
856        d_cfg.max_tasks_per_stage = max_tasks_per_stage;
857        Ok(())
858    }
859
860    fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError> {
861        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
862        d_cfg.partial_reduce = enabled;
863        Ok(())
864    }
865
866    fn set_distributed_worker_connection_buffer_budget_bytes(
867        &mut self,
868        budget_bytes: usize,
869    ) -> Result<(), DataFusionError> {
870        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
871        d_cfg.worker_connection_buffer_budget_bytes = budget_bytes;
872        Ok(())
873    }
874
875    fn set_distributed_work_unit_feed<T, P, F>(&mut self, getter: F)
876    where
877        T: ExecutionPlan + 'static,
878        P: WorkUnitFeedProvider + 'static,
879        P::WorkUnit: 'static,
880        F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static,
881    {
882        set_distributed_work_unit_feed(self, move |plan: &Arc<dyn ExecutionPlan>| {
883            plan.downcast_ref::<T>().and_then(&getter)
884        })
885    }
886
887    fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError> {
888        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
889        d_cfg.dynamic_task_count = enabled;
890        Ok(())
891    }
892
893    fn set_distributed_dynamic_bytes_per_partition(
894        &mut self,
895        dynamic_bytes_per_partition: usize,
896    ) -> Result<(), DataFusionError> {
897        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
898        d_cfg.dynamic_bytes_per_partition = dynamic_bytes_per_partition;
899        Ok(())
900    }
901
902    fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext) {
903        self.set_extension(Arc::new(local_worker_context));
904    }
905
906    fn set_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(&mut self, h: H) {
907        DesiredTaskCountHandlers::push_custom(self, Arc::new(h))
908    }
909
910    fn set_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(&mut self, h: H) {
911        ScaleUpLeafNodeHandlers::push_custom(self, Arc::new(h));
912    }
913
914    fn set_distributed_route_tasks_handler<H: RouteTasksHandler>(&mut self, h: H) {
915        RouteTasksHandlers::push_custom(self, Arc::new(h));
916    }
917
918    fn set_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(&mut self, h: H) {
919        WorkerPlanRewriteHandlers::push_custom(self, Arc::new(h));
920    }
921
922    delegate! {
923        to self {
924            #[call(set_distributed_option_extension)]
925            #[expr($;self)]
926            fn with_distributed_option_extension<T: ConfigExtension + Default>(mut self, t: T) -> Self;
927
928            #[call(set_distributed_option_extension_from_headers)]
929            #[expr($?;Ok(self))]
930            fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(mut self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
931
932            #[call(set_distributed_user_codec)]
933            #[expr($;self)]
934            fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
935
936            #[call(set_distributed_user_codec_arc)]
937            #[expr($;self)]
938            fn with_distributed_user_codec_arc(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
939
940            #[call(set_distributed_worker_resolver)]
941            #[expr($;self)]
942            fn with_distributed_worker_resolver<T: WorkerResolver + 'static>(mut self, resolver: T) -> Self;
943
944            #[call(set_distributed_channel_resolver)]
945            #[expr($;self)]
946            fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
947
948            #[call(set_distributed_file_scan_config_bytes_per_partition)]
949            #[expr($?;Ok(self))]
950            fn with_distributed_file_scan_config_bytes_per_partition(mut self, bytes_per_partition: usize) -> Result<Self, DataFusionError>;
951
952            #[call(set_distributed_cardinality_effect_task_scale_factor)]
953            #[expr($?;Ok(self))]
954            fn with_distributed_cardinality_effect_task_scale_factor(mut self, factor: f64) -> Result<Self, DataFusionError>;
955
956            #[call(set_distributed_metrics_collection)]
957            #[expr($?;Ok(self))]
958            fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result<Self, DataFusionError>;
959
960            #[call(set_distributed_children_isolator_unions)]
961            #[expr($?;Ok(self))]
962            fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result<Self, DataFusionError>;
963
964            #[call(set_distributed_broadcast_joins)]
965            #[expr($?;Ok(self))]
966            fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result<Self, DataFusionError>;
967
968            #[call(set_distributed_compression)]
969            #[expr($?;Ok(self))]
970            #[cfg(feature = "grpc")]
971            fn with_distributed_compression(mut self, compression: Option<arrow_ipc::CompressionType>) -> Result<Self, DataFusionError>;
972
973            #[call(set_distributed_shuffle_batch_size)]
974            #[expr($?;Ok(self))]
975            fn with_distributed_shuffle_batch_size(mut self, batch_size: usize) -> Result<Self, DataFusionError>;
976
977            #[call(set_distributed_passthrough_headers)]
978            #[expr($?;Ok(self))]
979            fn with_distributed_passthrough_headers(mut self, headers: HeaderMap) -> Result<Self, DataFusionError>;
980
981            #[call(set_distributed_max_tasks_per_stage)]
982            #[expr($?;Ok(self))]
983            fn with_distributed_max_tasks_per_stage(mut self, max_tasks_per_stage: usize) -> Result<Self, DataFusionError>;
984
985            #[call(set_distributed_partial_reduce)]
986            #[expr($?;Ok(self))]
987            fn with_distributed_partial_reduce(mut self, enabled: bool) -> Result<Self, DataFusionError>;
988
989            #[call(set_distributed_worker_connection_buffer_budget_bytes)]
990            #[expr($?;Ok(self))]
991            fn with_distributed_worker_connection_buffer_budget_bytes(mut self, budget_bytes: usize) -> Result<Self, DataFusionError>;
992
993            #[call(set_distributed_work_unit_feed)]
994            #[expr($;self)]
995            fn with_distributed_work_unit_feed<T, P, F>(mut self, getter: F) -> Self
996            where
997                T: ExecutionPlan + 'static,
998                P: WorkUnitFeedProvider + 'static,
999                P::WorkUnit: 'static,
1000                F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static;
1001
1002            #[call(set_distributed_dynamic_task_count)]
1003            #[expr($?;Ok(self))]
1004            fn with_distributed_dynamic_task_count(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1005
1006            #[call(set_distributed_dynamic_bytes_per_partition)]
1007            #[expr($?;Ok(self))]
1008            fn with_distributed_dynamic_bytes_per_partition(mut self, dynamic_bytes_per_partition: usize) -> Result<Self, DataFusionError>;
1009
1010            #[call(set_distributed_local_worker_context)]
1011            #[expr($;self)]
1012            fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self;
1013
1014            #[call(set_distributed_desired_task_count_handler)]
1015            #[expr($;self)]
1016            fn with_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(mut self, h: H) -> Self;
1017
1018            #[call(set_distributed_scale_up_leaf_node_handler)]
1019            #[expr($;self)]
1020            fn with_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(mut self, h: H) -> Self;
1021
1022            #[call(set_distributed_route_tasks_handler)]
1023            #[expr($;self)]
1024            fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(mut self, h: H) -> Self;
1025
1026            #[call(set_distributed_worker_plan_rewrite_handler)]
1027            #[expr($;self)]
1028            fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(mut self, h: H) -> Self;
1029        }
1030    }
1031}
1032impl DistributedGetterExt for SessionConfig {
1033    fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError> {
1034        get_distributed_worker_resolver(self)
1035    }
1036}
1037
1038impl DistributedExt for SessionStateBuilder {
1039    delegate! {
1040        to self.config().get_or_insert_default() {
1041            fn set_distributed_option_extension<T: ConfigExtension + Default>(&mut self, t: T);
1042            #[call(set_distributed_option_extension)]
1043            #[expr($;self)]
1044            fn with_distributed_option_extension<T: ConfigExtension + Default>(mut self, t: T) -> Self;
1045
1046            fn set_distributed_option_extension_from_headers<T: ConfigExtension + Default>(&mut self, h: &HeaderMap) -> Result<(), DataFusionError>;
1047            #[call(set_distributed_option_extension_from_headers)]
1048            #[expr($?;Ok(self))]
1049            fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(mut self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
1050
1051            fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
1052            #[call(set_distributed_user_codec)]
1053            #[expr($;self)]
1054            fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
1055
1056            fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>);
1057            #[call(set_distributed_user_codec_arc)]
1058            #[expr($;self)]
1059            fn with_distributed_user_codec_arc(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
1060
1061            fn set_distributed_worker_resolver<T: WorkerResolver + 'static>(&mut self, resolver: T);
1062            #[call(set_distributed_worker_resolver)]
1063            #[expr($;self)]
1064            fn with_distributed_worker_resolver<T: WorkerResolver + 'static>(mut self, resolver: T) -> Self;
1065
1066            fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
1067            #[call(set_distributed_channel_resolver)]
1068            #[expr($;self)]
1069            fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
1070
1071            fn set_distributed_file_scan_config_bytes_per_partition(&mut self, bytes_per_partition: usize) -> Result<(), DataFusionError>;
1072            #[call(set_distributed_file_scan_config_bytes_per_partition)]
1073            #[expr($?;Ok(self))]
1074            fn with_distributed_file_scan_config_bytes_per_partition(mut self, bytes_per_partition: usize) -> Result<Self, DataFusionError>;
1075
1076            fn set_distributed_cardinality_effect_task_scale_factor(&mut self, factor: f64) -> Result<(), DataFusionError>;
1077            #[call(set_distributed_cardinality_effect_task_scale_factor)]
1078            #[expr($?;Ok(self))]
1079            fn with_distributed_cardinality_effect_task_scale_factor(mut self, factor: f64) -> Result<Self, DataFusionError>;
1080
1081            fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1082            #[call(set_distributed_metrics_collection)]
1083            #[expr($?;Ok(self))]
1084            fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1085
1086            fn set_distributed_children_isolator_unions(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1087            #[call(set_distributed_children_isolator_unions)]
1088            #[expr($?;Ok(self))]
1089            fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1090
1091            fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1092            #[call(set_distributed_broadcast_joins)]
1093            #[expr($?;Ok(self))]
1094            fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1095
1096            #[cfg(feature = "grpc")]
1097            fn set_distributed_compression(&mut self, compression: Option<arrow_ipc::CompressionType>) -> Result<(), DataFusionError>;
1098            #[call(set_distributed_compression)]
1099            #[expr($?;Ok(self))]
1100            #[cfg(feature = "grpc")]
1101            fn with_distributed_compression(mut self, compression: Option<arrow_ipc::CompressionType>) -> Result<Self, DataFusionError>;
1102
1103            fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>;
1104            #[call(set_distributed_shuffle_batch_size)]
1105            #[expr($?;Ok(self))]
1106            fn with_distributed_shuffle_batch_size(mut self, batch_size: usize) -> Result<Self, DataFusionError>;
1107
1108            fn set_distributed_passthrough_headers(&mut self, headers: HeaderMap) -> Result<(), DataFusionError>;
1109            #[call(set_distributed_passthrough_headers)]
1110            #[expr($?;Ok(self))]
1111            fn with_distributed_passthrough_headers(mut self, headers: HeaderMap) -> Result<Self, DataFusionError>;
1112
1113            fn set_distributed_max_tasks_per_stage(&mut self, max_tasks_per_stage: usize) -> Result<(), DataFusionError>;
1114            #[call(set_distributed_max_tasks_per_stage)]
1115            #[expr($?;Ok(self))]
1116            fn with_distributed_max_tasks_per_stage(mut self, max_tasks_per_stage: usize) -> Result<Self, DataFusionError>;
1117
1118            fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1119            #[call(set_distributed_partial_reduce)]
1120            #[expr($?;Ok(self))]
1121            fn with_distributed_partial_reduce(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1122
1123            fn set_distributed_worker_connection_buffer_budget_bytes(&mut self, budget_bytes: usize) -> Result<(), DataFusionError>;
1124            #[call(set_distributed_worker_connection_buffer_budget_bytes)]
1125            #[expr($?;Ok(self))]
1126            fn with_distributed_worker_connection_buffer_budget_bytes(mut self, budget_bytes: usize) -> Result<Self, DataFusionError>;
1127
1128            fn set_distributed_work_unit_feed<T, P, F>(&mut self, getter: F)
1129            where
1130                T: ExecutionPlan + 'static,
1131                P: WorkUnitFeedProvider + 'static,
1132                P::WorkUnit: 'static,
1133                F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static;
1134            #[call(set_distributed_work_unit_feed)]
1135            #[expr($;self)]
1136            fn with_distributed_work_unit_feed<T, P, F>(mut self, getter: F) -> Self
1137            where
1138                T: ExecutionPlan + 'static,
1139                P: WorkUnitFeedProvider + 'static,
1140                P::WorkUnit: 'static,
1141                F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static;
1142
1143            fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1144            #[call(set_distributed_dynamic_task_count)]
1145            #[expr($?;Ok(self))]
1146            fn with_distributed_dynamic_task_count(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1147
1148            fn set_distributed_dynamic_bytes_per_partition(&mut self, dynamic_bytes_per_partition: usize) -> Result<(), DataFusionError>;
1149            #[call(set_distributed_dynamic_bytes_per_partition)]
1150            #[expr($?;Ok(self))]
1151            fn with_distributed_dynamic_bytes_per_partition(mut self, dynamic_bytes_per_partition: usize) -> Result<Self, DataFusionError>;
1152
1153            fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext);
1154            #[call(set_distributed_local_worker_context)]
1155            #[expr($;self)]
1156            fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self;
1157
1158            fn set_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(&mut self, h: H);
1159            #[call(set_distributed_desired_task_count_handler)]
1160            #[expr($;self)]
1161            fn with_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(mut self, h: H) -> Self;
1162
1163            fn set_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(&mut self, h: H);
1164            #[call(set_distributed_scale_up_leaf_node_handler)]
1165            #[expr($;self)]
1166            fn with_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(mut self, h: H) -> Self;
1167
1168            fn set_distributed_route_tasks_handler<H: RouteTasksHandler>(&mut self, h: H);
1169            #[call(set_distributed_route_tasks_handler)]
1170            #[expr($;self)]
1171            fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(mut self, h: H) -> Self;
1172
1173            fn set_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(&mut self, h: H);
1174            #[call(set_distributed_worker_plan_rewrite_handler)]
1175            #[expr($;self)]
1176            fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(mut self, h: H) -> Self;
1177        }
1178    }
1179}
1180impl DistributedGetterExt for SessionState {
1181    delegate! {
1182        to self.config() {
1183            fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError>;
1184        }
1185    }
1186}
1187
1188impl DistributedExt for SessionState {
1189    delegate! {
1190        to self.config_mut() {
1191            fn set_distributed_option_extension<T: ConfigExtension + Default>(&mut self, t: T);
1192            #[call(set_distributed_option_extension)]
1193            #[expr($;self)]
1194            fn with_distributed_option_extension<T: ConfigExtension + Default>(mut self, t: T) -> Self;
1195
1196            fn set_distributed_option_extension_from_headers<T: ConfigExtension + Default>(&mut self, h: &HeaderMap) -> Result<(), DataFusionError>;
1197            #[call(set_distributed_option_extension_from_headers)]
1198            #[expr($?;Ok(self))]
1199            fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(mut self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
1200
1201            fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
1202            #[call(set_distributed_user_codec)]
1203            #[expr($;self)]
1204            fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
1205
1206            fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>);
1207            #[call(set_distributed_user_codec_arc)]
1208            #[expr($;self)]
1209            fn with_distributed_user_codec_arc(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
1210
1211            fn set_distributed_worker_resolver<T: WorkerResolver + 'static>(&mut self, resolver: T);
1212            #[call(set_distributed_worker_resolver)]
1213            #[expr($;self)]
1214            fn with_distributed_worker_resolver<T: WorkerResolver + 'static>(mut self, resolver: T) -> Self;
1215
1216            fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
1217            #[call(set_distributed_channel_resolver)]
1218            #[expr($;self)]
1219            fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
1220
1221            fn set_distributed_file_scan_config_bytes_per_partition(&mut self, bytes_per_partition: usize) -> Result<(), DataFusionError>;
1222            #[call(set_distributed_file_scan_config_bytes_per_partition)]
1223            #[expr($?;Ok(self))]
1224            fn with_distributed_file_scan_config_bytes_per_partition(mut self, bytes_per_partition: usize) -> Result<Self, DataFusionError>;
1225
1226            fn set_distributed_cardinality_effect_task_scale_factor(&mut self, factor: f64) -> Result<(), DataFusionError>;
1227            #[call(set_distributed_cardinality_effect_task_scale_factor)]
1228            #[expr($?;Ok(self))]
1229            fn with_distributed_cardinality_effect_task_scale_factor(mut self, factor: f64) -> Result<Self, DataFusionError>;
1230
1231            fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1232            #[call(set_distributed_metrics_collection)]
1233            #[expr($?;Ok(self))]
1234            fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1235
1236            fn set_distributed_children_isolator_unions(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1237            #[call(set_distributed_children_isolator_unions)]
1238            #[expr($?;Ok(self))]
1239            fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1240
1241            fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1242            #[call(set_distributed_broadcast_joins)]
1243            #[expr($?;Ok(self))]
1244            fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1245
1246            #[cfg(feature = "grpc")]
1247            fn set_distributed_compression(&mut self, compression: Option<arrow_ipc::CompressionType>) -> Result<(), DataFusionError>;
1248            #[call(set_distributed_compression)]
1249            #[expr($?;Ok(self))]
1250            #[cfg(feature = "grpc")]
1251            fn with_distributed_compression(mut self, compression: Option<arrow_ipc::CompressionType>) -> Result<Self, DataFusionError>;
1252
1253            fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>;
1254            #[call(set_distributed_shuffle_batch_size)]
1255            #[expr($?;Ok(self))]
1256            fn with_distributed_shuffle_batch_size(mut self, batch_size: usize) -> Result<Self, DataFusionError>;
1257
1258            fn set_distributed_passthrough_headers(&mut self, headers: HeaderMap) -> Result<(), DataFusionError>;
1259            #[call(set_distributed_passthrough_headers)]
1260            #[expr($?;Ok(self))]
1261            fn with_distributed_passthrough_headers(mut self, headers: HeaderMap) -> Result<Self, DataFusionError>;
1262
1263            fn set_distributed_max_tasks_per_stage(&mut self, max_tasks_per_stage: usize) -> Result<(), DataFusionError>;
1264            #[call(set_distributed_max_tasks_per_stage)]
1265            #[expr($?;Ok(self))]
1266            fn with_distributed_max_tasks_per_stage(mut self, max_tasks_per_stage: usize) -> Result<Self, DataFusionError>;
1267
1268            fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1269            #[call(set_distributed_partial_reduce)]
1270            #[expr($?;Ok(self))]
1271            fn with_distributed_partial_reduce(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1272
1273            fn set_distributed_worker_connection_buffer_budget_bytes(&mut self, budget_bytes: usize) -> Result<(), DataFusionError>;
1274            #[call(set_distributed_worker_connection_buffer_budget_bytes)]
1275            #[expr($?;Ok(self))]
1276            fn with_distributed_worker_connection_buffer_budget_bytes(mut self, budget_bytes: usize) -> Result<Self, DataFusionError>;
1277
1278            fn set_distributed_work_unit_feed<T, P, F>(&mut self, getter: F)
1279            where
1280                T: ExecutionPlan + 'static,
1281                P: WorkUnitFeedProvider + 'static,
1282                P::WorkUnit: 'static,
1283                F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static;
1284            #[call(set_distributed_work_unit_feed)]
1285            #[expr($;self)]
1286            fn with_distributed_work_unit_feed<T, P, F>(mut self, getter: F) -> Self
1287            where
1288                T: ExecutionPlan + 'static,
1289                P: WorkUnitFeedProvider + 'static,
1290                P::WorkUnit: 'static,
1291                F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static;
1292
1293            fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1294            #[call(set_distributed_dynamic_task_count)]
1295            #[expr($?;Ok(self))]
1296            fn with_distributed_dynamic_task_count(mut self, enabled: bool) -> Result<Self, DataFusionError>;
1297
1298            fn set_distributed_dynamic_bytes_per_partition(&mut self, dynamic_bytes_per_partition: usize) -> Result<(), DataFusionError>;
1299            #[call(set_distributed_dynamic_bytes_per_partition)]
1300            #[expr($?;Ok(self))]
1301            fn with_distributed_dynamic_bytes_per_partition(mut self, dynamic_bytes_per_partition: usize) -> Result<Self, DataFusionError>;
1302
1303            fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext);
1304            #[call(set_distributed_local_worker_context)]
1305            #[expr($;self)]
1306            fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self;
1307
1308            fn set_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(&mut self, h: H);
1309            #[call(set_distributed_desired_task_count_handler)]
1310            #[expr($;self)]
1311            fn with_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(mut self, h: H) -> Self;
1312
1313            fn set_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(&mut self, h: H);
1314            #[call(set_distributed_scale_up_leaf_node_handler)]
1315            #[expr($;self)]
1316            fn with_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(mut self, h: H) -> Self;
1317
1318            fn set_distributed_route_tasks_handler<H: RouteTasksHandler>(&mut self, h: H);
1319            #[call(set_distributed_route_tasks_handler)]
1320            #[expr($;self)]
1321            fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(mut self, h: H) -> Self;
1322
1323            fn set_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(&mut self, h: H);
1324            #[call(set_distributed_worker_plan_rewrite_handler)]
1325            #[expr($;self)]
1326            fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(mut self, h: H) -> Self;
1327        }
1328    }
1329}
1330
1331impl DistributedExt for SessionContext {
1332    delegate! {
1333        to self.state_ref().write().config_mut() {
1334            fn set_distributed_option_extension<T: ConfigExtension + Default>(&mut self, t: T);
1335            #[call(set_distributed_option_extension)]
1336            #[expr($;self)]
1337            fn with_distributed_option_extension<T: ConfigExtension + Default>(self, t: T) -> Self;
1338
1339            fn set_distributed_option_extension_from_headers<T: ConfigExtension + Default>(&mut self, h: &HeaderMap) -> Result<(), DataFusionError>;
1340            #[call(set_distributed_option_extension_from_headers)]
1341            #[expr($?;Ok(self))]
1342            fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
1343
1344            fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
1345            #[call(set_distributed_user_codec)]
1346            #[expr($;self)]
1347            fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(self, codec: T) -> Self;
1348
1349            fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>);
1350            #[call(set_distributed_user_codec_arc)]
1351            #[expr($;self)]
1352            fn with_distributed_user_codec_arc(self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
1353
1354            fn set_distributed_worker_resolver<T: WorkerResolver + 'static>(&mut self, resolver: T);
1355            #[call(set_distributed_worker_resolver)]
1356            #[expr($;self)]
1357            fn with_distributed_worker_resolver<T: WorkerResolver + 'static>(self, resolver: T) -> Self;
1358
1359            fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
1360            #[call(set_distributed_channel_resolver)]
1361            #[expr($;self)]
1362            fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(self, resolver: T) -> Self;
1363
1364            fn set_distributed_file_scan_config_bytes_per_partition(&mut self, bytes_per_partition: usize) -> Result<(), DataFusionError>;
1365            #[call(set_distributed_file_scan_config_bytes_per_partition)]
1366            #[expr($?;Ok(self))]
1367            fn with_distributed_file_scan_config_bytes_per_partition(self, bytes_per_partition: usize) -> Result<Self, DataFusionError>;
1368
1369            fn set_distributed_cardinality_effect_task_scale_factor(&mut self, factor: f64) -> Result<(), DataFusionError>;
1370            #[call(set_distributed_cardinality_effect_task_scale_factor)]
1371            #[expr($?;Ok(self))]
1372            fn with_distributed_cardinality_effect_task_scale_factor(self, factor: f64) -> Result<Self, DataFusionError>;
1373
1374            fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1375            #[call(set_distributed_metrics_collection)]
1376            #[expr($?;Ok(self))]
1377            fn with_distributed_metrics_collection(self, enabled: bool) -> Result<Self, DataFusionError>;
1378
1379            fn set_distributed_children_isolator_unions(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1380            #[call(set_distributed_children_isolator_unions)]
1381            #[expr($?;Ok(self))]
1382            fn with_distributed_children_isolator_unions(self, enabled: bool) -> Result<Self, DataFusionError>;
1383
1384            fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1385            #[call(set_distributed_broadcast_joins)]
1386            #[expr($?;Ok(self))]
1387            fn with_distributed_broadcast_joins(self, enabled: bool) -> Result<Self, DataFusionError>;
1388
1389            #[cfg(feature = "grpc")]
1390            fn set_distributed_compression(&mut self, compression: Option<arrow_ipc::CompressionType>) -> Result<(), DataFusionError>;
1391            #[call(set_distributed_compression)]
1392            #[expr($?;Ok(self))]
1393            #[cfg(feature = "grpc")]
1394            fn with_distributed_compression(self, compression: Option<arrow_ipc::CompressionType>) -> Result<Self, DataFusionError>;
1395
1396            fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>;
1397            #[call(set_distributed_shuffle_batch_size)]
1398            #[expr($?;Ok(self))]
1399            fn with_distributed_shuffle_batch_size(self, batch_size: usize) -> Result<Self, DataFusionError>;
1400
1401            fn set_distributed_passthrough_headers(&mut self, headers: HeaderMap) -> Result<(), DataFusionError>;
1402            #[call(set_distributed_passthrough_headers)]
1403            #[expr($?;Ok(self))]
1404            fn with_distributed_passthrough_headers(self, headers: HeaderMap) -> Result<Self, DataFusionError>;
1405
1406            fn set_distributed_max_tasks_per_stage(&mut self, max_tasks_per_stage: usize) -> Result<(), DataFusionError>;
1407            #[call(set_distributed_max_tasks_per_stage)]
1408            #[expr($?;Ok(self))]
1409            fn with_distributed_max_tasks_per_stage(self, max_tasks_per_stage: usize) -> Result<Self, DataFusionError>;
1410
1411            fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1412            #[call(set_distributed_partial_reduce)]
1413            #[expr($?;Ok(self))]
1414            fn with_distributed_partial_reduce(self, enabled: bool) -> Result<Self, DataFusionError>;
1415
1416            fn set_distributed_worker_connection_buffer_budget_bytes(&mut self, budget_bytes: usize) -> Result<(), DataFusionError>;
1417            #[call(set_distributed_worker_connection_buffer_budget_bytes)]
1418            #[expr($?;Ok(self))]
1419            fn with_distributed_worker_connection_buffer_budget_bytes(self, budget_bytes: usize) -> Result<Self, DataFusionError>;
1420
1421            fn set_distributed_work_unit_feed<T, P, F>(&mut self, getter: F)
1422            where
1423                T: ExecutionPlan + 'static,
1424                P: WorkUnitFeedProvider + 'static,
1425                P::WorkUnit: 'static,
1426                F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static;
1427            #[call(set_distributed_work_unit_feed)]
1428            #[expr($;self)]
1429            fn with_distributed_work_unit_feed<T, P, F>(self, getter: F) -> Self
1430            where
1431                T: ExecutionPlan + 'static,
1432                P: WorkUnitFeedProvider + 'static,
1433                P::WorkUnit: 'static,
1434                F: Fn(&T) -> Option<&WorkUnitFeed<P>> + Send + Sync + 'static;
1435
1436            fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError>;
1437            #[call(set_distributed_dynamic_task_count)]
1438            #[expr($?;Ok(self))]
1439            fn with_distributed_dynamic_task_count(self, enabled: bool) -> Result<Self, DataFusionError>;
1440
1441            fn set_distributed_dynamic_bytes_per_partition(&mut self, dynamic_bytes_per_partition: usize) -> Result<(), DataFusionError>;
1442            #[call(set_distributed_dynamic_bytes_per_partition)]
1443            #[expr($?;Ok(self))]
1444            fn with_distributed_dynamic_bytes_per_partition(self, dynamic_bytes_per_partition: usize) -> Result<Self, DataFusionError>;
1445
1446            fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext);
1447            #[call(set_distributed_local_worker_context)]
1448            #[expr($;self)]
1449            fn with_distributed_local_worker_context(self, local_worker_context: LocalWorkerContext) -> Self;
1450
1451            fn set_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(&mut self, h: H);
1452            #[call(set_distributed_desired_task_count_handler)]
1453            #[expr($;self)]
1454            fn with_distributed_desired_task_count_handler<H: DesiredTaskCountHandler>(self, h: H) -> Self;
1455
1456            fn set_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(&mut self, h: H);
1457            #[call(set_distributed_scale_up_leaf_node_handler)]
1458            #[expr($;self)]
1459            fn with_distributed_scale_up_leaf_node_handler<H: ScaleUpLeafNodeHandler>(self, h: H) -> Self;
1460
1461            fn set_distributed_route_tasks_handler<H: RouteTasksHandler>(&mut self, h: H);
1462            #[call(set_distributed_route_tasks_handler)]
1463            #[expr($;self)]
1464            fn with_distributed_route_tasks_handler<H: RouteTasksHandler>(self, h: H) -> Self;
1465
1466            fn set_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(&mut self, h: H);
1467            #[call(set_distributed_worker_plan_rewrite_handler)]
1468            #[expr($;self)]
1469            fn with_distributed_worker_plan_rewrite_handler<H: WorkerPlanRewriteHandler>(self, h: H) -> Self;
1470        }
1471    }
1472}
1473
1474impl DistributedGetterExt for SessionContext {
1475    delegate! {
1476        to self.state_ref().read().config() {
1477            fn get_distributed_worker_resolver(&self) -> Result<Arc<dyn WorkerResolver>, DataFusionError>;
1478        }
1479    }
1480}
datafusion_distributed/distributed_ext.rs

datafusion_distributed/
distributed_ext.rs
