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

datafusion_datasource/
source.rs
