datafusion_datasource/source.rs
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//! [`DataSource`] and [`DataSourceExec`]
19
20use std::any::Any;
21use std::fmt;
22use std::fmt::{Debug, Formatter};
23use std::sync::{Arc, OnceLock};
24
25use datafusion_physical_expr::projection::ProjectionExprs;
26use datafusion_physical_plan::execution_plan::{
27 Boundedness, EmissionType, SchedulingType,
28};
29use datafusion_physical_plan::metrics::SplitMetrics;
30use datafusion_physical_plan::metrics::{
31 BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet,
32};
33use datafusion_physical_plan::projection::ProjectionExec;
34use datafusion_physical_plan::stream::BatchSplitStream;
35use datafusion_physical_plan::{
36 ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
37 ReplaceChildrenOptions,
38};
39use itertools::Itertools;
40
41use crate::file::FileSource;
42use crate::file_scan_config::FileScanConfig;
43use datafusion_common::config::ConfigOptions;
44use datafusion_common::tree_node::TreeNodeRecursion;
45use datafusion_common::{Constraints, Result, Statistics};
46use datafusion_execution::{SendableRecordBatchStream, TaskContext};
47use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
48use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
49use datafusion_physical_plan::SortOrderPushdownResult;
50use datafusion_physical_plan::StatisticsArgs;
51use datafusion_physical_plan::filter_pushdown::{
52 ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, PushedDown,
53};
54
55/// A source of data, typically a list of files or memory
56///
57/// This trait provides common behaviors for abstract sources of data. It has
58/// two common implementations:
59///
60/// 1. [`FileScanConfig`]: lists of files
61/// 2. [`MemorySourceConfig`]: in memory list of `RecordBatch`
62///
63/// File format specific behaviors are defined by [`FileSource`]
64///
65/// # See Also
66/// * [`FileSource`] for file format specific implementations (Parquet, Json, etc)
67/// * [`DataSourceExec`]: The [`ExecutionPlan`] that reads from a `DataSource`
68///
69/// # Notes
70///
71/// Requires `Debug` to assist debugging
72///
73/// [`FileScanConfig`]: https://docs.rs/datafusion/latest/datafusion/datasource/physical_plan/struct.FileScanConfig.html
74/// [`MemorySourceConfig`]: https://docs.rs/datafusion/latest/datafusion/datasource/memory/struct.MemorySourceConfig.html
75/// [`FileSource`]: crate::file::FileSource
76/// [`FileFormat``]: https://docs.rs/datafusion/latest/datafusion/datasource/file_format/index.html
77/// [`TableProvider`]: https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProvider.html
78///
79/// The following diagram shows how DataSource, FileSource, and DataSourceExec are related
80/// ```text
81/// ┌─────────────────────┐ -----► execute path
82/// │ │ ┄┄┄┄┄► init path
83/// │ DataSourceExec │
84/// │ │
85/// └───────▲─────────────┘
86/// ┊ │
87/// ┊ │
88/// ┌──────────▼──────────┐ ┌──────────-──────────┐
89/// │ │ | |
90/// │ DataSource(trait) │ | TableProvider(trait)|
91/// │ │ | |
92/// └───────▲─────────────┘ └─────────────────────┘
93/// ┊ │ ┊
94/// ┌───────────────┿──┴────────────────┐ ┊
95/// | ┌┄┄┄┄┄┄┄┄┄┄┄┘ | ┊
96/// | ┊ | ┊
97/// ┌──────────▼──────────┐ ┌──────────▼──────────┐ ┊
98/// │ │ │ │ ┌──────────▼──────────┐
99/// │ FileScanConfig │ │ MemorySourceConfig │ | |
100/// │ │ │ │ | FileFormat(trait) |
101/// └──────────────▲──────┘ └─────────────────────┘ | |
102/// │ ┊ └─────────────────────┘
103/// │ ┊ ┊
104/// │ ┊ ┊
105/// ┌──────────▼──────────┐ ┌──────────▼──────────┐
106/// │ │ │ ArrowSource │
107/// │ FileSource(trait) ◄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄│ ... │
108/// │ │ │ ParquetSource │
109/// └─────────────────────┘ └─────────────────────┘
110/// │
111/// │
112/// │
113/// │
114/// ┌──────────▼──────────┐
115/// │ ArrowSource │
116/// │ ... │
117/// │ ParquetSource │
118/// └─────────────────────┘
119/// |
120/// FileOpener (called by FileStream)
121/// │
122/// ┌──────────▼──────────┐
123/// │ │
124/// │ RecordBatch │
125/// │ │
126/// └─────────────────────┘
127/// ```
128pub trait DataSource: Any + Send + Sync + Debug {
129 /// Open the specified output partition and return its stream of
130 /// [`RecordBatch`]es.
131 ///
132 /// This should be used by data sources that do not need any sibling
133 /// coordination. Data sources that want to use per-execution shared state
134 /// (for example, to reorder work across partitions at runtime) should
135 /// implement [`Self::open_with_args`] instead.
136 ///
137 /// [`RecordBatch`]: arrow::record_batch::RecordBatch
138 fn open(
139 &self,
140 partition: usize,
141 context: Arc<TaskContext>,
142 ) -> Result<SendableRecordBatchStream>;
143
144 /// Format this source for display in explain plans
145 fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result;
146
147 /// Return a copy of this DataSource with a new partitioning scheme.
148 ///
149 /// Returns `Ok(None)` (the default) if the partitioning cannot be changed.
150 /// Refer to [`ExecutionPlan::repartitioned`] for details on when None should be returned.
151 ///
152 /// Repartitioning should not change the output ordering, if this ordering exists.
153 /// Refer to [`MemorySourceConfig::repartition_preserving_order`](crate::memory::MemorySourceConfig)
154 /// and the FileSource's
155 /// [`FileGroupPartitioner::repartition_file_groups`](crate::file_groups::FileGroupPartitioner::repartition_file_groups)
156 /// for examples.
157 fn repartitioned(
158 &self,
159 _target_partitions: usize,
160 _repartition_file_min_size: usize,
161 _output_ordering: Option<LexOrdering>,
162 ) -> Result<Option<Arc<dyn DataSource>>> {
163 Ok(None)
164 }
165
166 fn output_partitioning(&self) -> Partitioning;
167 fn eq_properties(&self) -> EquivalenceProperties;
168 fn scheduling_type(&self) -> SchedulingType {
169 SchedulingType::NonCooperative
170 }
171
172 /// Returns statistics for a specific partition, or aggregate statistics
173 /// across all partitions if `partition` is `None`.
174 fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>>;
175
176 /// Return a copy of this DataSource with a new fetch limit
177 fn with_fetch(&self, _limit: Option<usize>) -> Option<Arc<dyn DataSource>>;
178 fn fetch(&self) -> Option<usize>;
179 fn metrics(&self) -> ExecutionPlanMetricsSet {
180 ExecutionPlanMetricsSet::new()
181 }
182 fn try_swapping_with_projection(
183 &self,
184 _projection: &ProjectionExprs,
185 ) -> Result<Option<Arc<dyn DataSource>>>;
186
187 /// Try to push down filters into this DataSource.
188 ///
189 /// These filters are in terms of the output schema of this DataSource (e.g.
190 /// [`Self::eq_properties`] and output of any projections pushed into the
191 /// source), not the original table schema.
192 ///
193 /// See [`ExecutionPlan::handle_child_pushdown_result`] for more details.
194 ///
195 /// [`ExecutionPlan::handle_child_pushdown_result`]: datafusion_physical_plan::ExecutionPlan::handle_child_pushdown_result
196 fn try_pushdown_filters(
197 &self,
198 filters: Vec<Arc<dyn PhysicalExpr>>,
199 _config: &ConfigOptions,
200 ) -> Result<FilterPushdownPropagation<Arc<dyn DataSource>>> {
201 Ok(FilterPushdownPropagation::with_parent_pushdown_result(
202 vec![PushedDown::No; filters.len()],
203 ))
204 }
205
206 /// Try to create a new DataSource that produces data in the specified sort order.
207 ///
208 /// # Arguments
209 /// * `order` - The desired output ordering
210 ///
211 /// # Returns
212 /// * `Ok(SortOrderPushdownResult::Exact { .. })` - Created a source that guarantees exact ordering
213 /// * `Ok(SortOrderPushdownResult::Inexact { .. })` - Created a source optimized for the ordering
214 /// * `Ok(SortOrderPushdownResult::Unsupported)` - Cannot optimize for this ordering
215 /// * `Err(e)` - Error occurred
216 ///
217 /// Default implementation returns `Unsupported`.
218 fn try_pushdown_sort(
219 &self,
220 _order: &[PhysicalSortExpr],
221 ) -> Result<SortOrderPushdownResult<Arc<dyn DataSource>>> {
222 Ok(SortOrderPushdownResult::Unsupported)
223 }
224
225 /// Returns a variant of this `DataSource` that is aware of order-sensitivity.
226 fn with_preserve_order(&self, _preserve_order: bool) -> Option<Arc<dyn DataSource>> {
227 None
228 }
229
230 /// Apply a closure to each expression used by this data source.
231 ///
232 /// This includes filter predicates (which may contain dynamic filters) and any
233 /// other expressions used during data scanning.
234 ///
235 /// The function `f` should be called once per expression unless the function returns
236 /// [`TreeNodeRecursion::Stop`] to stop iteration.
237 ///
238 /// See [`ExecutionPlan::apply_expressions`] for more details and implementation examples.
239 ///
240 /// [`ExecutionPlan::apply_expressions`]: datafusion_physical_plan::ExecutionPlan::apply_expressions
241 fn apply_expressions(
242 &self,
243 f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
244 ) -> Result<TreeNodeRecursion>;
245
246 /// Injects arbitrary run-time state into this DataSource, returning a new instance
247 /// that incorporates that state *if* it is relevant to the concrete DataSource implementation.
248 ///
249 /// This is a generic entry point: the `state` can be any type wrapped in
250 /// `Arc<dyn Any + Send + Sync>`. A data source that cares about the state should
251 /// down-cast it to the concrete type it expects and, if successful, return a
252 /// modified copy of itself that captures the provided value. If the state is
253 /// not applicable, the default behaviour is to return `None` so that parent
254 /// nodes can continue propagating the attempt further down the plan tree.
255 fn with_new_state(
256 &self,
257 _state: Arc<dyn Any + Send + Sync>,
258 ) -> Option<Arc<dyn DataSource>> {
259 None
260 }
261
262 /// Create per execution state to share across sibling instances of this
263 /// data source during one execution.
264 ///
265 /// `config` is the session configuration, so implementations can honor
266 /// options that disable sibling sharing (returning `None`) for consumers
267 /// that cannot poll all partitions in one process.
268 ///
269 /// Returns `None` (the default) if this data source has
270 /// no sibling-shared execution state.
271 fn create_sibling_state(
272 &self,
273 _config: &ConfigOptions,
274 ) -> Option<Arc<dyn Any + Send + Sync>> {
275 None
276 }
277
278 /// Open a partition using optional sibling-shared execution state.
279 ///
280 /// The default implementation ignores the additional state and delegates to
281 /// [`Self::open`].
282 fn open_with_args(&self, args: OpenArgs) -> Result<SendableRecordBatchStream> {
283 self.open(args.partition, args.context)
284 }
285
286 /// Serialize this data source to a full [`PhysicalPlanNode`] (a
287 /// `DataSourceExec` wrapping this source), if it knows how.
288 ///
289 /// This is the `DataSource` analog of
290 /// [`ExecutionPlan::try_to_proto`].
291 /// [`DataSourceExec::try_to_proto`](crate::source::DataSourceExec) delegates
292 /// to this hook, which for file scans forwards to
293 /// [`FileSource::try_to_proto`]
294 /// through the shared [`FileScanConfig`]
295 /// spine.
296 ///
297 /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller falls
298 /// back to the central downcast chain in `datafusion-proto`.
299 /// * `Ok(Some(node))` — fully serialized; the caller must not fall back.
300 ///
301 /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode
302 #[cfg(feature = "proto")]
303 fn try_to_proto(
304 &self,
305 _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
306 ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
307 Ok(None)
308 }
309}
310
311/// Arguments for [`DataSource::open_with_args`]
312#[derive(Debug, Clone)]
313pub struct OpenArgs {
314 /// Which partition to open
315 pub partition: usize,
316 /// The task context for execution
317 pub context: Arc<TaskContext>,
318 /// Optional sibling-shared execution state, see
319 /// [`DataSource::create_sibling_state`] for details.
320 pub sibling_state: Option<Arc<dyn Any + Send + Sync>>,
321}
322
323impl OpenArgs {
324 /// Create a new OpenArgs with required arguments
325 pub fn new(partition: usize, context: Arc<TaskContext>) -> Self {
326 Self {
327 partition,
328 context,
329 sibling_state: None,
330 }
331 }
332
333 /// Set sibling shared state
334 pub fn with_shared_state(
335 mut self,
336 sibling_state: Option<Arc<dyn Any + Send + Sync>>,
337 ) -> Self {
338 self.sibling_state = sibling_state;
339 self
340 }
341}
342
343impl dyn DataSource {
344 pub fn is<T: DataSource>(&self) -> bool {
345 (self as &dyn Any).is::<T>()
346 }
347
348 pub fn downcast_ref<T: DataSource>(&self) -> Option<&T> {
349 (self as &dyn Any).downcast_ref()
350 }
351}
352
353/// [`ExecutionPlan`] that reads one or more files
354///
355/// `DataSourceExec` implements common functionality such as applying
356/// projections, and caching plan properties.
357///
358/// The [`DataSource`] describes where to find the data for this data source
359/// (for example in files or what in memory partitions).
360///
361/// For file based [`DataSource`]s, format specific behavior is implemented in
362/// the [`FileSource`] trait.
363///
364/// [`FileSource`]: crate::file::FileSource
365#[derive(Clone, Debug)]
366pub struct DataSourceExec {
367 /// The source of the data -- for example, `FileScanConfig` or `MemorySourceConfig`
368 data_source: Arc<dyn DataSource>,
369 /// Cached plan properties such as sort order
370 cache: Arc<PlanProperties>,
371 /// Per execution state shared across partitions of this plan.
372 ///
373 /// Created by [`DataSource::create_sibling_state`]
374 /// and then passed to
375 /// [`DataSource::open_with_args`].
376 execution_state: Arc<OnceLock<Option<Arc<dyn Any + Send + Sync>>>>,
377}
378
379impl DisplayAs for DataSourceExec {
380 fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result {
381 match t {
382 DisplayFormatType::Default | DisplayFormatType::Verbose => {
383 write!(f, "DataSourceExec: ")?;
384 }
385 DisplayFormatType::TreeRender => {}
386 }
387 self.data_source.fmt_as(t, f)
388 }
389}
390
391impl ExecutionPlan for DataSourceExec {
392 fn name(&self) -> &'static str {
393 "DataSourceExec"
394 }
395
396 fn properties(&self) -> &Arc<PlanProperties> {
397 &self.cache
398 }
399
400 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
401 Vec::new()
402 }
403
404 fn replace_children(
405 self: Arc<Self>,
406 _: Vec<Arc<dyn ExecutionPlan>>,
407 _: ReplaceChildrenOptions,
408 ) -> Result<Arc<dyn ExecutionPlan>> {
409 Ok(self)
410 }
411
412 fn with_new_children(
413 self: Arc<Self>,
414 children: Vec<Arc<dyn ExecutionPlan>>,
415 ) -> Result<Arc<dyn ExecutionPlan>> {
416 self.replace_children(
417 children,
418 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
419 )
420 }
421
422 fn apply_expressions(
423 &self,
424 f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
425 ) -> Result<TreeNodeRecursion> {
426 // Delegate to the underlying data source
427 self.data_source.apply_expressions(f)
428 }
429
430 /// Implementation of [`ExecutionPlan::repartitioned`] which relies upon the inner [`DataSource::repartitioned`].
431 ///
432 /// If the data source does not support changing its partitioning, returns `Ok(None)` (the default). Refer
433 /// to [`ExecutionPlan::repartitioned`] for more details.
434 fn repartitioned(
435 &self,
436 target_partitions: usize,
437 config: &ConfigOptions,
438 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
439 let data_source = self.data_source.repartitioned(
440 target_partitions,
441 config.optimizer.repartition_file_min_size,
442 self.properties().eq_properties.output_ordering(),
443 )?;
444
445 Ok(data_source.map(|source| {
446 let output_partitioning = source.output_partitioning();
447 let plan = self
448 .clone()
449 .with_data_source(source)
450 // Changing source partitioning may invalidate output partitioning. Update it also
451 .with_partitioning(output_partitioning);
452 Arc::new(plan) as _
453 }))
454 }
455
456 fn execute(
457 &self,
458 partition: usize,
459 context: Arc<TaskContext>,
460 ) -> Result<SendableRecordBatchStream> {
461 let shared_state = self
462 .execution_state
463 .get_or_init(|| {
464 self.data_source
465 .create_sibling_state(context.session_config().options())
466 })
467 .clone();
468 let args = OpenArgs::new(partition, Arc::clone(&context))
469 .with_shared_state(shared_state);
470 let stream = self.data_source.open_with_args(args)?;
471 let batch_size = context.session_config().batch_size();
472
473 log::debug!(
474 "Batch splitting enabled for partition {partition}: batch_size={batch_size}"
475 );
476 let metrics = self.data_source.metrics();
477 let split_metrics = SplitMetrics::new(&metrics, partition);
478 Ok(Box::pin(BatchSplitStream::new(
479 stream,
480 batch_size,
481 split_metrics,
482 )))
483 }
484
485 fn metrics(&self) -> Option<MetricsSet> {
486 let mut metrics = self.data_source.metrics().clone_inner();
487
488 // Add `output_rows_skew` metric to the metrics set.
489 // Done here because it's a derived metric from output_rows metric.
490 if let Some(file_scan_config) = self.data_source.downcast_ref::<FileScanConfig>()
491 && file_scan_config.file_source().file_type() == "parquet"
492 && let Some(output_rows_skew) =
493 BaselineMetrics::output_rows_skew_metric(&metrics)
494 {
495 metrics.push(output_rows_skew);
496 }
497
498 Some(metrics)
499 }
500
501 fn statistics_from_inputs(
502 &self,
503 _input_stats: &[Arc<Statistics>],
504 args: &StatisticsArgs,
505 ) -> Result<Arc<Statistics>> {
506 self.data_source.partition_statistics(args.partition())
507 }
508
509 fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
510 let data_source = self.data_source.with_fetch(limit)?;
511 let cache = Arc::clone(&self.cache);
512 let execution_state = Arc::new(OnceLock::new());
513
514 Some(Arc::new(Self {
515 data_source,
516 cache,
517 execution_state,
518 }))
519 }
520
521 fn fetch(&self) -> Option<usize> {
522 self.data_source.fetch()
523 }
524
525 fn try_swapping_with_projection(
526 &self,
527 projection: &ProjectionExec,
528 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
529 match self
530 .data_source
531 .try_swapping_with_projection(projection.projection_expr())?
532 {
533 Some(new_data_source) => {
534 Ok(Some(Arc::new(DataSourceExec::new(new_data_source))))
535 }
536 None => Ok(None),
537 }
538 }
539
540 fn handle_child_pushdown_result(
541 &self,
542 _phase: FilterPushdownPhase,
543 child_pushdown_result: ChildPushdownResult,
544 config: &ConfigOptions,
545 ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
546 // Push any remaining filters into our data source
547 let parent_filters = child_pushdown_result
548 .parent_filters
549 .into_iter()
550 .map(|f| f.filter)
551 .collect_vec();
552 let res = self
553 .data_source
554 .try_pushdown_filters(parent_filters, config)?;
555 match res.updated_node {
556 Some(data_source) => {
557 let mut new_node = self.clone();
558 new_node.data_source = data_source;
559 // Re-compute properties since we have new filters which will impact equivalence info
560 new_node.cache =
561 Arc::new(Self::compute_properties(&new_node.data_source));
562
563 Ok(FilterPushdownPropagation {
564 filters: res.filters,
565 updated_node: Some(Arc::new(new_node)),
566 })
567 }
568 None => Ok(FilterPushdownPropagation {
569 filters: res.filters,
570 updated_node: None,
571 }),
572 }
573 }
574
575 fn try_pushdown_sort(
576 &self,
577 order: &[PhysicalSortExpr],
578 ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> {
579 // Delegate to the data source and wrap result with DataSourceExec
580 self.data_source
581 .try_pushdown_sort(order)?
582 .try_map(|new_data_source| {
583 let new_exec = self.clone().with_data_source(new_data_source);
584 Ok(Arc::new(new_exec) as Arc<dyn ExecutionPlan>)
585 })
586 }
587
588 fn with_preserve_order(
589 &self,
590 preserve_order: bool,
591 ) -> Option<Arc<dyn ExecutionPlan>> {
592 self.data_source
593 .with_preserve_order(preserve_order)
594 .map(|new_data_source| {
595 Arc::new(self.clone().with_data_source(new_data_source))
596 as Arc<dyn ExecutionPlan>
597 })
598 }
599
600 fn with_new_state(
601 &self,
602 state: Arc<dyn Any + Send + Sync>,
603 ) -> Option<Arc<dyn ExecutionPlan>> {
604 self.data_source
605 .with_new_state(state)
606 .map(|new_data_source| {
607 Arc::new(self.clone().with_data_source(new_data_source))
608 as Arc<dyn ExecutionPlan>
609 })
610 }
611
612 fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> {
613 let mut new_exec = Arc::unwrap_or_clone(self);
614 new_exec.execution_state = Arc::new(OnceLock::new());
615 Ok(Arc::new(new_exec))
616 }
617
618 /// Delegates serialization to the wrapped [`DataSource`]. For file scans the
619 /// concrete [`FileSource`] emits the node via its
620 /// own `try_to_proto` hook, keeping the format-specific wire logic in the
621 /// format crate.
622 #[cfg(feature = "proto")]
623 fn try_to_proto(
624 &self,
625 ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
626 ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
627 self.data_source().try_to_proto(ctx)
628 }
629}
630
631impl DataSourceExec {
632 pub fn from_data_source(data_source: impl DataSource + 'static) -> Arc<Self> {
633 Arc::new(Self::new(Arc::new(data_source)))
634 }
635
636 // Default constructor for `DataSourceExec`, setting the `cooperative` flag to `true`.
637 pub fn new(data_source: Arc<dyn DataSource>) -> Self {
638 let cache = Self::compute_properties(&data_source);
639 Self {
640 data_source,
641 cache: Arc::new(cache),
642 execution_state: Arc::new(OnceLock::new()),
643 }
644 }
645
646 /// Return the source object
647 pub fn data_source(&self) -> &Arc<dyn DataSource> {
648 &self.data_source
649 }
650
651 pub fn with_data_source(mut self, data_source: Arc<dyn DataSource>) -> Self {
652 self.cache = Arc::new(Self::compute_properties(&data_source));
653 self.data_source = data_source;
654 self.execution_state = Arc::new(OnceLock::new());
655 self
656 }
657
658 /// Assign constraints
659 pub fn with_constraints(mut self, constraints: Constraints) -> Self {
660 Arc::make_mut(&mut self.cache).set_constraints(constraints);
661 self
662 }
663
664 /// Assign output partitioning
665 pub fn with_partitioning(mut self, partitioning: Partitioning) -> Self {
666 Arc::make_mut(&mut self.cache).partitioning = partitioning;
667 self
668 }
669
670 fn compute_properties(data_source: &Arc<dyn DataSource>) -> PlanProperties {
671 PlanProperties::new(
672 data_source.eq_properties(),
673 data_source.output_partitioning(),
674 EmissionType::Incremental,
675 Boundedness::Bounded,
676 )
677 .with_scheduling_type(data_source.scheduling_type())
678 }
679
680 /// Downcast the `DataSourceExec`'s `data_source` to a specific file source
681 ///
682 /// Returns `None` if
683 /// 1. the datasource is not scanning files (`FileScanConfig`)
684 /// 2. The [`FileScanConfig::file_source`] is not of type `T`
685 pub fn downcast_to_file_source<T: FileSource>(
686 &self,
687 ) -> Option<(&FileScanConfig, &T)> {
688 self.data_source()
689 .downcast_ref::<FileScanConfig>()
690 .and_then(|file_scan_conf| {
691 file_scan_conf
692 .file_source()
693 .downcast_ref::<T>()
694 .map(|source| (file_scan_conf, source))
695 })
696 }
697}
698
699/// Create a new `DataSourceExec` from a `DataSource`
700impl<S> From<S> for DataSourceExec
701where
702 S: DataSource + 'static,
703{
704 fn from(source: S) -> Self {
705 Self::new(Arc::new(source))
706 }
707}