Skip to main content

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}