chore: optimize binary_expr_normalize + align stack to 8mb for testing on macOS

Author: ovr

rust/cubesql/cubesql/src/compile/engine/df/optimizers/plan_normalize.rs

@@ -1218,23 +1218,15 @@ fn binary_expr_normalize(
     // and cast the `DATE` to `TIMESTAMP` to match the types.
     match (&left_type, &right_type) {
         (DataType::Timestamp(_, _), DataType::Date32) => {
-            let new_right = Box::new(evaluate_expr(
-                optimizer,
-                right.cast_to(&left_type, schema)?,
-            )?);
             return Ok(Box::new(Expr::BinaryExpr {
                 left,
                 op,
-                right: new_right,
+                right: evaluate_expr(optimizer, right.cast_to(&left_type, schema)?)?,
             }));
         }
         (DataType::Date32, DataType::Timestamp(_, _)) => {
-            let new_left = Box::new(evaluate_expr(
-                optimizer,
-                left.cast_to(&right_type, schema)?,
-            )?);
             return Ok(Box::new(Expr::BinaryExpr {
-                left: new_left,
+                left: evaluate_expr(optimizer, left.cast_to(&right_type, schema)?)?,
                 op,
                 right,
             }));
@@ -1256,21 +1248,16 @@ fn binary_expr_normalize(
     };
 
     if literal_on_the_left {
-        let new_left = Box::new(evaluate_expr(optimizer, left.cast_to(&cast_type, schema)?)?);
         Ok(Box::new(Expr::BinaryExpr {
-            left: new_left,
+            left: evaluate_expr(optimizer, left.cast_to(&cast_type, schema)?)?,
             op,
             right,
         }))
     } else {
-        let new_right = Box::new(evaluate_expr(
-            optimizer,
-            right.cast_to(&cast_type, schema)?,
-        )?);
         Ok(Box::new(Expr::BinaryExpr {
             left,
             op,
-            right: new_right,
+            right: evaluate_expr(optimizer, right.cast_to(&cast_type, schema)?)?,
         }))
     }
 }
@@ -1357,7 +1344,7 @@ fn in_list_expr_normalize(
                 return Ok(list_expr_normalized);
             }
 
-            evaluate_expr(optimizer, list_expr_normalized.cast_to(&expr_type, schema)?)
+            evaluate_expr_stacked(optimizer, list_expr_normalized.cast_to(&expr_type, schema)?)
         })
         .collect::<Result<Vec<_>>>()?;
 
@@ -1368,18 +1355,22 @@ fn in_list_expr_normalize(
     }))
 }
 
-/// Evaluates an expression to a constant if possible.
-fn evaluate_expr(optimizer: &PlanNormalize, expr: Expr) -> Result<Expr> {
+fn evaluate_expr_stacked(optimizer: &PlanNormalize, expr: Expr) -> Result<Expr> {
     let execution_props = &optimizer.cube_ctx.state.execution_props;
     let mut const_evaluator = ConstEvaluator::new(execution_props);
     expr.rewrite(&mut const_evaluator)
 }
 
+/// Evaluates an expression to a constant if possible.
+fn evaluate_expr(optimizer: &PlanNormalize, expr: Expr) -> Result<Box<Expr>> {
+    Ok(Box::new(evaluate_expr_stacked(optimizer, expr)?))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::compile::test::{
-        get_test_tenant_ctx, rewrite_engine::create_test_postgresql_cube_context,
+        get_test_tenant_ctx, rewrite_engine::create_test_postgresql_cube_context, run_async_test,
     };
     use datafusion::{
         arrow::datatypes::{DataType, Field, Schema},
@@ -1402,35 +1393,39 @@ mod tests {
         expr
     }
 
-    #[tokio::test]
-    async fn test_stack_overflow_deeply_nested_or() -> Result<()> {
-        let meta = get_test_tenant_ctx();
-        let cube_ctx = create_test_postgresql_cube_context(meta)
-            .await
-            .expect("Failed to create cube context");
-
-        // Create a simple table
-        let schema = Schema::new(vec![
-            Field::new("id", DataType::Int32, false),
-            Field::new("value", DataType::Int32, true),
-        ]);
-
-        let table_scan = LogicalPlanBuilder::scan_empty(Some("test_table"), &schema, None)
-            .expect("Failed to create table scan")
-            .build()
-            .expect("Failed to build plan");
-
-        // Create a deeply nested OR expression (should cause stack overflow)
-        let deeply_nested_filter = create_deeply_nested_or_expr("value", 1_000);
-
-        let plan = LogicalPlanBuilder::from(table_scan)
-            .filter(deeply_nested_filter)
-            .expect("Failed to add filter")
-            .build()
-            .expect("Failed to build plan");
-
-        let optimizer = PlanNormalize::new(&cube_ctx);
-        optimizer.optimize(&plan, &OptimizerConfig::new())?;
+    // plan_normalize is recursive, at the same time ExprRewriter from DF is too
+    // let's guard it with test, that our code in dev profile is optimized to rewrite N nodes
+    #[test]
+    fn test_stack_overflow_deeply_nested_or() -> Result<()> {
+        run_async_test(async move {
+            let meta = get_test_tenant_ctx();
+            let cube_ctx = create_test_postgresql_cube_context(meta)
+                .await
+                .expect("Failed to create cube context");
+
+            // Create a simple table
+            let schema = Schema::new(vec![
+                Field::new("id", DataType::Int32, false),
+                Field::new("value", DataType::Int32, true),
+            ]);
+
+            let table_scan = LogicalPlanBuilder::scan_empty(Some("test_table"), &schema, None)
+                .expect("Failed to create table scan")
+                .build()
+                .expect("Failed to build plan");
+
+            // Create a deeply nested OR expression (should cause stack overflow)
+            let deeply_nested_filter = create_deeply_nested_or_expr("value", 200);
+
+            let plan = LogicalPlanBuilder::from(table_scan)
+                .filter(deeply_nested_filter)
+                .expect("Failed to add filter")
+                .build()
+                .expect("Failed to build plan");
+
+            let optimizer = PlanNormalize::new(&cube_ctx);
+            optimizer.optimize(&plan, &OptimizerConfig::new()).unwrap();
+        });
 
         Ok(())
     }

rust/cubesql/cubesql/src/compile/test/mod.rs

@@ -1,5 +1,3 @@
-use std::{collections::HashMap, env, ops::Deref, sync::Arc};
-
 use super::{convert_sql_to_cube_query, CompilationResult, QueryPlan};
 use crate::{
     compile::{
@@ -23,6 +21,8 @@ use crate::{
 use async_trait::async_trait;
 use cubeclient::models::V1CubeMetaType;
 use datafusion::{arrow::datatypes::SchemaRef, dataframe::DataFrame as DFDataFrame};
+use std::future::Future;
+use std::{collections::HashMap, env, ops::Deref, sync::Arc};
 use uuid::Uuid;
 
 pub mod rewrite_engine;
@@ -1201,3 +1201,45 @@ pub async fn convert_select_to_query_plan_with_meta(
 
     query.unwrap()
 }
+
+pub struct AsyncTestOptions {
+    stack_bytes: usize,
+}
+
+impl Default for AsyncTestOptions {
+    fn default() -> Self {
+        Self {
+            // By default tokio stack size is aligned with OS default:
+            // 1mb for Windows, 2mb for macOS, 8mb for linux
+            // Let's align everything to 8mb
+            stack_bytes: 8 * 1024 * 1024,
+        }
+    }
+}
+
+pub fn run_async_test<F>(fut: F)
+where
+    F: Future<Output = ()> + Send + 'static,
+{
+    run_async_test_opt(AsyncTestOptions::default(), fut)
+}
+
+pub fn run_async_test_opt<F>(opts: AsyncTestOptions, fut: F)
+where
+    F: Future<Output = ()> + Send + 'static,
+{
+    std::thread::Builder::new()
+        .name("test-rt".into())
+        .stack_size(opts.stack_bytes)
+        .spawn(|| {
+            let rt = tokio::runtime::Builder::new_multi_thread()
+                .enable_all()
+                .build()
+                .unwrap();
+
+            rt.block_on(fut);
+        })
+        .unwrap()
+        .join()
+        .unwrap();
+}