Skip to content

Commit aa5231b

Browse files
Rich-T-kidRich-T-kid
committed
revert changes
1 parent 302551b commit aa5231b

22 files changed

Lines changed: 352 additions & 498 deletions

‎src/config_extension_ext.rs‎

Lines changed: 0 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
use crate::DistributedConfig;
21
use datafusion::common::{DataFusionError, HashSet, internal_datafusion_err};
32
use datafusion::config::ConfigExtension;
43
use datafusion::prelude::SessionConfig;
@@ -129,115 +128,6 @@ pub(crate) fn get_config_extension_propagation_headers(
129128
Ok(headers)
130129
}
131130

132-
/// Serializes the scalar fields of [DistributedConfig] from `SessionConfig.extensions`
133-
/// into gRPC metadata headers so workers can reconstruct the config.
134-
pub(crate) fn get_distributed_config_propagation_headers(
135-
cfg: &SessionConfig,
136-
) -> Result<HeaderMap, DataFusionError> {
137-
fn insert(headers: &mut HeaderMap, field: &str, value: &str) -> Result<(), DataFusionError> {
138-
let name = format!("{FLIGHT_METADATA_CONFIG_PREFIX}distributed.{field}");
139-
headers.insert(
140-
HeaderName::from_str(&name)
141-
.map_err(|e| internal_datafusion_err!("invalid header name: {e}"))?,
142-
value
143-
.parse()
144-
.map_err(|e| internal_datafusion_err!("invalid header value: {e}"))?,
145-
);
146-
Ok(())
147-
}
148-
149-
let Some(d) = cfg.get_extension::<DistributedConfig>() else {
150-
return Ok(HeaderMap::new());
151-
};
152-
let mut headers = HeaderMap::new();
153-
let h = &mut headers;
154-
insert(
155-
h,
156-
"file_scan_config_bytes_per_partition",
157-
&d.file_scan_config_bytes_per_partition.to_string(),
158-
)?;
159-
insert(
160-
h,
161-
"cardinality_task_count_factor",
162-
&d.cardinality_task_count_factor.to_string(),
163-
)?;
164-
insert(
165-
h,
166-
"children_isolator_unions",
167-
&d.children_isolator_unions.to_string(),
168-
)?;
169-
insert(h, "collect_metrics", &d.collect_metrics.to_string())?;
170-
insert(h, "broadcast_joins", &d.broadcast_joins.to_string())?;
171-
insert(h, "compression", &d.compression)?;
172-
insert(h, "shuffle_batch_size", &d.shuffle_batch_size.to_string())?;
173-
insert(h, "max_tasks_per_stage", &d.max_tasks_per_stage.to_string())?;
174-
insert(h, "partial_reduce", &d.partial_reduce.to_string())?;
175-
insert(
176-
h,
177-
"worker_connection_buffer_budget_bytes",
178-
&d.worker_connection_buffer_budget_bytes.to_string(),
179-
)?;
180-
insert(h, "dynamic_task_count", &d.dynamic_task_count.to_string())?;
181-
insert(
182-
h,
183-
"bytes_per_partition_per_second",
184-
&d.bytes_per_partition_per_second.to_string(),
185-
)?;
186-
Ok(headers)
187-
}
188-
189-
/// Deserializes [DistributedConfig] scalar fields from gRPC metadata headers into
190-
/// `SessionConfig.extensions`. Called on the worker side when a task arrives.
191-
pub(crate) fn set_distributed_config_from_headers(
192-
cfg: &mut SessionConfig,
193-
headers: &HeaderMap,
194-
) -> Result<(), DataFusionError> {
195-
let prefix = format!("{FLIGHT_METADATA_CONFIG_PREFIX}distributed.");
196-
let mut d = cfg
197-
.get_extension::<DistributedConfig>()
198-
.map(|a| a.as_ref().clone())
199-
.unwrap_or_default();
200-
201-
for (k, v) in headers.iter() {
202-
let key = k.as_str();
203-
if !key.starts_with(&prefix) {
204-
continue;
205-
}
206-
let field = &key[prefix.len()..];
207-
let value = v
208-
.to_str()
209-
.map_err(|e| internal_datafusion_err!("invalid header value: {e}"))?;
210-
macro_rules! parse {
211-
($target:expr) => {{
212-
$target = value
213-
.parse()
214-
.map_err(|e| internal_datafusion_err!("bad header {field}: {e}"))?;
215-
}};
216-
}
217-
match field {
218-
"file_scan_config_bytes_per_partition" => {
219-
parse!(d.file_scan_config_bytes_per_partition)
220-
}
221-
"cardinality_task_count_factor" => parse!(d.cardinality_task_count_factor),
222-
"children_isolator_unions" => parse!(d.children_isolator_unions),
223-
"collect_metrics" => parse!(d.collect_metrics),
224-
"broadcast_joins" => parse!(d.broadcast_joins),
225-
"compression" => d.compression = value.to_string(),
226-
"shuffle_batch_size" => parse!(d.shuffle_batch_size),
227-
"max_tasks_per_stage" => parse!(d.max_tasks_per_stage),
228-
"partial_reduce" => parse!(d.partial_reduce),
229-
"worker_connection_buffer_budget_bytes" => {
230-
parse!(d.worker_connection_buffer_budget_bytes)
231-
}
232-
"dynamic_task_count" => parse!(d.dynamic_task_count),
233-
"bytes_per_partition_per_second" => parse!(d.bytes_per_partition_per_second),
234-
_ => {}
235-
}
236-
}
237-
cfg.set_extension(Arc::new(d));
238-
Ok(())
239-
}
240-
241131
#[cfg(test)]
242132
mod tests {
243133
use crate::config_extension_ext::{

‎src/coordinator/distributed.rs‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -209,10 +209,7 @@ impl ExecutionPlan for DistributedExec {
209209
// channel is closed.
210210
let guard = query_coordinator.end_query_guard();
211211

212-
let d_cfg = context
213-
.session_config()
214-
.get_extension::<DistributedConfig>()
215-
.expect("DistributedConfig should be set");
212+
let d_cfg = DistributedConfig::from_config_options(context.session_config().options())?;
216213
let result = match d_cfg.dynamic_task_count {
217214
true => prepare_dynamic_plan(&query_coordinator, &base_plan).await?,
218215
false => prepare_static_plan(&query_coordinator, &base_plan)?,

‎src/coordinator/prepare_dynamic_plan.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ pub(super) async fn prepare_dynamic_plan(
140140
})
141141
})
142142
},
143-
query_coordinator.session_config(),
143+
query_coordinator.session_config().options(),
144144
)
145145
.await?;
146146

‎src/coordinator/query_coordinator.rs‎

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension};
2-
use crate::config_extension_ext::{
3-
get_config_extension_propagation_headers, get_distributed_config_propagation_headers,
4-
};
2+
use crate::config_extension_ext::get_config_extension_propagation_headers;
53
use crate::coordinator::MetricsStore;
64
use crate::coordinator::latency_metric::LatencyMetric;
75
use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec};
@@ -173,7 +171,6 @@ impl<'a> StageCoordinator<'a> {
173171
let channel_resolver = get_distributed_channel_resolver(self.task_ctx.as_ref());
174172

175173
let mut headers = get_config_extension_propagation_headers(session_config)?;
176-
headers.extend(get_distributed_config_propagation_headers(session_config)?);
177174
headers.extend(get_passthrough_headers(session_config));
178175

179176
let coordinator_to_worker_stream = futures::stream::once(async { msg })
@@ -264,9 +261,7 @@ impl<'a> StageCoordinator<'a> {
264261
tx: UnboundedSender<CoordinatorToWorkerMsg>,
265262
) -> Result<()> {
266263
let session_config = self.task_ctx.session_config();
267-
let d_cfg = session_config
268-
.get_extension::<DistributedConfig>()
269-
.expect("DistributedConfig should be set");
264+
let d_cfg = DistributedConfig::from_config_options(session_config.options())?;
270265
let wuf_registry = &d_cfg.__private_work_unit_feed_registry;
271266

272267
let d_ctx = DistributedTaskContext {
@@ -338,9 +333,7 @@ impl<'a> StageCoordinator<'a> {
338333
task_i: usize,
339334
) -> Result<(Arc<dyn ExecutionPlan>, Vec<WorkUnitFeedDeclaration>)> {
340335
let session_config = self.task_ctx.session_config();
341-
let d_cfg = session_config
342-
.get_extension::<DistributedConfig>()
343-
.expect("DistributedConfig should be set");
336+
let d_cfg = DistributedConfig::from_config_options(session_config.options())?;
344337
let wuf_registry = &d_cfg.__private_work_unit_feed_registry;
345338

346339
let mut work_unit_feed_declarations = vec![];
@@ -380,9 +373,7 @@ impl<'a> StageCoordinator<'a> {
380373
/// [TaskEstimator::route_tasks] method.
381374
pub(super) fn routed_urls(&self) -> Result<Vec<Url>> {
382375
let session_config = self.task_ctx.session_config();
383-
let d_cfg = session_config
384-
.get_extension::<DistributedConfig>()
385-
.expect("DistributedConfig should be set");
376+
let d_cfg = DistributedConfig::from_config_options(session_config.options())?;
386377
let worker_resolver = get_distributed_worker_resolver(session_config)?;
387378
let task_estimator = &d_cfg.__private_task_estimator;
388379

‎src/distributed_ext.rs‎

Lines changed: 12 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -645,58 +645,38 @@ impl DistributedExt for SessionConfig {
645645
&mut self,
646646
bytes_per_partition: usize,
647647
) -> Result<(), DataFusionError> {
648-
let mut d_cfg = self
649-
.get_extension::<DistributedConfig>()
650-
.map(|a| a.as_ref().clone())
651-
.unwrap_or_default();
648+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
652649
d_cfg.file_scan_config_bytes_per_partition = bytes_per_partition;
653-
self.set_extension(Arc::new(d_cfg));
654650
Ok(())
655651
}
656652

657653
fn set_distributed_cardinality_effect_task_scale_factor(
658654
&mut self,
659655
factor: f64,
660656
) -> Result<(), DataFusionError> {
661-
let mut d_cfg = self
662-
.get_extension::<DistributedConfig>()
663-
.map(|a| a.as_ref().clone())
664-
.unwrap_or_default();
657+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
665658
d_cfg.cardinality_task_count_factor = factor;
666-
self.set_extension(Arc::new(d_cfg));
667659
Ok(())
668660
}
669661

670662
fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError> {
671-
let mut d_cfg = self
672-
.get_extension::<DistributedConfig>()
673-
.map(|a| a.as_ref().clone())
674-
.unwrap_or_default();
663+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
675664
d_cfg.collect_metrics = enabled;
676-
self.set_extension(Arc::new(d_cfg));
677665
Ok(())
678666
}
679667

680668
fn set_distributed_children_isolator_unions(
681669
&mut self,
682670
enabled: bool,
683671
) -> Result<(), DataFusionError> {
684-
let mut d_cfg = self
685-
.get_extension::<DistributedConfig>()
686-
.map(|a| a.as_ref().clone())
687-
.unwrap_or_default();
672+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
688673
d_cfg.children_isolator_unions = enabled;
689-
self.set_extension(Arc::new(d_cfg));
690674
Ok(())
691675
}
692676

693677
fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError> {
694-
let mut d_cfg = self
695-
.get_extension::<DistributedConfig>()
696-
.map(|a| a.as_ref().clone())
697-
.unwrap_or_default();
678+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
698679
d_cfg.broadcast_joins = enabled;
699-
self.set_extension(Arc::new(d_cfg));
700680
Ok(())
701681
}
702682

@@ -705,29 +685,21 @@ impl DistributedExt for SessionConfig {
705685
&mut self,
706686
compression: Option<arrow_ipc::CompressionType>,
707687
) -> Result<(), DataFusionError> {
708-
let mut d_cfg = self
709-
.get_extension::<DistributedConfig>()
710-
.map(|a| a.as_ref().clone())
711-
.unwrap_or_default();
688+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
712689
d_cfg.compression = match compression {
713690
Some(arrow_ipc::CompressionType::ZSTD) => "zstd".to_string(),
714691
Some(arrow_ipc::CompressionType::LZ4_FRAME) => "lz4".to_string(),
715692
_ => "none".to_string(),
716693
};
717-
self.set_extension(Arc::new(d_cfg));
718694
Ok(())
719695
}
720696

721697
fn set_distributed_shuffle_batch_size(
722698
&mut self,
723699
batch_size: usize,
724700
) -> Result<(), DataFusionError> {
725-
let mut d_cfg = self
726-
.get_extension::<DistributedConfig>()
727-
.map(|a| a.as_ref().clone())
728-
.unwrap_or_default();
701+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
729702
d_cfg.shuffle_batch_size = batch_size;
730-
self.set_extension(Arc::new(d_cfg));
731703
Ok(())
732704
}
733705

@@ -742,35 +714,23 @@ impl DistributedExt for SessionConfig {
742714
&mut self,
743715
max_tasks_per_stage: usize,
744716
) -> Result<(), DataFusionError> {
745-
let mut d_cfg = self
746-
.get_extension::<DistributedConfig>()
747-
.map(|a| a.as_ref().clone())
748-
.unwrap_or_default();
717+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
749718
d_cfg.max_tasks_per_stage = max_tasks_per_stage;
750-
self.set_extension(Arc::new(d_cfg));
751719
Ok(())
752720
}
753721

754722
fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError> {
755-
let mut d_cfg = self
756-
.get_extension::<DistributedConfig>()
757-
.map(|a| a.as_ref().clone())
758-
.unwrap_or_default();
723+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
759724
d_cfg.partial_reduce = enabled;
760-
self.set_extension(Arc::new(d_cfg));
761725
Ok(())
762726
}
763727

764728
fn set_distributed_worker_connection_buffer_budget_bytes(
765729
&mut self,
766730
budget_bytes: usize,
767731
) -> Result<(), DataFusionError> {
768-
let mut d_cfg = self
769-
.get_extension::<DistributedConfig>()
770-
.map(|a| a.as_ref().clone())
771-
.unwrap_or_default();
732+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
772733
d_cfg.worker_connection_buffer_budget_bytes = budget_bytes;
773-
self.set_extension(Arc::new(d_cfg));
774734
Ok(())
775735
}
776736

@@ -787,25 +747,17 @@ impl DistributedExt for SessionConfig {
787747
}
788748

789749
fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError> {
790-
let mut d_cfg = self
791-
.get_extension::<DistributedConfig>()
792-
.map(|a| a.as_ref().clone())
793-
.unwrap_or_default();
750+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
794751
d_cfg.dynamic_task_count = enabled;
795-
self.set_extension(Arc::new(d_cfg));
796752
Ok(())
797753
}
798754

799755
fn set_distributed_bytes_per_partition_per_second(
800756
&mut self,
801757
bytes_per_partition_per_second: usize,
802758
) -> Result<(), DataFusionError> {
803-
let mut d_cfg = self
804-
.get_extension::<DistributedConfig>()
805-
.map(|a| a.as_ref().clone())
806-
.unwrap_or_default();
759+
let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
807760
d_cfg.bytes_per_partition_per_second = bytes_per_partition_per_second;
808-
self.set_extension(Arc::new(d_cfg));
809761
Ok(())
810762
}
811763

0 commit comments

Comments
 (0)