diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs index 10f7a4587ba28..78db203a33dc9 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/rust/cubesql/cubesql/src/compile/mod.rs @@ -40,7 +40,6 @@ mod tests { rewrite::rewriter::Rewriter, test::{get_sixteen_char_member_cube, get_string_cube_meta}, }, - config::ConfigObjImpl, CubeError, }; use chrono::Datelike; @@ -53,13 +52,12 @@ mod tests { use pretty_assertions::assert_eq; use regex::Regex; use serde_json::json; - use std::{env, sync::Arc}; + use std::env; use crate::compile::test::{ convert_select_to_query_plan, convert_select_to_query_plan_customized, - convert_select_to_query_plan_with_config, convert_select_to_query_plan_with_meta, - execute_queries_with_flags, execute_query, init_testing_logger, LogicalPlanTestUtils, - TestContext, + convert_select_to_query_plan_with_meta, execute_queries_with_flags, execute_query, + init_testing_logger, LogicalPlanTestUtils, TestContext, }; #[tokio::test] @@ -14137,956 +14135,6 @@ ORDER BY "source"."str0" ASC .contains("LEFT")); } - #[tokio::test] - async fn test_simple_wrapper() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT COALESCE(customer_gender, 'N/A', 'NN'), AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("COALESCE")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_wrapper_group_by_rollup() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, notes, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, ROLLUP(2)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("Rollup")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_wrapper_group_by_rollup_with_aliases() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender as \"customer_gender1\", notes as \"notes\", AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY ROLLUP(1, 2)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("Rollup")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_wrapper_group_by_rollup_nested() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, notes, avg(mp) from (SELECT customer_gender, notes, avg(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, 2) b GROUP BY ROLLUP(1, 2)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("ROLLUP(1, 2)")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_wrapper_group_by_rollup_nested_from_asterisk() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, notes, avg(avgPrice) from (SELECT * FROM KibanaSampleDataEcommerce) b GROUP BY ROLLUP(1, 2) ORDER BY 1" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("Rollup")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_wrapper_group_by_rollup_nested_with_aliases() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender as \"gender\", notes as \"notes\", avg(mp) from (SELECT customer_gender, notes, avg(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, 2) b GROUP BY ROLLUP(1, 2)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("ROLLUP(1, 2)")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_wrapper_group_by_rollup_nested_complex() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, notes, order_date, last_mod, avg(mp) from \ - (SELECT customer_gender, notes, order_date, last_mod, avg(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, 2, 3, 4) b \ - GROUP BY ROLLUP(1), ROLLUP(2), 3, CUBE(4)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("ROLLUP(1), ROLLUP(2), 3, CUBE(4)")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_wrapper_group_by_rollup_placeholders() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, notes, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY ROLLUP(1, 2)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("Rollup")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_wrapper_group_by_cube() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, notes, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY CUBE(customer_gender, notes)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("Cube")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_wrapper_group_by_rollup_complex() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, notes, has_subscription, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY ROLLUP(customer_gender, notes), has_subscription" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("Rollup")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_projection_empty_source() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT (SELECT 'male' where 1 group by 'male' having 1 order by 'male' limit 1) as gender, avgPrice FROM KibanaSampleDataEcommerce a" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("(SELECT")); - assert!(sql.contains("utf8__male__")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - //println!("phys plan {:?}", physical_plan); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_filter_empty_source() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT avgPrice FROM KibanaSampleDataEcommerce a where customer_gender = (SELECT 'male' )" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("(SELECT")); - assert!(sql.contains("utf8__male__")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - //println!("phys plan {:?}", physical_plan); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_projection_aggregate_empty_source() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT (SELECT 'male'), avg(avgPrice) FROM KibanaSampleDataEcommerce a GROUP BY 1" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("(SELECT")); - assert!(sql.contains("utf8__male__")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_filter_in_empty_source() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, avgPrice FROM KibanaSampleDataEcommerce a where customer_gender in (select 'male')" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("IN (SELECT")); - assert!(sql.contains("utf8__male__")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_filter_and_projection_empty_source() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT (select 'male'), avgPrice FROM KibanaSampleDataEcommerce a where customer_gender in (select 'female')" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - - let sql = logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql; - assert!(sql.contains("IN (SELECT")); - assert!(sql.contains("(SELECT")); - assert!(sql.contains("utf8__male__")); - assert!(sql.contains("utf8__female__")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_projection() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT (SELECT customer_gender FROM KibanaSampleDataEcommerce LIMIT 1) as gender, avgPrice FROM KibanaSampleDataEcommerce a" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("(SELECT")); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("\\\\\\\"limit\\\\\\\":1")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_projection_aggregate() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT (SELECT customer_gender FROM KibanaSampleDataEcommerce WHERE customer_gender = 'male' LIMIT 1), avg(avgPrice) FROM KibanaSampleDataEcommerce a GROUP BY 1" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("(SELECT")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_filter_equal() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, avgPrice FROM KibanaSampleDataEcommerce a where customer_gender = (select customer_gender from KibanaSampleDataEcommerce limit 1)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("(SELECT")); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("\\\\\\\"limit\\\\\\\":1")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_filter_in() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT customer_gender, avgPrice FROM KibanaSampleDataEcommerce a where customer_gender in (select customer_gender from KibanaSampleDataEcommerce)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("IN (SELECT")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_simple_subquery_wrapper_filter_and_projection() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT (select customer_gender from KibanaSampleDataEcommerce limit 1), avgPrice FROM KibanaSampleDataEcommerce a where customer_gender in (select customer_gender from KibanaSampleDataEcommerce)" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("IN (SELECT")); - - let _physical_plan = query_plan.as_physical_plan().await.unwrap(); - } - - #[tokio::test] - async fn test_case_wrapper() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT CASE WHEN COALESCE(customer_gender, 'N/A', 'NN') = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("CASE WHEN")); - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_distinct() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - r#"SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, COUNT(DISTINCT countDistinct) mp - FROM KibanaSampleDataEcommerce a - WHERE - ( - ( - ( a.order_date ) >= '2024-01-01' - AND ( a.order_date ) < '2024-02-01' - ) - ) - GROUP BY 1"# - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("CASE WHEN")); - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_alias_with_order() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END AS \"f822c516-3515-11c2-8464-5d4845a02f73\", AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END ORDER BY CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END NULLS FIRST LIMIT 500" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("ORDER BY \"case_when_a_cust\"")); - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_ungrouped() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("CASE WHEN")); - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_non_strict_match() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let mut config = ConfigObjImpl::default(); - - config.disable_strict_agg_type_match = true; - - let query_plan = convert_select_to_query_plan_with_config( - "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, SUM(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1" - .to_string(), - DatabaseProtocol::PostgreSQL, - Arc::new(config) - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("CASE WHEN")); - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_ungrouped_sorted() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1 ORDER BY 1 DESC" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("ORDER BY")); - } - - #[tokio::test] - async fn test_case_wrapper_ungrouped_sorted_aliased() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT x FROM (SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END x, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1 ORDER BY 1 DESC) b" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - // TODO test without depend on column name - .contains("ORDER BY \"case_when")); - } - - #[tokio::test] - async fn test_case_wrapper_with_internal_limit() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1 LIMIT 1123" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("CASE WHEN")); - - assert!( - logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("1123"), - "SQL contains 1123: {}", - logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - ); - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_with_system_fields() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, __user, __cubeJoinField, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, 2, 3 LIMIT 1123" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - - assert!( - logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("\\\"cube_name\\\":\\\"KibanaSampleDataEcommerce\\\",\\\"alias\\\":\\\"user\\\""), - r#"SQL contains `\"cube_name\":\"KibanaSampleDataEcommerce\",\"alias\":\"user\"` {}"#, - logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - ); - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_with_limit() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT * FROM (SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1) q LIMIT 1123" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("CASE WHEN")); - - assert!( - logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("1123"), - "SQL contains 1123: {}", - logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - ); - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_with_null() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT CASE WHEN taxful_total_price IS NULL THEN NULL WHEN taxful_total_price < taxful_total_price * 2 THEN COALESCE(taxful_total_price, 0, 0) END, AVG(avgPrice) FROM KibanaSampleDataEcommerce GROUP BY 1" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - .contains("CASE WHEN")); - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_ungrouped_on_dimension() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan( - "SELECT CASE WHEN SUM(taxful_total_price) > 0 THEN SUM(taxful_total_price) ELSE 0 END FROM KibanaSampleDataEcommerce a" - .to_string(), - DatabaseProtocol::PostgreSQL, - ) - .await; - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - } - - #[tokio::test] - async fn test_case_wrapper_escaping() { - if !Rewriter::sql_push_down_enabled() { - return; - } - init_testing_logger(); - - let query_plan = convert_select_to_query_plan_customized( - "SELECT CASE WHEN customer_gender = '\\`' THEN COALESCE(customer_gender, 'N/A', 'NN') ELSE 'N/A' END as \"\\`\", AVG(avgPrice) FROM KibanaSampleDataEcommerce a GROUP BY 1".to_string(), - DatabaseProtocol::PostgreSQL, - vec![ - ("expressions/binary".to_string(), "{{ left }} \\`{{ op }} {{ right }}".to_string()) - ], - ).await; - - let physical_plan = query_plan.as_physical_plan().await.unwrap(); - println!( - "Physical plan: {}", - displayable(physical_plan.as_ref()).indent() - ); - - let logical_plan = query_plan.as_logical_plan(); - assert!(logical_plan - .find_cube_scan_wrapper() - .wrapped_sql - .unwrap() - .sql - // Expect 6 backslashes as output is JSON and it's escaped one more time - .contains("\\\\\\\\\\\\`")); - } - #[tokio::test] async fn test_extract_epoch_from_dimension() { if !Rewriter::sql_push_down_enabled() { @@ -18428,7 +17476,7 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), "quicksight_sql_implementation_info", execute_query( r#" - SELECT character_value, version() + SELECT character_value, version() FROM INFORMATION_SCHEMA.SQL_IMPLEMENTATION_INFO WHERE implementation_info_id IN ('17','18') "# @@ -18535,7 +17583,7 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), ) select case - when diststyle = 'ALL' then size/cast(nodes.node_count as float) + when diststyle = 'ALL' then size/cast(nodes.node_count as float) else size end as sizeMBs from SVV_TABLE_INFO @@ -18561,7 +17609,7 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), r#" select nspname from pg_external_schema pe - join pg_namespace pn on pe.esoid = pn.oid + join pg_namespace pn on pe.esoid = pn.oid where nspowner != 1 and nspname = 'public' diff --git a/rust/cubesql/cubesql/src/compile/test/mod.rs b/rust/cubesql/cubesql/src/compile/test/mod.rs index 544ff648bdf4b..8f014f7b21807 100644 --- a/rust/cubesql/cubesql/src/compile/test/mod.rs +++ b/rust/cubesql/cubesql/src/compile/test/mod.rs @@ -38,6 +38,8 @@ pub mod test_introspection; pub mod test_udfs; #[cfg(test)] pub mod test_user_change; +#[cfg(test)] +pub mod test_wrapper; pub mod utils; pub use utils::*; diff --git a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs new file mode 100644 index 0000000000000..1eb78f92fa5e6 --- /dev/null +++ b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs @@ -0,0 +1,966 @@ +use datafusion::physical_plan::displayable; +use std::sync::Arc; + +use crate::{ + compile::{ + rewrite::rewriter::Rewriter, + test::{ + convert_select_to_query_plan, convert_select_to_query_plan_customized, + convert_select_to_query_plan_with_config, init_testing_logger, LogicalPlanTestUtils, + }, + DatabaseProtocol, + }, + config::ConfigObjImpl, +}; + +#[tokio::test] +async fn test_simple_wrapper() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT COALESCE(customer_gender, 'N/A', 'NN'), AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("COALESCE")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_wrapper_group_by_rollup() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, notes, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, ROLLUP(2)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("Rollup")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_wrapper_group_by_rollup_with_aliases() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender as \"customer_gender1\", notes as \"notes\", AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY ROLLUP(1, 2)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("Rollup")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_wrapper_group_by_rollup_nested() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, notes, avg(mp) from (SELECT customer_gender, notes, avg(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, 2) b GROUP BY ROLLUP(1, 2)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("ROLLUP(1, 2)")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_wrapper_group_by_rollup_nested_from_asterisk() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, notes, avg(avgPrice) from (SELECT * FROM KibanaSampleDataEcommerce) b GROUP BY ROLLUP(1, 2) ORDER BY 1" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("Rollup")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_wrapper_group_by_rollup_nested_with_aliases() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender as \"gender\", notes as \"notes\", avg(mp) from (SELECT customer_gender, notes, avg(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, 2) b GROUP BY ROLLUP(1, 2)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("ROLLUP(1, 2)")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_wrapper_group_by_rollup_nested_complex() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, notes, order_date, last_mod, avg(mp) from \ + (SELECT customer_gender, notes, order_date, last_mod, avg(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, 2, 3, 4) b \ + GROUP BY ROLLUP(1), ROLLUP(2), 3, CUBE(4)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("ROLLUP(1), ROLLUP(2), 3, CUBE(4)")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_wrapper_group_by_rollup_placeholders() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, notes, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY ROLLUP(1, 2)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("Rollup")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_wrapper_group_by_cube() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, notes, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY CUBE(customer_gender, notes)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("Cube")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_wrapper_group_by_rollup_complex() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, notes, has_subscription, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY ROLLUP(customer_gender, notes), has_subscription" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("Rollup")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_projection_empty_source() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT (SELECT 'male' where 1 group by 'male' having 1 order by 'male' limit 1) as gender, avgPrice FROM KibanaSampleDataEcommerce a" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("(SELECT")); + assert!(sql.contains("utf8__male__")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); + //println!("phys plan {:?}", physical_plan); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_filter_empty_source() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT avgPrice FROM KibanaSampleDataEcommerce a where customer_gender = (SELECT 'male' )" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("(SELECT")); + assert!(sql.contains("utf8__male__")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); + //println!("phys plan {:?}", physical_plan); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_projection_aggregate_empty_source() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT (SELECT 'male'), avg(avgPrice) FROM KibanaSampleDataEcommerce a GROUP BY 1" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("(SELECT")); + assert!(sql.contains("utf8__male__")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_filter_in_empty_source() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, avgPrice FROM KibanaSampleDataEcommerce a where customer_gender in (select 'male')" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("IN (SELECT")); + assert!(sql.contains("utf8__male__")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_filter_and_projection_empty_source() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT (select 'male'), avgPrice FROM KibanaSampleDataEcommerce a where customer_gender in (select 'female')" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + + let sql = logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("IN (SELECT")); + assert!(sql.contains("(SELECT")); + assert!(sql.contains("utf8__male__")); + assert!(sql.contains("utf8__female__")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_projection() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT (SELECT customer_gender FROM KibanaSampleDataEcommerce LIMIT 1) as gender, avgPrice FROM KibanaSampleDataEcommerce a" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("(SELECT")); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("\\\\\\\"limit\\\\\\\":1")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_projection_aggregate() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT (SELECT customer_gender FROM KibanaSampleDataEcommerce WHERE customer_gender = 'male' LIMIT 1), avg(avgPrice) FROM KibanaSampleDataEcommerce a GROUP BY 1" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("(SELECT")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_filter_equal() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, avgPrice FROM KibanaSampleDataEcommerce a where customer_gender = (select customer_gender from KibanaSampleDataEcommerce limit 1)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("(SELECT")); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("\\\\\\\"limit\\\\\\\":1")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_filter_in() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT customer_gender, avgPrice FROM KibanaSampleDataEcommerce a where customer_gender in (select customer_gender from KibanaSampleDataEcommerce)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("IN (SELECT")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_simple_subquery_wrapper_filter_and_projection() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT (select customer_gender from KibanaSampleDataEcommerce limit 1), avgPrice FROM KibanaSampleDataEcommerce a where customer_gender in (select customer_gender from KibanaSampleDataEcommerce)" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("IN (SELECT")); + + let _physical_plan = query_plan.as_physical_plan().await.unwrap(); +} + +#[tokio::test] +async fn test_case_wrapper() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT CASE WHEN COALESCE(customer_gender, 'N/A', 'NN') = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("CASE WHEN")); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_distinct() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + r#"SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, COUNT(DISTINCT countDistinct) mp + FROM KibanaSampleDataEcommerce a + WHERE + ( + ( + ( a.order_date ) >= '2024-01-01' + AND ( a.order_date ) < '2024-02-01' + ) + ) + GROUP BY 1"# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("CASE WHEN")); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_alias_with_order() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END AS \"f822c516-3515-11c2-8464-5d4845a02f73\", AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END ORDER BY CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END NULLS FIRST LIMIT 500" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("ORDER BY \"case_when_a_cust\"")); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_ungrouped() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("CASE WHEN")); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_non_strict_match() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let mut config = ConfigObjImpl::default(); + + config.disable_strict_agg_type_match = true; + + let query_plan = convert_select_to_query_plan_with_config( + "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, SUM(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1" + .to_string(), + DatabaseProtocol::PostgreSQL, + Arc::new(config) + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("CASE WHEN")); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_ungrouped_sorted() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1 ORDER BY 1 DESC" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("ORDER BY")); +} + +#[tokio::test] +async fn test_case_wrapper_ungrouped_sorted_aliased() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT x FROM (SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END x, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1 ORDER BY 1 DESC) b" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + // TODO test without depend on column name + .contains("ORDER BY \"case_when")); +} + +#[tokio::test] +async fn test_case_wrapper_with_internal_limit() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1 LIMIT 1123" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("CASE WHEN")); + + assert!( + logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("1123"), + "SQL contains 1123: {}", + logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + ); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_with_system_fields() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, __user, __cubeJoinField, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1, 2, 3 LIMIT 1123" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + + assert!( + logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains( + "\\\"cube_name\\\":\\\"KibanaSampleDataEcommerce\\\",\\\"alias\\\":\\\"user\\\"" + ), + r#"SQL contains `\"cube_name\":\"KibanaSampleDataEcommerce\",\"alias\":\"user\"` {}"#, + logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + ); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_with_limit() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT * FROM (SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1) q LIMIT 1123" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("CASE WHEN")); + + assert!( + logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("1123"), + "SQL contains 1123: {}", + logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + ); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_with_null() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT CASE WHEN taxful_total_price IS NULL THEN NULL WHEN taxful_total_price < taxful_total_price * 2 THEN COALESCE(taxful_total_price, 0, 0) END, AVG(avgPrice) FROM KibanaSampleDataEcommerce GROUP BY 1" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + .contains("CASE WHEN")); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_ungrouped_on_dimension() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + "SELECT CASE WHEN SUM(taxful_total_price) > 0 THEN SUM(taxful_total_price) ELSE 0 END FROM KibanaSampleDataEcommerce a" + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +#[tokio::test] +async fn test_case_wrapper_escaping() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan_customized( + "SELECT CASE WHEN customer_gender = '\\`' THEN COALESCE(customer_gender, 'N/A', 'NN') ELSE 'N/A' END as \"\\`\", AVG(avgPrice) FROM KibanaSampleDataEcommerce a GROUP BY 1".to_string(), + DatabaseProtocol::PostgreSQL, + vec![ + ("expressions/binary".to_string(), "{{ left }} \\`{{ op }} {{ right }}".to_string()) + ], + ).await; + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); + + let logical_plan = query_plan.as_logical_plan(); + assert!(logical_plan + .find_cube_scan_wrapper() + .wrapped_sql + .unwrap() + .sql + // Expect 6 backslashes as output is JSON and it's escaped one more time + .contains("\\\\\\\\\\\\`")); +}