/Users/andrewlamb/Software/datafusion/datafusion/execution/src/config.rs
Line | Count | Source (jump to first uncovered line) |
1 | | // Licensed to the Apache Software Foundation (ASF) under one |
2 | | // or more contributor license agreements. See the NOTICE file |
3 | | // distributed with this work for additional information |
4 | | // regarding copyright ownership. The ASF licenses this file |
5 | | // to you under the Apache License, Version 2.0 (the |
6 | | // "License"); you may not use this file except in compliance |
7 | | // with the License. You may obtain a copy of the License at |
8 | | // |
9 | | // http://www.apache.org/licenses/LICENSE-2.0 |
10 | | // |
11 | | // Unless required by applicable law or agreed to in writing, |
12 | | // software distributed under the License is distributed on an |
13 | | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
14 | | // KIND, either express or implied. See the License for the |
15 | | // specific language governing permissions and limitations |
16 | | // under the License. |
17 | | |
18 | | use std::{ |
19 | | any::{Any, TypeId}, |
20 | | collections::HashMap, |
21 | | hash::{BuildHasherDefault, Hasher}, |
22 | | sync::Arc, |
23 | | }; |
24 | | |
25 | | use datafusion_common::{ |
26 | | config::{ConfigExtension, ConfigOptions}, |
27 | | Result, ScalarValue, |
28 | | }; |
29 | | |
30 | | /// Configuration options for [`SessionContext`]. |
31 | | /// |
32 | | /// Can be passed to [`SessionContext::new_with_config`] to customize the configuration of DataFusion. |
33 | | /// |
34 | | /// Options can be set using namespaces keys with `.` as the separator, where the |
35 | | /// namespace determines which configuration struct the value to routed to. All |
36 | | /// built-in options are under the `datafusion` namespace. |
37 | | /// |
38 | | /// For example, the key `datafusion.execution.batch_size` will set [ExecutionOptions::batch_size][datafusion_common::config::ExecutionOptions::batch_size], |
39 | | /// because [ConfigOptions::execution] is [ExecutionOptions][datafusion_common::config::ExecutionOptions]. Similarly, the key |
40 | | /// `datafusion.execution.parquet.pushdown_filters` will set [ParquetOptions::pushdown_filters][datafusion_common::config::ParquetOptions::pushdown_filters], |
41 | | /// since [ExecutionOptions::parquet][datafusion_common::config::ExecutionOptions::parquet] is [ParquetOptions][datafusion_common::config::ParquetOptions]. |
42 | | /// |
43 | | /// Some options have convenience methods. For example [SessionConfig::with_batch_size] is |
44 | | /// shorthand for setting `datafusion.execution.batch_size`. |
45 | | /// |
46 | | /// ``` |
47 | | /// use datafusion_execution::config::SessionConfig; |
48 | | /// use datafusion_common::ScalarValue; |
49 | | /// |
50 | | /// let config = SessionConfig::new() |
51 | | /// .set("datafusion.execution.batch_size", &ScalarValue::UInt64(Some(1234))) |
52 | | /// .set_bool("datafusion.execution.parquet.pushdown_filters", true); |
53 | | /// |
54 | | /// assert_eq!(config.batch_size(), 1234); |
55 | | /// assert_eq!(config.options().execution.batch_size, 1234); |
56 | | /// assert_eq!(config.options().execution.parquet.pushdown_filters, true); |
57 | | /// ``` |
58 | | /// |
59 | | /// You can also directly mutate the options via [SessionConfig::options_mut]. |
60 | | /// So the following is equivalent to the above: |
61 | | /// |
62 | | /// ``` |
63 | | /// # use datafusion_execution::config::SessionConfig; |
64 | | /// # use datafusion_common::ScalarValue; |
65 | | /// # |
66 | | /// let mut config = SessionConfig::new(); |
67 | | /// config.options_mut().execution.batch_size = 1234; |
68 | | /// config.options_mut().execution.parquet.pushdown_filters = true; |
69 | | /// # |
70 | | /// # assert_eq!(config.batch_size(), 1234); |
71 | | /// # assert_eq!(config.options().execution.batch_size, 1234); |
72 | | /// # assert_eq!(config.options().execution.parquet.pushdown_filters, true); |
73 | | /// ``` |
74 | | /// |
75 | | /// ## Built-in options |
76 | | /// |
77 | | /// | Namespace | Config struct | |
78 | | /// | --------- | ------------- | |
79 | | /// | `datafusion.catalog` | [CatalogOptions][datafusion_common::config::CatalogOptions] | |
80 | | /// | `datafusion.execution` | [ExecutionOptions][datafusion_common::config::ExecutionOptions] | |
81 | | /// | `datafusion.execution.parquet` | [ParquetOptions][datafusion_common::config::ParquetOptions] | |
82 | | /// | `datafusion.optimizer` | [OptimizerOptions][datafusion_common::config::OptimizerOptions] | |
83 | | /// | `datafusion.sql_parser` | [SqlParserOptions][datafusion_common::config::SqlParserOptions] | |
84 | | /// | `datafusion.explain` | [ExplainOptions][datafusion_common::config::ExplainOptions] | |
85 | | /// |
86 | | /// ## Custom configuration |
87 | | /// |
88 | | /// Configuration options can be extended. See [SessionConfig::with_extension] for details. |
89 | | /// |
90 | | /// [`SessionContext`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html |
91 | | /// [`SessionContext::new_with_config`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html#method.new_with_config |
92 | | #[derive(Clone, Debug)] |
93 | | pub struct SessionConfig { |
94 | | /// Configuration options |
95 | | options: ConfigOptions, |
96 | | /// Opaque extensions. |
97 | | extensions: AnyMap, |
98 | | } |
99 | | |
100 | | impl Default for SessionConfig { |
101 | 1.38k | fn default() -> Self { |
102 | 1.38k | Self { |
103 | 1.38k | options: ConfigOptions::new(), |
104 | 1.38k | // Assume no extensions by default. |
105 | 1.38k | extensions: HashMap::with_capacity_and_hasher( |
106 | 1.38k | 0, |
107 | 1.38k | BuildHasherDefault::default(), |
108 | 1.38k | ), |
109 | 1.38k | } |
110 | 1.38k | } |
111 | | } |
112 | | |
113 | | impl SessionConfig { |
114 | | /// Create an execution config with default setting |
115 | 1.07k | pub fn new() -> Self { |
116 | 1.07k | Default::default() |
117 | 1.07k | } |
118 | | |
119 | | /// Create an execution config with config options read from the environment |
120 | 0 | pub fn from_env() -> Result<Self> { |
121 | 0 | Ok(ConfigOptions::from_env()?.into()) |
122 | 0 | } |
123 | | |
124 | | /// Create new ConfigOptions struct, taking values from a string hash map. |
125 | 0 | pub fn from_string_hash_map(settings: &HashMap<String, String>) -> Result<Self> { |
126 | 0 | Ok(ConfigOptions::from_string_hash_map(settings)?.into()) |
127 | 0 | } |
128 | | |
129 | | /// Return a handle to the configuration options. |
130 | | /// |
131 | | /// Can be used to read the current configuration. |
132 | | /// |
133 | | /// ``` |
134 | | /// use datafusion_execution::config::SessionConfig; |
135 | | /// |
136 | | /// let config = SessionConfig::new(); |
137 | | /// assert!(config.options().execution.batch_size > 0); |
138 | | /// ``` |
139 | 72 | pub fn options(&self) -> &ConfigOptions { |
140 | 72 | &self.options |
141 | 72 | } |
142 | | |
143 | | /// Return a mutable handle to the configuration options. |
144 | | /// |
145 | | /// Can be used to set configuration options. |
146 | | /// |
147 | | /// ``` |
148 | | /// use datafusion_execution::config::SessionConfig; |
149 | | /// |
150 | | /// let mut config = SessionConfig::new(); |
151 | | /// config.options_mut().execution.batch_size = 1024; |
152 | | /// assert_eq!(config.options().execution.batch_size, 1024); |
153 | | /// ``` |
154 | 0 | pub fn options_mut(&mut self) -> &mut ConfigOptions { |
155 | 0 | &mut self.options |
156 | 0 | } |
157 | | |
158 | | /// Set a configuration option |
159 | 4 | pub fn set(self, key: &str, value: &ScalarValue) -> Self { |
160 | 4 | self.set_str(key, &value.to_string()) |
161 | 4 | } |
162 | | |
163 | | /// Set a boolean configuration option |
164 | 0 | pub fn set_bool(self, key: &str, value: bool) -> Self { |
165 | 0 | self.set_str(key, &value.to_string()) |
166 | 0 | } |
167 | | |
168 | | /// Set a generic `u64` configuration option |
169 | 0 | pub fn set_u64(self, key: &str, value: u64) -> Self { |
170 | 0 | self.set_str(key, &value.to_string()) |
171 | 0 | } |
172 | | |
173 | | /// Set a generic `usize` configuration option |
174 | 0 | pub fn set_usize(self, key: &str, value: usize) -> Self { |
175 | 0 | self.set_str(key, &value.to_string()) |
176 | 0 | } |
177 | | |
178 | | /// Set a generic `str` configuration option |
179 | 4 | pub fn set_str(mut self, key: &str, value: &str) -> Self { |
180 | 4 | self.options.set(key, value).unwrap(); |
181 | 4 | self |
182 | 4 | } |
183 | | |
184 | | /// Customize batch size |
185 | 319 | pub fn with_batch_size(mut self, n: usize) -> Self { |
186 | 319 | // batch size must be greater than zero |
187 | 319 | assert!(n > 0); |
188 | 319 | self.options.execution.batch_size = n; |
189 | 319 | self |
190 | 319 | } |
191 | | |
192 | | /// Customize [`target_partitions`] |
193 | | /// |
194 | | /// [`target_partitions`]: datafusion_common::config::ExecutionOptions::target_partitions |
195 | 2 | pub fn with_target_partitions(mut self, n: usize) -> Self { |
196 | 2 | // partition count must be greater than zero |
197 | 2 | assert!(n > 0); |
198 | 2 | self.options.execution.target_partitions = n; |
199 | 2 | self |
200 | 2 | } |
201 | | |
202 | | /// Insert new [ConfigExtension] |
203 | 0 | pub fn with_option_extension<T: ConfigExtension>(mut self, extension: T) -> Self { |
204 | 0 | self.options_mut().extensions.insert(extension); |
205 | 0 | self |
206 | 0 | } |
207 | | |
208 | | /// Get [`target_partitions`] |
209 | | /// |
210 | | /// [`target_partitions`]: datafusion_common::config::ExecutionOptions::target_partitions |
211 | 0 | pub fn target_partitions(&self) -> usize { |
212 | 0 | self.options.execution.target_partitions |
213 | 0 | } |
214 | | |
215 | | /// Is the information schema enabled? |
216 | 0 | pub fn information_schema(&self) -> bool { |
217 | 0 | self.options.catalog.information_schema |
218 | 0 | } |
219 | | |
220 | | /// Should the context create the default catalog and schema? |
221 | 0 | pub fn create_default_catalog_and_schema(&self) -> bool { |
222 | 0 | self.options.catalog.create_default_catalog_and_schema |
223 | 0 | } |
224 | | |
225 | | /// Are joins repartitioned during execution? |
226 | 0 | pub fn repartition_joins(&self) -> bool { |
227 | 0 | self.options.optimizer.repartition_joins |
228 | 0 | } |
229 | | |
230 | | /// Are aggregates repartitioned during execution? |
231 | 0 | pub fn repartition_aggregations(&self) -> bool { |
232 | 0 | self.options.optimizer.repartition_aggregations |
233 | 0 | } |
234 | | |
235 | | /// Are window functions repartitioned during execution? |
236 | 0 | pub fn repartition_window_functions(&self) -> bool { |
237 | 0 | self.options.optimizer.repartition_windows |
238 | 0 | } |
239 | | |
240 | | /// Do we execute sorts in a per-partition fashion and merge afterwards, |
241 | | /// or do we coalesce partitions first and sort globally? |
242 | 0 | pub fn repartition_sorts(&self) -> bool { |
243 | 0 | self.options.optimizer.repartition_sorts |
244 | 0 | } |
245 | | |
246 | | /// Prefer existing sort (true) or maximize parallelism (false). See |
247 | | /// [prefer_existing_sort] for more details |
248 | | /// |
249 | | /// [prefer_existing_sort]: datafusion_common::config::OptimizerOptions::prefer_existing_sort |
250 | 0 | pub fn prefer_existing_sort(&self) -> bool { |
251 | 0 | self.options.optimizer.prefer_existing_sort |
252 | 0 | } |
253 | | |
254 | | /// Are statistics collected during execution? |
255 | 0 | pub fn collect_statistics(&self) -> bool { |
256 | 0 | self.options.execution.collect_statistics |
257 | 0 | } |
258 | | |
259 | | /// Selects a name for the default catalog and schema |
260 | 0 | pub fn with_default_catalog_and_schema( |
261 | 0 | mut self, |
262 | 0 | catalog: impl Into<String>, |
263 | 0 | schema: impl Into<String>, |
264 | 0 | ) -> Self { |
265 | 0 | self.options.catalog.default_catalog = catalog.into(); |
266 | 0 | self.options.catalog.default_schema = schema.into(); |
267 | 0 | self |
268 | 0 | } |
269 | | |
270 | | /// Controls whether the default catalog and schema will be automatically created |
271 | 0 | pub fn with_create_default_catalog_and_schema(mut self, create: bool) -> Self { |
272 | 0 | self.options.catalog.create_default_catalog_and_schema = create; |
273 | 0 | self |
274 | 0 | } |
275 | | |
276 | | /// Enables or disables the inclusion of `information_schema` virtual tables |
277 | 0 | pub fn with_information_schema(mut self, enabled: bool) -> Self { |
278 | 0 | self.options.catalog.information_schema = enabled; |
279 | 0 | self |
280 | 0 | } |
281 | | |
282 | | /// Enables or disables the use of repartitioning for joins to improve parallelism |
283 | 165 | pub fn with_repartition_joins(mut self, enabled: bool) -> Self { |
284 | 165 | self.options.optimizer.repartition_joins = enabled; |
285 | 165 | self |
286 | 165 | } |
287 | | |
288 | | /// Enables or disables the use of repartitioning for aggregations to improve parallelism |
289 | 0 | pub fn with_repartition_aggregations(mut self, enabled: bool) -> Self { |
290 | 0 | self.options.optimizer.repartition_aggregations = enabled; |
291 | 0 | self |
292 | 0 | } |
293 | | |
294 | | /// Sets minimum file range size for repartitioning scans |
295 | 0 | pub fn with_repartition_file_min_size(mut self, size: usize) -> Self { |
296 | 0 | self.options.optimizer.repartition_file_min_size = size; |
297 | 0 | self |
298 | 0 | } |
299 | | |
300 | | /// Enables or disables the allowing unordered symmetric hash join |
301 | 0 | pub fn with_allow_symmetric_joins_without_pruning(mut self, enabled: bool) -> Self { |
302 | 0 | self.options.optimizer.allow_symmetric_joins_without_pruning = enabled; |
303 | 0 | self |
304 | 0 | } |
305 | | |
306 | | /// Enables or disables the use of repartitioning for file scans |
307 | 0 | pub fn with_repartition_file_scans(mut self, enabled: bool) -> Self { |
308 | 0 | self.options.optimizer.repartition_file_scans = enabled; |
309 | 0 | self |
310 | 0 | } |
311 | | |
312 | | /// Enables or disables the use of repartitioning for window functions to improve parallelism |
313 | 0 | pub fn with_repartition_windows(mut self, enabled: bool) -> Self { |
314 | 0 | self.options.optimizer.repartition_windows = enabled; |
315 | 0 | self |
316 | 0 | } |
317 | | |
318 | | /// Enables or disables the use of per-partition sorting to improve parallelism |
319 | 0 | pub fn with_repartition_sorts(mut self, enabled: bool) -> Self { |
320 | 0 | self.options.optimizer.repartition_sorts = enabled; |
321 | 0 | self |
322 | 0 | } |
323 | | |
324 | | /// Prefer existing sort (true) or maximize parallelism (false). See |
325 | | /// [prefer_existing_sort] for more details |
326 | | /// |
327 | | /// [prefer_existing_sort]: datafusion_common::config::OptimizerOptions::prefer_existing_sort |
328 | 0 | pub fn with_prefer_existing_sort(mut self, enabled: bool) -> Self { |
329 | 0 | self.options.optimizer.prefer_existing_sort = enabled; |
330 | 0 | self |
331 | 0 | } |
332 | | |
333 | | /// Prefer existing union (true). See [prefer_existing_union] for more details |
334 | | /// |
335 | | /// [prefer_existing_union]: datafusion_common::config::OptimizerOptions::prefer_existing_union |
336 | 0 | pub fn with_prefer_existing_union(mut self, enabled: bool) -> Self { |
337 | 0 | self.options.optimizer.prefer_existing_union = enabled; |
338 | 0 | self |
339 | 0 | } |
340 | | |
341 | | /// Enables or disables the use of pruning predicate for parquet readers to skip row groups |
342 | 0 | pub fn with_parquet_pruning(mut self, enabled: bool) -> Self { |
343 | 0 | self.options.execution.parquet.pruning = enabled; |
344 | 0 | self |
345 | 0 | } |
346 | | |
347 | | /// Returns true if pruning predicate should be used to skip parquet row groups |
348 | 0 | pub fn parquet_pruning(&self) -> bool { |
349 | 0 | self.options.execution.parquet.pruning |
350 | 0 | } |
351 | | |
352 | | /// Returns true if bloom filter should be used to skip parquet row groups |
353 | 0 | pub fn parquet_bloom_filter_pruning(&self) -> bool { |
354 | 0 | self.options.execution.parquet.bloom_filter_on_read |
355 | 0 | } |
356 | | |
357 | | /// Enables or disables the use of bloom filter for parquet readers to skip row groups |
358 | 0 | pub fn with_parquet_bloom_filter_pruning(mut self, enabled: bool) -> Self { |
359 | 0 | self.options.execution.parquet.bloom_filter_on_read = enabled; |
360 | 0 | self |
361 | 0 | } |
362 | | |
363 | | /// Returns true if page index should be used to skip parquet data pages |
364 | 0 | pub fn parquet_page_index_pruning(&self) -> bool { |
365 | 0 | self.options.execution.parquet.enable_page_index |
366 | 0 | } |
367 | | |
368 | | /// Enables or disables the use of page index for parquet readers to skip parquet data pages |
369 | 0 | pub fn with_parquet_page_index_pruning(mut self, enabled: bool) -> Self { |
370 | 0 | self.options.execution.parquet.enable_page_index = enabled; |
371 | 0 | self |
372 | 0 | } |
373 | | |
374 | | /// Enables or disables the collection of statistics after listing files |
375 | 0 | pub fn with_collect_statistics(mut self, enabled: bool) -> Self { |
376 | 0 | self.options.execution.collect_statistics = enabled; |
377 | 0 | self |
378 | 0 | } |
379 | | |
380 | | /// Get the currently configured batch size |
381 | 1.94k | pub fn batch_size(&self) -> usize { |
382 | 1.94k | self.options.execution.batch_size |
383 | 1.94k | } |
384 | | |
385 | | /// Enables or disables the coalescence of small batches into larger batches |
386 | 0 | pub fn with_coalesce_batches(mut self, enabled: bool) -> Self { |
387 | 0 | self.options.execution.coalesce_batches = enabled; |
388 | 0 | self |
389 | 0 | } |
390 | | |
391 | | /// Returns true if record batches will be examined between each operator |
392 | | /// and small batches will be coalesced into larger batches. |
393 | 0 | pub fn coalesce_batches(&self) -> bool { |
394 | 0 | self.options.execution.coalesce_batches |
395 | 0 | } |
396 | | |
397 | | /// Enables or disables the round robin repartition for increasing parallelism |
398 | 1 | pub fn with_round_robin_repartition(mut self, enabled: bool) -> Self { |
399 | 1 | self.options.optimizer.enable_round_robin_repartition = enabled; |
400 | 1 | self |
401 | 1 | } |
402 | | |
403 | | /// Returns true if the physical plan optimizer will try to |
404 | | /// add round robin repartition to increase parallelism to leverage more CPU cores. |
405 | 0 | pub fn round_robin_repartition(&self) -> bool { |
406 | 0 | self.options.optimizer.enable_round_robin_repartition |
407 | 0 | } |
408 | | |
409 | | /// Set the size of [`sort_spill_reservation_bytes`] to control |
410 | | /// memory pre-reservation |
411 | | /// |
412 | | /// [`sort_spill_reservation_bytes`]: datafusion_common::config::ExecutionOptions::sort_spill_reservation_bytes |
413 | 0 | pub fn with_sort_spill_reservation_bytes( |
414 | 0 | mut self, |
415 | 0 | sort_spill_reservation_bytes: usize, |
416 | 0 | ) -> Self { |
417 | 0 | self.options.execution.sort_spill_reservation_bytes = |
418 | 0 | sort_spill_reservation_bytes; |
419 | 0 | self |
420 | 0 | } |
421 | | |
422 | | /// Set the size of [`sort_in_place_threshold_bytes`] to control |
423 | | /// how sort does things. |
424 | | /// |
425 | | /// [`sort_in_place_threshold_bytes`]: datafusion_common::config::ExecutionOptions::sort_in_place_threshold_bytes |
426 | 0 | pub fn with_sort_in_place_threshold_bytes( |
427 | 0 | mut self, |
428 | 0 | sort_in_place_threshold_bytes: usize, |
429 | 0 | ) -> Self { |
430 | 0 | self.options.execution.sort_in_place_threshold_bytes = |
431 | 0 | sort_in_place_threshold_bytes; |
432 | 0 | self |
433 | 0 | } |
434 | | |
435 | | /// Convert configuration options to name-value pairs with values |
436 | | /// converted to strings. |
437 | | /// |
438 | | /// Note that this method will eventually be deprecated and |
439 | | /// replaced by [`options`]. |
440 | | /// |
441 | | /// [`options`]: Self::options |
442 | 0 | pub fn to_props(&self) -> HashMap<String, String> { |
443 | 0 | let mut map = HashMap::new(); |
444 | | // copy configs from config_options |
445 | 0 | for entry in self.options.entries() { |
446 | 0 | map.insert(entry.key, entry.value.unwrap_or_default()); |
447 | 0 | } |
448 | | |
449 | 0 | map |
450 | 0 | } |
451 | | |
452 | | /// Add extensions. |
453 | | /// |
454 | | /// Extensions can be used to attach extra data to the session config -- e.g. tracing information or caches. |
455 | | /// Extensions are opaque and the types are unknown to DataFusion itself, which makes them extremely flexible. [^1] |
456 | | /// |
457 | | /// Extensions are stored within an [`Arc`] so they do NOT require [`Clone`]. The are immutable. If you need to |
458 | | /// modify their state over their lifetime -- e.g. for caches -- you need to establish some for of interior mutability. |
459 | | /// |
460 | | /// Extensions are indexed by their type `T`. If multiple values of the same type are provided, only the last one |
461 | | /// will be kept. |
462 | | /// |
463 | | /// You may use [`get_extension`](Self::get_extension) to retrieve extensions. |
464 | | /// |
465 | | /// # Example |
466 | | /// ``` |
467 | | /// use std::sync::Arc; |
468 | | /// use datafusion_execution::config::SessionConfig; |
469 | | /// |
470 | | /// // application-specific extension types |
471 | | /// struct Ext1(u8); |
472 | | /// struct Ext2(u8); |
473 | | /// struct Ext3(u8); |
474 | | /// |
475 | | /// let ext1a = Arc::new(Ext1(10)); |
476 | | /// let ext1b = Arc::new(Ext1(11)); |
477 | | /// let ext2 = Arc::new(Ext2(2)); |
478 | | /// |
479 | | /// let cfg = SessionConfig::default() |
480 | | /// // will only remember the last Ext1 |
481 | | /// .with_extension(Arc::clone(&ext1a)) |
482 | | /// .with_extension(Arc::clone(&ext1b)) |
483 | | /// .with_extension(Arc::clone(&ext2)); |
484 | | /// |
485 | | /// let ext1_received = cfg.get_extension::<Ext1>().unwrap(); |
486 | | /// assert!(!Arc::ptr_eq(&ext1_received, &ext1a)); |
487 | | /// assert!(Arc::ptr_eq(&ext1_received, &ext1b)); |
488 | | /// |
489 | | /// let ext2_received = cfg.get_extension::<Ext2>().unwrap(); |
490 | | /// assert!(Arc::ptr_eq(&ext2_received, &ext2)); |
491 | | /// |
492 | | /// assert!(cfg.get_extension::<Ext3>().is_none()); |
493 | | /// ``` |
494 | | /// |
495 | | /// [^1]: Compare that to [`ConfigOptions`] which only supports [`ScalarValue`] payloads. |
496 | 0 | pub fn with_extension<T>(mut self, ext: Arc<T>) -> Self |
497 | 0 | where |
498 | 0 | T: Send + Sync + 'static, |
499 | 0 | { |
500 | 0 | self.set_extension(ext); |
501 | 0 | self |
502 | 0 | } |
503 | | |
504 | | /// Set extension. Pretty much the same as [`with_extension`](Self::with_extension), but take |
505 | | /// mutable reference instead of owning it. Useful if you want to add another extension after |
506 | | /// the [`SessionConfig`] is created. |
507 | | /// |
508 | | /// # Example |
509 | | /// ``` |
510 | | /// use std::sync::Arc; |
511 | | /// use datafusion_execution::config::SessionConfig; |
512 | | /// |
513 | | /// // application-specific extension types |
514 | | /// struct Ext1(u8); |
515 | | /// struct Ext2(u8); |
516 | | /// struct Ext3(u8); |
517 | | /// |
518 | | /// let ext1a = Arc::new(Ext1(10)); |
519 | | /// let ext1b = Arc::new(Ext1(11)); |
520 | | /// let ext2 = Arc::new(Ext2(2)); |
521 | | /// |
522 | | /// let mut cfg = SessionConfig::default(); |
523 | | /// |
524 | | /// // will only remember the last Ext1 |
525 | | /// cfg.set_extension(Arc::clone(&ext1a)); |
526 | | /// cfg.set_extension(Arc::clone(&ext1b)); |
527 | | /// cfg.set_extension(Arc::clone(&ext2)); |
528 | | /// |
529 | | /// let ext1_received = cfg.get_extension::<Ext1>().unwrap(); |
530 | | /// assert!(!Arc::ptr_eq(&ext1_received, &ext1a)); |
531 | | /// assert!(Arc::ptr_eq(&ext1_received, &ext1b)); |
532 | | /// |
533 | | /// let ext2_received = cfg.get_extension::<Ext2>().unwrap(); |
534 | | /// assert!(Arc::ptr_eq(&ext2_received, &ext2)); |
535 | | /// |
536 | | /// assert!(cfg.get_extension::<Ext3>().is_none()); |
537 | | /// ``` |
538 | 0 | pub fn set_extension<T>(&mut self, ext: Arc<T>) |
539 | 0 | where |
540 | 0 | T: Send + Sync + 'static, |
541 | 0 | { |
542 | 0 | let ext = ext as Arc<dyn Any + Send + Sync + 'static>; |
543 | 0 | let id = TypeId::of::<T>(); |
544 | 0 | self.extensions.insert(id, ext); |
545 | 0 | } |
546 | | |
547 | | /// Get extension, if any for the specified type `T` exists. |
548 | | /// |
549 | | /// See [`with_extension`](Self::with_extension) on how to add attach extensions. |
550 | 0 | pub fn get_extension<T>(&self) -> Option<Arc<T>> |
551 | 0 | where |
552 | 0 | T: Send + Sync + 'static, |
553 | 0 | { |
554 | 0 | let id = TypeId::of::<T>(); |
555 | 0 | self.extensions |
556 | 0 | .get(&id) |
557 | 0 | .cloned() |
558 | 0 | .map(|ext| Arc::downcast(ext).expect("TypeId unique")) |
559 | 0 | } |
560 | | } |
561 | | |
562 | | impl From<ConfigOptions> for SessionConfig { |
563 | 0 | fn from(options: ConfigOptions) -> Self { |
564 | 0 | Self { |
565 | 0 | options, |
566 | 0 | ..Default::default() |
567 | 0 | } |
568 | 0 | } |
569 | | } |
570 | | |
571 | | /// Map that holds opaque objects indexed by their type. |
572 | | /// |
573 | | /// Data is wrapped into an [`Arc`] to enable [`Clone`] while still being [object safe]. |
574 | | /// |
575 | | /// [object safe]: https://doc.rust-lang.org/reference/items/traits.html#object-safety |
576 | | type AnyMap = |
577 | | HashMap<TypeId, Arc<dyn Any + Send + Sync + 'static>, BuildHasherDefault<IdHasher>>; |
578 | | |
579 | | /// Hasher for [`AnyMap`]. |
580 | | /// |
581 | | /// With [`TypeId`]s as keys, there's no need to hash them. They are already hashes themselves, coming from the compiler. |
582 | | /// The [`IdHasher`] just holds the [`u64`] of the [`TypeId`], and then returns it, instead of doing any bit fiddling. |
583 | | #[derive(Default)] |
584 | | struct IdHasher(u64); |
585 | | |
586 | | impl Hasher for IdHasher { |
587 | 0 | fn write(&mut self, _: &[u8]) { |
588 | 0 | unreachable!("TypeId calls write_u64"); |
589 | | } |
590 | | |
591 | | #[inline] |
592 | 0 | fn write_u64(&mut self, id: u64) { |
593 | 0 | self.0 = id; |
594 | 0 | } |
595 | | |
596 | | #[inline] |
597 | 0 | fn finish(&self) -> u64 { |
598 | 0 | self.0 |
599 | 0 | } |
600 | | } |