From cb744851edf7d87f44bc5dbf7866e1f4481004b8 Mon Sep 17 00:00:00 2001 From: Igal Shilman Date: Thu, 17 Oct 2024 11:22:36 +0200 Subject: [PATCH] [datafusion] prepend physical_optimizer_rule before other rules (#2114) * [datafusion] prepend physical_optimizer_rule before the other rules This is a bug fix that was introduced in https://github.com/restatedev/restate/pull/2004 Which change the order that optimizers are registered. The reason for this is described as a comment few lines above. --- crates/storage-query-datafusion/src/context.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/crates/storage-query-datafusion/src/context.rs b/crates/storage-query-datafusion/src/context.rs index 86afa1629..796e6a661 100644 --- a/crates/storage-query-datafusion/src/context.rs +++ b/crates/storage-query-datafusion/src/context.rs @@ -18,7 +18,7 @@ use datafusion::error::DataFusionError; use datafusion::execution::context::SQLOptions; use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv}; use datafusion::execution::SessionStateBuilder; -use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_optimizer::optimizer::PhysicalOptimizer; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::prelude::{SessionConfig, SessionContext}; @@ -225,13 +225,14 @@ impl QueryContext { // A far more involved but potentially more robust solution would be wrap the SymmetricHashJoin in a ProjectionExec // If this would become an issue for any reason, then we can explore that alternative. // - let physical_optimizers: Vec> = - vec![Arc::new(physical_optimizer::JoinRewrite::new())]; + let join_rewrite = Arc::new(physical_optimizer::JoinRewrite::new()); + let mut default_physical_optimizer_rules = PhysicalOptimizer::default().rules; + default_physical_optimizer_rules.insert(0, join_rewrite); - state_builder = state_builder.with_physical_optimizer_rules(physical_optimizers); + state_builder = + state_builder.with_physical_optimizer_rules(default_physical_optimizer_rules); let state = state_builder.build(); - let ctx = SessionContext::new_with_state(state); let sql_options = SQLOptions::new()