Fix regression with TPC-H benchmark (#1060)

* Fix regression with TPC-H benchmark * add link to issue
apache · Sep 24, 2024 · 51adfb0 · 51adfb0
1 parent e39a7e6
commit 51adfb0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs
@@ -32,6 +32,8 @@ use ballista_core::error::BallistaError;
 use ballista_core::serde::scheduler::{ExecutorSpecification, PartitionId};
 use ballista_core::serde::BallistaCodec;
 use datafusion::execution::context::TaskContext;
+use datafusion::functions::datetime::date_part;
+use datafusion::functions::unicode::substr;
 use datafusion::functions_aggregate::covariance::{covar_pop_udaf, covar_samp_udaf};
 use datafusion::functions_aggregate::sum::sum_udaf;
 use datafusion::functions_aggregate::variance::var_samp_udaf;
@@ -199,6 +201,10 @@ async fn run_received_task<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
     task_aggregate_functions.insert("covar_pop".to_string(), covar_pop_udaf());
     task_aggregate_functions.insert("SUM".to_string(), sum_udaf());
 
+    // TODO which other functions need adding here?
+    task_scalar_functions.insert("date_part".to_string(), date_part());
+    task_scalar_functions.insert("substr".to_string(), substr());
+
     for window_func in executor.window_functions.clone() {
         task_window_functions.insert(window_func.0, window_func.1);
     }

diff --git a/ballista/scheduler/src/state/execution_graph/execution_stage.rs b/ballista/scheduler/src/state/execution_graph/execution_stage.rs
@@ -23,7 +23,6 @@ use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use datafusion::physical_optimizer::aggregate_statistics::AggregateStatistics;
-use datafusion::physical_optimizer::join_selection::JoinSelection;
 use datafusion::physical_optimizer::PhysicalOptimizerRule;
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::physical_plan::metrics::{MetricValue, MetricsSet};
@@ -349,10 +348,13 @@ impl UnresolvedStage {
             &input_locations,
         )?;
 
+        // TODO reinstate this logic once https://github.com/apache/datafusion/issues/10978
+        // is fixed
         // Optimize join order and statistics based on new resolved statistics
-        let optimize_join = JoinSelection::new();
-        let config = SessionConfig::default();
-        let plan = optimize_join.optimize(plan, config.options())?;
+        // let optimize_join = JoinSelection::new();
+        // let config = SessionConfig::default();
+        // let plan = optimize_join.optimize(plan, config.options())?;
+
         let optimize_aggregate = AggregateStatistics::new();
         let plan =
             optimize_aggregate.optimize(plan, SessionConfig::default().options())?;