datafusion_federation/schema_cast/
mod.rs

1use async_stream::stream;
2use datafusion::arrow::datatypes::SchemaRef;
3use datafusion::common::tree_node::TreeNodeRecursion;
4use datafusion::common::Statistics;
5use datafusion::config::ConfigOptions;
6use datafusion::error::{DataFusionError, Result};
7use datafusion::execution::{SendableRecordBatchStream, TaskContext};
8use datafusion::physical_plan::filter_pushdown::{FilterDescription, FilterPushdownPhase};
9use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
10use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
11use datafusion::physical_plan::{
12    ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PhysicalExpr,
13    PlanProperties, StatisticsArgs,
14};
15use futures::StreamExt;
16use std::clone::Clone;
17use std::fmt;
18use std::sync::Arc;
19
20mod intervals_cast;
21mod lists_cast;
22pub mod record_convert;
23mod struct_cast;
24
25#[derive(Debug)]
26#[allow(clippy::module_name_repetitions)]
27pub struct SchemaCastScanExec {
28    input: Arc<dyn ExecutionPlan>,
29    schema: SchemaRef,
30    properties: Arc<PlanProperties>,
31    metrics_set: ExecutionPlanMetricsSet,
32}
33
34impl SchemaCastScanExec {
35    pub fn new(input: Arc<dyn ExecutionPlan>, schema: SchemaRef) -> Self {
36        let eq_properties = input.equivalence_properties().clone();
37        let emission_type = input.pipeline_behavior();
38        let boundedness = input.boundedness();
39        let properties = Arc::new(PlanProperties::new(
40            eq_properties,
41            input.output_partitioning().clone(),
42            emission_type,
43            boundedness,
44        ));
45        Self {
46            input,
47            schema,
48            properties,
49            metrics_set: ExecutionPlanMetricsSet::new(),
50        }
51    }
52}
53
54impl DisplayAs for SchemaCastScanExec {
55    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
56        write!(f, "SchemaCastScanExec")
57    }
58}
59
60impl ExecutionPlan for SchemaCastScanExec {
61    fn name(&self) -> &str {
62        "SchemaCastScanExec"
63    }
64
65    fn properties(&self) -> &Arc<PlanProperties> {
66        &self.properties
67    }
68
69    fn schema(&self) -> SchemaRef {
70        Arc::clone(&self.schema)
71    }
72
73    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
74        vec![&self.input]
75    }
76
77    /// Prevents the introduction of additional `RepartitionExec` and processing input in parallel.
78    /// This guarantees that the input is processed as a single stream, preserving the order of the data.
79    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
80        vec![false]
81    }
82
83    fn with_new_children(
84        self: Arc<Self>,
85        children: Vec<Arc<dyn ExecutionPlan>>,
86    ) -> Result<Arc<dyn ExecutionPlan>> {
87        if children.len() == 1 {
88            Ok(Arc::new(Self::new(
89                Arc::clone(&children[0]),
90                Arc::clone(&self.schema),
91            )))
92        } else {
93            Err(DataFusionError::Execution(
94                "SchemaCastScanExec expects exactly one input".to_string(),
95            ))
96        }
97    }
98
99    fn execute(
100        &self,
101        partition: usize,
102        context: Arc<TaskContext>,
103    ) -> Result<SendableRecordBatchStream> {
104        let mut stream = self.input.execute(partition, context)?;
105        let schema = Arc::clone(&self.schema);
106        let baseline_metrics = BaselineMetrics::new(&self.metrics_set, partition);
107
108        Ok(Box::pin(RecordBatchStreamAdapter::new(
109            Arc::clone(&schema),
110            {
111                stream! {
112                    while let Some(batch) = stream.next().await {
113                        let _timer = baseline_metrics.elapsed_compute().timer();
114                        let batch = record_convert::try_cast_to(batch?, Arc::clone(&schema));
115                        let batch = batch.map_err(|e| { DataFusionError::External(Box::new(e)) });
116                        if let Ok(ref b) = batch {
117                            baseline_metrics.output_rows().add(b.num_rows());
118                        }
119                        yield batch;
120                    }
121                }
122            },
123        )))
124    }
125
126    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
127        vec![ChildStats::At(partition)]
128    }
129
130    fn statistics_from_inputs(
131        &self,
132        input_stats: &[Arc<Statistics>],
133        _args: &StatisticsArgs,
134    ) -> Result<Arc<Statistics>> {
135        if input_stats.len() == 1 {
136            Ok(Arc::clone(&input_stats[0]))
137        } else {
138            Err(DataFusionError::Execution(format!(
139                "{} expects exactly one input",
140                self.name()
141            )))
142        }
143    }
144
145    fn metrics(&self) -> Option<MetricsSet> {
146        Some(self.metrics_set.clone_inner())
147    }
148
149    fn gather_filters_for_pushdown(
150        &self,
151        _phase: FilterPushdownPhase,
152        parent_filters: Vec<Arc<dyn PhysicalExpr>>,
153        _config: &ConfigOptions,
154    ) -> Result<FilterDescription> {
155        FilterDescription::from_children(parent_filters, &self.children())
156    }
157
158    fn apply_expressions(
159        &self,
160        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
161    ) -> Result<TreeNodeRecursion> {
162        // We have no expressions, so per the docs for `apply_expressions`, we should return `Continue`
163        Ok(TreeNodeRecursion::Continue)
164    }
165}
datafusion_federation/schema_cast/mod.rs

datafusion_federation/schema_cast/
mod.rs
