delta_funnel/orchestrator/session/
source_report.rs

1use std::collections::BTreeMap;
2
3use crate::{
4    DeltaFunnelError, ReportReasonCode, collect_delta_provider_read_stats,
5    report::delta::{DeltaProviderSchedulingReport, DeltaSourceReport, SourceUsageStatus},
6};
7
8use super::{
9    DeltaFunnelSession, LazyTable, MssqlDryRunOutputReport, PlannedMssqlOutput,
10    RegisteredSessionSource, errors::datafusion_handoff_setup_error,
11};
12
13impl DeltaFunnelSession {
14    /// Returns registered Delta source reports in registration order.
15    #[must_use]
16    pub fn sources(&self) -> &[RegisteredSessionSource] {
17        &self.sources
18    }
19
20    /// Returns metadata-only Delta source readiness reports in registration order.
21    #[must_use]
22    pub fn source_reports(&self) -> Vec<DeltaSourceReport> {
23        let scheduling = self.delta_source_scheduling_report();
24        self.sources
25            .iter()
26            .map(|source| delta_source_report_metadata_only(source, scheduling))
27            .collect()
28    }
29
30    /// Returns Delta source reports enriched from the physical plan for `table`.
31    ///
32    /// This method resolves the lazy table and builds a DataFusion physical plan
33    /// so provider planning metadata can be reported. It does not execute row
34    /// streams or contact SQL Server.
35    ///
36    /// # Errors
37    ///
38    /// Returns an error when the table is unknown or DataFusion cannot build the
39    /// physical plan.
40    pub async fn source_reports_for_lazy_table_plan(
41        &self,
42        table: &LazyTable,
43    ) -> Result<Vec<DeltaSourceReport>, DeltaFunnelError> {
44        let dataframe = self.dataframe_for_lazy_table(table).await?;
45        let physical_plan = dataframe
46            .create_physical_plan()
47            .await
48            .map_err(|error| datafusion_handoff_setup_error("physical_plan", error))?;
49        let provider_stats = collect_delta_provider_read_stats(physical_plan.as_ref());
50
51        Ok(self.source_reports_with_provider_read_stats(provider_stats))
52    }
53
54    fn source_reports_with_provider_read_stats(
55        &self,
56        provider_stats: Vec<crate::DeltaProviderReadStatsSnapshot>,
57    ) -> Vec<DeltaSourceReport> {
58        let scheduling = self.delta_source_scheduling_report();
59        let mut provider_stats_by_source = BTreeMap::new();
60        for stats in provider_stats {
61            provider_stats_by_source.insert(stats.source_name.clone(), stats);
62        }
63
64        self.sources
65            .iter()
66            .map(|source| {
67                let report = delta_source_report_metadata_only(source, scheduling);
68                if let Some(stats) = provider_stats_by_source.remove(source.name()) {
69                    report.with_provider_read_stats(stats)
70                } else {
71                    report
72                }
73            })
74            .collect()
75    }
76
77    pub(super) fn source_reports_for_dry_run_outputs(
78        &self,
79        outputs: &[MssqlDryRunOutputReport],
80    ) -> Result<Vec<DeltaSourceReport>, DeltaFunnelError> {
81        self.source_reports_for_output_tables(outputs.iter().map(|output| {
82            (
83                output.output_name().to_owned(),
84                output.planned_output().table(),
85            )
86        }))
87    }
88
89    pub(super) fn source_reports_for_dry_run_outputs_with_provider_stats(
90        &self,
91        outputs: &[MssqlDryRunOutputReport],
92        provider_stats: Vec<crate::DeltaProviderReadStatsSnapshot>,
93    ) -> Result<Vec<DeltaSourceReport>, DeltaFunnelError> {
94        let mut provider_stats_by_source = BTreeMap::new();
95        for stats in provider_stats {
96            provider_stats_by_source.insert(stats.source_name.clone(), stats);
97        }
98
99        Ok(self
100            .source_reports_for_dry_run_outputs(outputs)?
101            .into_iter()
102            .map(|report| {
103                if let Some(stats) = provider_stats_by_source.remove(report.source_name()) {
104                    report.with_provider_read_stats(stats)
105                } else {
106                    report.with_provider_stats_reason(ReportReasonCode::CapabilityUnavailable)
107                }
108            })
109            .collect())
110    }
111
112    pub(super) async fn provider_read_stats_for_dry_run_outputs(
113        &self,
114        outputs: &[MssqlDryRunOutputReport],
115    ) -> Result<Vec<crate::DeltaProviderReadStatsSnapshot>, DeltaFunnelError> {
116        let mut provider_stats = Vec::new();
117        for output in outputs {
118            provider_stats.extend(
119                self.provider_read_stats_for_lazy_table(output.planned_output().table())
120                    .await?,
121            );
122        }
123        Ok(provider_stats)
124    }
125
126    async fn provider_read_stats_for_lazy_table(
127        &self,
128        table: &LazyTable,
129    ) -> Result<Vec<crate::DeltaProviderReadStatsSnapshot>, DeltaFunnelError> {
130        let dataframe = self.dataframe_for_lazy_table(table).await?;
131        let physical_plan = dataframe
132            .create_physical_plan()
133            .await
134            .map_err(|error| datafusion_handoff_setup_error("physical_plan", error))?;
135
136        Ok(collect_delta_provider_read_stats(physical_plan.as_ref()))
137    }
138
139    pub(super) fn source_reports_for_planned_outputs_with_provider_stats(
140        &self,
141        outputs: &[PlannedMssqlOutput],
142        provider_stats: Vec<crate::DeltaProviderReadStatsSnapshot>,
143    ) -> Result<Vec<DeltaSourceReport>, DeltaFunnelError> {
144        let mut provider_stats_by_source = BTreeMap::new();
145        for stats in provider_stats {
146            provider_stats_by_source.insert(stats.source_name.clone(), stats);
147        }
148
149        Ok(self
150            .source_reports_for_output_tables(outputs.iter().map(|output| {
151                (
152                    output.output_plan().output_name().to_owned(),
153                    output.table(),
154                )
155            }))?
156            .into_iter()
157            .map(|report| {
158                if let Some(stats) = provider_stats_by_source.remove(report.source_name()) {
159                    report.with_provider_read_stats(stats)
160                } else {
161                    report.with_provider_stats_reason(ReportReasonCode::CapabilityUnavailable)
162                }
163            })
164            .collect())
165    }
166
167    fn source_reports_for_output_tables<'a>(
168        &self,
169        outputs: impl IntoIterator<Item = (String, &'a LazyTable)>,
170    ) -> Result<Vec<DeltaSourceReport>, DeltaFunnelError> {
171        let outputs = outputs.into_iter().collect::<Vec<_>>();
172        let mut output_sources = Vec::with_capacity(outputs.len());
173        let mut all_usage_known = true;
174
175        for (output_name, table) in outputs {
176            match self.known_source_dependencies_for_table(table)? {
177                Some(source_ids) => {
178                    output_sources.push((output_name, source_ids));
179                }
180                None => {
181                    all_usage_known = false;
182                }
183            }
184        }
185
186        Ok(self
187            .sources
188            .iter()
189            .map(|source| {
190                let used_by_output_names = output_sources
191                    .iter()
192                    .filter(|(_, source_ids)| source_ids.contains(&source.table().id()))
193                    .map(|(output_name, _)| output_name.clone())
194                    .collect::<Vec<_>>();
195                let usage_status = if used_by_output_names.is_empty() {
196                    if all_usage_known {
197                        SourceUsageStatus::NotUsed
198                    } else {
199                        SourceUsageStatus::Unknown
200                    }
201                } else {
202                    SourceUsageStatus::Used
203                };
204
205                delta_source_report_metadata_only(source, self.delta_source_scheduling_report())
206                    .with_usage(usage_status, used_by_output_names)
207            })
208            .collect())
209    }
210
211    fn delta_source_scheduling_report(&self) -> DeltaProviderSchedulingReport {
212        DeltaProviderSchedulingReport::from_options(
213            self.options.query_options(),
214            self.options.provider_scan_options(),
215        )
216    }
217}
218
219fn delta_source_report_metadata_only(
220    source: &RegisteredSessionSource,
221    scheduling: DeltaProviderSchedulingReport,
222) -> DeltaSourceReport {
223    DeltaSourceReport::metadata_only(
224        source.name(),
225        source.source_uri(),
226        source.snapshot_version(),
227        source.protocol().clone(),
228        scheduling,
229    )
230    .with_phase_timings(source.phase_timings().to_vec())
231}
232
233#[cfg(test)]
234mod tests {
235    use super::super::{SessionOptions, test_support::DeltaLogTable};
236    use super::DeltaFunnelSession;
237    use crate::{DeltaSourceConfig, FileCount, ReportReasonCode};
238
239    #[tokio::test]
240    async fn source_reports_for_lazy_table_plan_include_provider_stats_without_execution()
241    -> Result<(), Box<dyn std::error::Error>> {
242        let table = DeltaLogTable::new("orders")?;
243        let mut session = DeltaFunnelSession::new(SessionOptions::default())?;
244        let source = session.delta_lake(DeltaSourceConfig::new("orders", table.uri()))?;
245
246        let reports = session.source_reports_for_lazy_table_plan(&source).await?;
247
248        assert_eq!(reports.len(), 1);
249        let report = &reports[0];
250        assert_eq!(report.source_name(), "orders");
251        assert_eq!(report.provider_stats_reason(), None);
252        let stats = report
253            .provider_read_stats()
254            .ok_or("expected provider read stats")?;
255        assert_eq!(stats.source_name, "orders");
256        assert_eq!(stats.snapshot_version, report.snapshot_version());
257        assert_eq!(stats.files_started, 0);
258        assert_eq!(stats.files_completed, 0);
259        assert_eq!(stats.batches_produced, 0);
260        assert_eq!(stats.rows_produced, 0);
261        match stats.scan_metadata_exhausted {
262            Some(true) => {
263                assert_eq!(report.file_count(), FileCount::exact(stats.files_planned));
264                assert_eq!(report.file_count_reason(), None);
265                assert!(report.scan_metadata_exhausted());
266            }
267            Some(false) => {
268                assert_eq!(
269                    report.file_count(),
270                    FileCount::estimated(stats.files_planned)
271                );
272                assert_eq!(report.file_count_reason(), None);
273                assert!(!report.scan_metadata_exhausted());
274            }
275            None => {
276                assert_eq!(report.file_count(), FileCount::unavailable());
277                assert_eq!(
278                    report.file_count_reason(),
279                    Some(ReportReasonCode::CapabilityUnavailable)
280                );
281                assert!(!report.scan_metadata_exhausted());
282            }
283        }
284        Ok(())
285    }
286}
delta_funnel/orchestrator/session/source_report.rs

delta_funnel/orchestrator/session/
source_report.rs
