datafusion_federation/schema_cast/
mod.rs

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

datafusion_federation/schema_cast/
mod.rs
