in work

Author: waralexrom

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs

@@ -80,7 +80,7 @@ impl PreAggregationOptimizer {
     ) -> Result<Option<Rc<Query>>, CubeError> {
         let rewriter = LogicalPlanRewriter::new();
         for multi_stage in &query.multistage_members {
-            let rewritten = rewriter.rewrite_top_down(multi_stage.clone(), &mut |&plan_node| {
+            let rewritten = rewriter.rewrite_top_down_with(multi_stage.clone(), |plan_node| {
                 Ok(NodeRewriteResult::stop())
             })?;
         }

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/visitor/rewriter.rs

@@ -34,14 +34,6 @@ pub trait LogicalNodeRewriter {
     fn process_node(&mut self, node: &PlanNode) -> Result<NodeRewriteResult, CubeError>;
 }
 
-impl<F> LogicalNodeRewriter for F
-where
-    F: FnMut(&PlanNode) -> Result<NodeRewriteResult, CubeError>,
-{
-    fn process_node(&mut self, node: &PlanNode) -> Result<NodeRewriteResult, CubeError> {
-        self(node)
-    }
-}
 
 pub struct LogicalPlanRewriter {}
 
@@ -65,6 +57,29 @@ impl LogicalPlanRewriter {
         Ok(res)
     }
 
+    pub fn rewrite_top_down_with<F, N: LogicalNode>(
+        &self,
+        node: Rc<N>,
+        f: F,
+    ) -> Result<Rc<N>, CubeError>
+    where
+        F: FnMut(&PlanNode) -> Result<NodeRewriteResult, CubeError>,
+    {
+        struct FnWrapper<F>(F);
+        
+        impl<F> LogicalNodeRewriter for FnWrapper<F>
+        where
+            F: FnMut(&PlanNode) -> Result<NodeRewriteResult, CubeError>,
+        {
+            fn process_node(&mut self, node: &PlanNode) -> Result<NodeRewriteResult, CubeError> {
+                (self.0)(node)
+            }
+        }
+        
+        let mut wrapper = FnWrapper(f);
+        self.rewrite_top_down(node, &mut wrapper)
+    }
+
     fn rewrite_top_down_impl<T: LogicalNodeRewriter>(
         &self,
         node: PlanNode,

rust/cubesqlplanner/cubesqlplanner/src/planner/base_query.rs

@@ -5,6 +5,7 @@ use crate::cube_bridge::base_query_options::BaseQueryOptions;
 use crate::cube_bridge::pre_aggregation_obj::NativePreAggregationObj;
 //use crate::logical_plan::optimizers::*;
 use crate::logical_plan::PreAggregation;
+use crate::logical_plan::PreAggregationOptimizer;
 use crate::logical_plan::Query;
 use crate::physical_plan_builder::PhysicalPlanBuilder;
 use cubenativeutils::wrappers::inner_types::InnerTypes;
@@ -147,8 +148,7 @@ impl<IT: InnerTypes> BaseQuery<IT> {
         &self,
         plan: Rc<Query>,
     ) -> Result<(Rc<Query>, Vec<Rc<PreAggregation>>), CubeError> {
-        Ok((plan.clone(), Vec::new()))
-        /* let result = if !self.request.is_pre_aggregation_query() {
+        let result = if !self.request.is_pre_aggregation_query() {
             let mut pre_aggregation_optimizer = PreAggregationOptimizer::new(
                 self.query_tools.clone(),
                 self.cubestore_support_multistage,
@@ -169,6 +169,6 @@ impl<IT: InnerTypes> BaseQuery<IT> {
         } else {
             (plan.clone(), Vec::new())
         };
-        Ok(result) */
+        Ok(result)
     }
 }