chore(cubestore): Upgrade DF 46: Pass customizer more completely and avoid ParquetExec

Author: srh

rust/cubestore/Cargo.lock

@@ -264,6 +264,7 @@ impl QueryPlannerImpl {
     }
 
     pub fn make_execution_context() -> SessionContext {
+        // TODO upgrade DF: Remove this -- use metadata_cache_factory.make_session_config()
         Self::execution_context_helper(SessionConfig::new())
     }
 

rust/cubestore/cubestore/src/queryplanner/mod.rs

@@ -14,7 +14,7 @@ pub fn add_check_memory_exec(
     mem_handler: Arc<dyn MemoryHandler>,
 ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
     let p_any = p.as_any();
-    // TODO upgrade DF: Do we use ParquetExec?  Or just DataSourceExec?  It's fine to have both here.
+    // We supposedly don't use ParquetExec, which is deprecated in DF 46, anymore but we keep the check here in case we do.
     if p_any.is::<DataSourceExec>() || p_any.is::<ParquetExec>() || p_any.is::<MemoryExec>() || p_any.is::<ClusterSendExec>() {
         let memory_check = Arc::new(CheckMemoryExec::new(p, mem_handler.clone()));
         Ok(memory_check)

rust/cubestore/cubestore/src/queryplanner/optimizations/check_memory.rs

@@ -169,7 +169,7 @@ fn finalize_physical_plan(
     let p = rewrite_physical_plan(p, &mut |p| add_check_memory_exec(p, memory_handler.clone()))?;
     let p = if let Some(data_loaded_size) = data_loaded_size {
         rewrite_physical_plan(p, &mut |p| {
-            add_trace_data_loaded_exec(p, data_loaded_size.clone())
+            add_trace_data_loaded_exec(p, &data_loaded_size)
         })?
     } else {
         p

rust/cubestore/cubestore/src/queryplanner/optimizations/mod.rs

@@ -1,19 +1,30 @@
 use crate::queryplanner::trace_data_loaded::{DataLoadedSize, TraceDataLoadedExec};
-use datafusion::datasource::physical_plan::ParquetExec;
+use datafusion::datasource::physical_plan::{ParquetExec, ParquetSource};
 use datafusion::error::DataFusionError;
 use datafusion::physical_plan::ExecutionPlan;
+use datafusion_datasource::file_scan_config::FileScanConfig;
+use datafusion_datasource::source::DataSourceExec;
 use std::sync::Arc;
 
-/// Add `TraceDataLoadedExec` behind ParquetExec nodes.
+/// Add `TraceDataLoadedExec` behind ParquetExec or DataSourceExec (with File hence Parquet source) nodes.
 pub fn add_trace_data_loaded_exec(
     p: Arc<dyn ExecutionPlan>,
-    data_loaded_size: Arc<DataLoadedSize>,
+    data_loaded_size: &Arc<DataLoadedSize>,
 ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+    fn do_wrap(p: Arc<dyn ExecutionPlan>, data_loaded_size: &Arc<DataLoadedSize>) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+        Ok(Arc::new(TraceDataLoadedExec::new(p, data_loaded_size.clone())))
+    }
+
     let p_any = p.as_any();
     if p_any.is::<ParquetExec>() {
-        let trace_data_loaded = Arc::new(TraceDataLoadedExec::new(p, data_loaded_size.clone()));
-        Ok(trace_data_loaded)
-    } else {
-        Ok(p)
+        // ParquetExec is deprecated in DF 46 and we don't use it; we shouldn't hit this case, but we keep it just in case.
+        return do_wrap(p, data_loaded_size);
+    } else if let Some(dse) = p_any.downcast_ref::<DataSourceExec>() {
+        if let Some(file_scan) = dse.data_source().as_any().downcast_ref::<FileScanConfig>() {
+            if file_scan.file_source().as_any().is::<ParquetSource>() {
+                return do_wrap(p, data_loaded_size);
+            }
+        }
     }
+    Ok(p)
 }

rust/cubestore/cubestore/src/queryplanner/optimizations/trace_data_loaded.rs

@@ -4,7 +4,7 @@ use bigdecimal::ToPrimitive;
 use datafusion::arrow::datatypes::Schema;
 use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
 use datafusion::common::DFSchema;
-use datafusion::datasource::physical_plan::ParquetExec;
+use datafusion::datasource::physical_plan::{ParquetExec, ParquetSource};
 use datafusion::datasource::{DefaultTableSource, TableProvider};
 use datafusion::error::DataFusionError;
 use datafusion::logical_expr::{
@@ -15,9 +15,11 @@ use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::filter::FilterExec;
 use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
-use datafusion::physical_plan::{ExecutionPlan, InputOrderMode, PlanProperties};
+use datafusion::physical_plan::{DefaultDisplay, DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, PlanProperties};
 use datafusion::prelude::Expr;
+use datafusion_datasource::file_scan_config::FileScanConfig;
 use datafusion_datasource::memory::MemoryExec;
+use datafusion_datasource::source::DataSourceExec;
 use itertools::{repeat_n, Itertools};
 use std::sync::Arc;
 
@@ -656,15 +658,31 @@ fn pp_phys_plan_indented(p: &dyn ExecutionPlan, indent: usize, o: &PPOptions, ou
         } else if let Some(_) = a.downcast_ref::<FilterByKeyRangeExec>() {
             *out += "FilterByKeyRange";
         } else if let Some(p) = a.downcast_ref::<ParquetExec>() {
+            // We don't use ParquetExec any more.
             *out += &format!(
-                "ParquetScan, files: {}",
+                "ParquetExec (ERROR: deprecated), files: {}",
                 p.base_config()
                     .file_groups
                     .iter()
                     .flatten()
                     .map(|p| p.object_meta.location.to_string())
                     .join(",")
             );
+        } else if let Some(dse) = a.downcast_ref::<DataSourceExec>() {
+            let data_source = dse.data_source();
+            if let Some(fse) = data_source.as_any().downcast_ref::<FileScanConfig>() {
+                if let Some(p) = fse.file_source().as_any().downcast_ref::<ParquetSource>() {
+                    *out += &format!(
+                        "ParquetScan, files: {}",
+                        fse.file_groups.iter().flatten().map(|p| p.object_meta.location.to_string()).join(","),
+                    );
+                } else {
+                    *out += &format!("{}", DefaultDisplay(dse));
+                }
+            } else {
+                *out += &format!("{}", DefaultDisplay(dse));
+            }
+
             // TODO upgrade DF
             // } else if let Some(_) = a.downcast_ref::<SkipExec>() {
             //     *out += "SkipRows";