delta_funnel/query_engine/
datafusion.rs

1//! DataFusion integration.
2
3use std::sync::Arc;
4
5use datafusion::common::DataFusionError;
6use datafusion::execution::TaskContext;
7use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream};
8
9use crate::DeltaFunnelError;
10
11mod catalog;
12mod execution;
13mod planning;
14mod session;
15
16pub use catalog::registration::{
17    DeltaTableProviderConfig, RegisteredDeltaSource, RegisteredDeltaSources,
18    register_delta_sources, register_delta_sources_with_scan_execution_options,
19};
20pub use execution::{
21    DeltaProviderReadStatsSnapshot, DeltaProviderReaderBackend, DeltaProviderScanExecutionOptions,
22};
23pub use planning::partition_target::{
24    DeltaScanPartitionTargetDiagnosticInput, DeltaScanPartitionTargetDiagnosticOutput,
25    DeltaScanPartitionTargetDiagnosticSource, DeltaScanPartitionTargetLocalEnvironmentDiagnostic,
26    DeltaScanPartitionTargetLocalUnixFileDescriptorLimitStatus,
27    delta_scan_partition_target_local_environment_diagnostic,
28    derive_delta_scan_partition_target_diagnostic,
29};
30pub use session::{QueryOptions, datafusion_session_config, datafusion_session_context};
31
32impl From<DeltaFunnelError> for DataFusionError {
33    fn from(error: DeltaFunnelError) -> Self {
34        Self::External(Box::new(error))
35    }
36}
37
38/// Collects provider-owned Delta read stats snapshots from a DataFusion
39/// physical plan.
40#[must_use]
41pub fn collect_delta_provider_read_stats(
42    plan: &dyn ExecutionPlan,
43) -> Vec<DeltaProviderReadStatsSnapshot> {
44    let mut found = Vec::new();
45    collect_delta_provider_read_stats_into(plan, &mut found);
46    found
47}
48
49fn collect_delta_provider_read_stats_into(
50    plan: &dyn ExecutionPlan,
51    found: &mut Vec<DeltaProviderReadStatsSnapshot>,
52) {
53    if let Some(scan) = plan
54        .as_any()
55        .downcast_ref::<execution::DeltaScanPlanningExec>()
56    {
57        found.push(scan.read_stats_snapshot());
58    }
59    for child in plan.children() {
60        collect_delta_provider_read_stats_into(child.as_ref(), found);
61    }
62}
63
64/// Executes one selected DataFusion query output as a single merged stream.
65///
66/// DataFusion physical plans can have multiple output partitions. This helper
67/// uses DataFusion's own `execute_stream` behavior, which merges those
68/// partitions into one `RecordBatch` stream while still letting partition tasks
69/// run concurrently. DeltaFunnel's downstream MSSQL writer can then stay a
70/// single awaited consumer without forcing serial partition execution.
71pub fn datafusion_query_output_stream(
72    plan: Arc<dyn ExecutionPlan>,
73    task_context: Arc<TaskContext>,
74) -> Result<SendableRecordBatchStream, DataFusionError> {
75    datafusion::physical_plan::execute_stream(plan, task_context)
76}
77
78#[cfg(test)]
79mod test_support {
80    #![allow(missing_docs)]
81
82    use std::any::Any;
83    use std::collections::HashMap;
84    use std::fs;
85    use std::path::{Path, PathBuf};
86    use std::sync::{Arc, Mutex, MutexGuard};
87    use std::time::{SystemTime, UNIX_EPOCH};
88
89    use async_trait::async_trait;
90    use datafusion::catalog::{CatalogProvider, SchemaProvider};
91    use datafusion::common::{DataFusionError, Result as DataFusionResult};
92    use datafusion::datasource::TableProvider;
93    use datafusion::physical_plan::ExecutionPlan;
94    use datafusion::prelude::SessionContext;
95
96    use crate::query_engine::datafusion::catalog::registration::{
97        DeltaTableProviderConfig, register_delta_sources,
98    };
99    use crate::query_engine::datafusion::execution::DeltaScanPlanningExec;
100    use crate::{DeltaSourceConfig, load_delta_source, preflight_delta_protocol};
101
102    pub(crate) struct DeltaLogTable {
103        path: PathBuf,
104    }
105
106    impl Drop for DeltaLogTable {
107        fn drop(&mut self) {
108            let _ = fs::remove_dir_all(&self.path);
109        }
110    }
111
112    impl DeltaLogTable {
113        pub(crate) fn new(name: &str) -> Result<Self, Box<dyn std::error::Error>> {
114            Self::new_with_schema(
115                name,
116                DEFAULT_SCHEMA_FIELDS_JSON,
117                "[]",
118                r#""partitionValues":{}"#,
119            )
120        }
121
122        pub(crate) fn new_with_schema(
123            name: &str,
124            schema_fields_json: &str,
125            partition_columns_json: &str,
126            add_partition_values_json: &str,
127        ) -> Result<Self, Box<dyn std::error::Error>> {
128            Self::new_with_schema_and_adds(
129                name,
130                schema_fields_json,
131                partition_columns_json,
132                &[add_partition_values_json],
133            )
134        }
135
136        pub(crate) fn new_with_schema_and_adds(
137            name: &str,
138            schema_fields_json: &str,
139            partition_columns_json: &str,
140            add_partition_values_jsons: &[&str],
141        ) -> Result<Self, Box<dyn std::error::Error>> {
142            Self::new_with_schema_protocol_and_adds(
143                name,
144                PROTOCOL_JSON,
145                schema_fields_json,
146                partition_columns_json,
147                add_partition_values_jsons,
148            )
149        }
150
151        pub(crate) fn new_with_schema_protocol_and_adds(
152            name: &str,
153            protocol_json: &str,
154            schema_fields_json: &str,
155            partition_columns_json: &str,
156            add_partition_values_jsons: &[&str],
157        ) -> Result<Self, Box<dyn std::error::Error>> {
158            let add_partition_values_and_sizes = add_partition_values_jsons
159                .iter()
160                .map(|partition_values_json| (*partition_values_json, 0))
161                .collect::<Vec<_>>();
162
163            Self::new_with_schema_protocol_and_sized_adds(
164                name,
165                protocol_json,
166                schema_fields_json,
167                partition_columns_json,
168                &add_partition_values_and_sizes,
169            )
170        }
171
172        pub(crate) fn new_with_schema_and_sized_adds(
173            name: &str,
174            schema_fields_json: &str,
175            partition_columns_json: &str,
176            add_partition_values_and_sizes: &[(&str, u64)],
177        ) -> Result<Self, Box<dyn std::error::Error>> {
178            Self::new_with_schema_protocol_and_sized_adds(
179                name,
180                PROTOCOL_JSON,
181                schema_fields_json,
182                partition_columns_json,
183                add_partition_values_and_sizes,
184            )
185        }
186
187        pub(crate) fn new_with_schema_protocol_and_sized_adds(
188            name: &str,
189            protocol_json: &str,
190            schema_fields_json: &str,
191            partition_columns_json: &str,
192            add_partition_values_and_sizes: &[(&str, u64)],
193        ) -> Result<Self, Box<dyn std::error::Error>> {
194            let path = Path::new("target")
195                .join("delta-funnel-datafusion-provider-tests")
196                .join(unique_name(name)?);
197            let log_path = path.join("_delta_log");
198            fs::create_dir_all(&log_path)?;
199            fs::write(
200                log_path.join("00000000000000000000.json"),
201                format!(
202                    "{}\n{}\n",
203                    protocol_json,
204                    metadata_json(schema_fields_json, partition_columns_json)
205                ),
206            )?;
207            let add_actions = add_partition_values_and_sizes
208                .iter()
209                .enumerate()
210                .map(|(index, (partition_values_json, size))| {
211                    add_json(
212                        &format!("part-{index:05}.parquet"),
213                        partition_values_json,
214                        *size,
215                    )
216                })
217                .collect::<Vec<_>>()
218                .join("\n");
219            fs::write(
220                log_path.join("00000000000000000001.json"),
221                format!("{add_actions}\n"),
222            )?;
223
224            Ok(Self { path })
225        }
226
227        pub(crate) fn path(&self) -> &Path {
228            &self.path
229        }
230    }
231
232    const PROTOCOL_JSON: &str = r#"{"protocol":{"minReaderVersion":1,"minWriterVersion":2}}"#;
233    pub(crate) const DEFAULT_SCHEMA_FIELDS_JSON: &str = r#"[{\"name\":\"id\",\"type\":\"integer\",\"nullable\":false,\"metadata\":{}},{\"name\":\"customer_name\",\"type\":\"string\",\"nullable\":true,\"metadata\":{}}]"#;
234    pub(crate) const PARTITIONED_SCHEMA_FIELDS_JSON: &str = r#"[{\"name\":\"id\",\"type\":\"integer\",\"nullable\":false,\"metadata\":{}},{\"name\":\"region\",\"type\":\"string\",\"nullable\":true,\"metadata\":{}}]"#;
235    pub(crate) const NESTED_SCHEMA_FIELDS_JSON: &str = r#"[{\"name\":\"id\",\"type\":\"integer\",\"nullable\":false,\"metadata\":{}},{\"name\":\"profile\",\"type\":{\"type\":\"struct\",\"fields\":[{\"name\":\"age\",\"type\":\"integer\",\"nullable\":true,\"metadata\":{}},{\"name\":\"tags\",\"type\":{\"type\":\"array\",\"elementType\":\"string\",\"containsNull\":true},\"nullable\":true,\"metadata\":{}}]},\"nullable\":true,\"metadata\":{}}]"#;
236    pub(crate) const DEEP_NESTED_WITH_CITY_SCHEMA_FIELDS_JSON: &str = r#"[{\"name\":\"profile\",\"type\":{\"type\":\"struct\",\"fields\":[{\"name\":\"address\",\"type\":{\"type\":\"struct\",\"fields\":[{\"name\":\"city\",\"type\":\"string\",\"nullable\":true,\"metadata\":{}}]},\"nullable\":true,\"metadata\":{}}]},\"nullable\":true,\"metadata\":{}},{\"name\":\"city\",\"type\":\"string\",\"nullable\":true,\"metadata\":{}}]"#;
237    pub(crate) const INVALID_NESTED_IDS_SCHEMA_FIELDS_JSON: &str = r#"[{\"name\":\"bad_array\",\"type\":{\"type\":\"array\",\"elementType\":\"string\",\"containsNull\":true},\"nullable\":true,\"metadata\":{\"delta.columnMapping.nested.ids\":\"not an object\"}}]"#;
238
239    fn metadata_json(schema_fields_json: &str, partition_columns_json: &str) -> String {
240        format!(
241            r#"{{"metaData":{{"id":"delta-funnel-test","format":{{"provider":"parquet","options":{{}}}},"schemaString":"{{\"type\":\"struct\",\"fields\":{schema_fields_json}}}","partitionColumns":{partition_columns_json},"configuration":{{}},"createdTime":1587968585495}}}}"#
242        )
243    }
244
245    fn add_json(path: &str, partition_values_json: &str, size: u64) -> String {
246        format!(
247            r#"{{"add":{{"path":"{path}",{partition_values_json},"size":{size},"modificationTime":1587968586000,"dataChange":true}}}}"#
248        )
249    }
250
251    fn unique_name(name: &str) -> Result<String, Box<dyn std::error::Error>> {
252        let nanos = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos();
253
254        Ok(format!("{}-{}-{nanos}", std::process::id(), name))
255    }
256
257    pub(crate) fn register_fixture_source(
258        ctx: &SessionContext,
259        source_name: &str,
260        fixture_name: &str,
261    ) -> Result<DeltaLogTable, Box<dyn std::error::Error>> {
262        let table = DeltaLogTable::new(fixture_name)?;
263        let source = load_delta_source(DeltaSourceConfig {
264            name: source_name.to_owned(),
265            table_uri: table.path.to_string_lossy().to_string(),
266            version: None,
267            storage_options: Default::default(),
268        })?;
269        let preflight = preflight_delta_protocol(&source)?;
270
271        register_delta_sources(
272            ctx,
273            vec![DeltaTableProviderConfig {
274                source,
275                protocol: preflight,
276                scan_target_partitions: None,
277            }],
278        )?;
279
280        Ok(table)
281    }
282
283    pub(crate) fn find_delta_scan_plans<'a>(
284        plan: &'a dyn ExecutionPlan,
285        found: &mut Vec<&'a DeltaScanPlanningExec>,
286    ) {
287        if let Some(scan) = plan.as_any().downcast_ref::<DeltaScanPlanningExec>() {
288            found.push(scan);
289        }
290        for child in plan.children() {
291            find_delta_scan_plans(child.as_ref(), found);
292        }
293    }
294
295    #[derive(Debug, Default)]
296    pub(crate) struct FailsOnCustomersSchemaProvider {
297        tables: Mutex<HashMap<String, Arc<dyn TableProvider>>>,
298    }
299
300    impl FailsOnCustomersSchemaProvider {
301        fn tables(&self) -> MutexGuard<'_, HashMap<String, Arc<dyn TableProvider>>> {
302            self.tables
303                .lock()
304                .unwrap_or_else(|poisoned| poisoned.into_inner())
305        }
306    }
307
308    #[async_trait]
309    impl SchemaProvider for FailsOnCustomersSchemaProvider {
310        fn as_any(&self) -> &dyn Any {
311            self
312        }
313
314        fn table_names(&self) -> Vec<String> {
315            self.tables().keys().cloned().collect()
316        }
317
318        async fn table(
319            &self,
320            name: &str,
321        ) -> DataFusionResult<Option<Arc<dyn TableProvider>>, DataFusionError> {
322            Ok(self.tables().get(name).cloned())
323        }
324
325        fn register_table(
326            &self,
327            name: String,
328            table: Arc<dyn TableProvider>,
329        ) -> DataFusionResult<Option<Arc<dyn TableProvider>>> {
330            if name == "customers" {
331                return Err(DataFusionError::Execution(
332                    "forced customers registration failure".to_owned(),
333                ));
334            }
335
336            Ok(self.tables().insert(name, table))
337        }
338
339        fn deregister_table(&self, name: &str) -> DataFusionResult<Option<Arc<dyn TableProvider>>> {
340            Ok(self.tables().remove(name))
341        }
342
343        fn table_exist(&self, name: &str) -> bool {
344            self.tables().contains_key(name)
345        }
346    }
347
348    #[derive(Debug)]
349    pub(crate) struct SingleSchemaCatalogProvider {
350        schema: Arc<dyn SchemaProvider>,
351    }
352
353    impl SingleSchemaCatalogProvider {
354        pub(crate) fn new(schema: Arc<dyn SchemaProvider>) -> Self {
355            Self { schema }
356        }
357    }
358
359    impl CatalogProvider for SingleSchemaCatalogProvider {
360        fn as_any(&self) -> &dyn Any {
361            self
362        }
363
364        fn schema_names(&self) -> Vec<String> {
365            vec!["public".to_owned()]
366        }
367
368        fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {
369            (name == "public").then(|| Arc::clone(&self.schema))
370        }
371    }
372}
373
374#[cfg(test)]
375mod tests {
376    use std::{error::Error, sync::Arc};
377
378    use datafusion::{
379        arrow::{
380            array::Int32Array,
381            datatypes::{DataType, Field, Schema, SchemaRef},
382            record_batch::RecordBatch,
383        },
384        execution::TaskContext,
385        physical_plan::{ExecutionPlan, test::TestMemoryExec},
386    };
387    use futures_util::StreamExt;
388
389    use super::datafusion_query_output_stream;
390
391    #[tokio::test]
392    async fn query_output_stream_merges_multi_partition_plan() -> Result<(), Box<dyn Error>> {
393        let schema = schema();
394        let plan = TestMemoryExec::try_new_exec(
395            &[
396                vec![int_batch(Arc::clone(&schema), &[1, 2])?],
397                vec![int_batch(Arc::clone(&schema), &[3, 4])?],
398            ],
399            Arc::clone(&schema),
400            None,
401        )?;
402        assert_eq!(plan.properties().output_partitioning().partition_count(), 2);
403
404        let plan: Arc<dyn ExecutionPlan> = plan;
405        let mut stream = datafusion_query_output_stream(plan, Arc::new(TaskContext::default()))?;
406        let mut values = Vec::new();
407
408        while let Some(batch) = stream.next().await {
409            values.extend(batch_values(&batch?)?);
410        }
411        values.sort_unstable();
412
413        assert_eq!(values, vec![1, 2, 3, 4]);
414        Ok(())
415    }
416
417    fn int_batch(schema: SchemaRef, values: &[i32]) -> Result<RecordBatch, Box<dyn Error>> {
418        RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values.to_vec()))])
419            .map_err(Into::into)
420    }
421
422    fn batch_values(batch: &RecordBatch) -> Result<Vec<i32>, Box<dyn Error>> {
423        let values = batch
424            .column(0)
425            .as_any()
426            .downcast_ref::<Int32Array>()
427            .ok_or("expected Int32Array")?;
428
429        Ok((0..values.len()).map(|index| values.value(index)).collect())
430    }
431
432    fn schema() -> SchemaRef {
433        Arc::new(Schema::new(vec![Field::new(
434            "value",
435            DataType::Int32,
436            false,
437        )]))
438    }
439}
delta_funnel/query_engine/datafusion.rs

delta_funnel/query_engine/
datafusion.rs
