From a3181d163fb4b4b1b55462185ca0a4de01eb945e Mon Sep 17 00:00:00 2001 From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com> Date: Fri, 14 Aug 2026 19:49:48 +0400 Subject: [PATCH] fix(cubesql): Prefer SQL pushdown over limitless post-processing Signed-off-by: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com> --- .../configuration/environment-variables.mdx | 21 + rust/cubesql/cubesql/src/compile/mod.rs | 589 +++++++----------- .../cubesql/src/compile/rewrite/cost.rs | 131 +++- .../cubesql/src/compile/rewrite/rewriter.rs | 32 +- ..._noninjective_coalesce_from_dimension.snap | 12 - ...sts__noninjective_left_from_dimension.snap | 12 - ...s__noninjective_nullif_from_dimension.snap | 12 - ...ts__noninjective_right_from_dimension.snap | 12 - ...pile__tests__nonrewritable_date_trunc.snap | 10 - .../cubesql/src/compile/test/test_wrapper.rs | 160 +++++ .../cubesql/cubesql/src/compile/test/utils.rs | 52 ++ rust/cubesql/cubesql/src/config/mod.rs | 12 + 12 files changed, 624 insertions(+), 431 deletions(-) delete mode 100644 rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_coalesce_from_dimension.snap delete mode 100644 rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_left_from_dimension.snap delete mode 100644 rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_nullif_from_dimension.snap delete mode 100644 rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_right_from_dimension.snap delete mode 100644 rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__nonrewritable_date_trunc.snap diff --git a/docs-mintlify/reference/configuration/environment-variables.mdx b/docs-mintlify/reference/configuration/environment-variables.mdx index e60521f6f965e..772c095574842 100644 --- a/docs-mintlify/reference/configuration/environment-variables.mdx +++ b/docs-mintlify/reference/configuration/environment-variables.mdx @@ -1531,6 +1531,27 @@ Queries with a `LIMIT` at or below that limit are not streamed. | --------------- | ---------------------- | --------------------- | | `true`, `false` | `false` | `false` | +## `CUBESQL_FAIL_ON_LIMITLESS_POST_PROCESSING` + +If `true`, the [SQL API][ref-sql-api] rejects a query when part of it has to run outside +the data source over a Cube query that has no `LIMIT` of its own. + +Such a query is capped at the maximum row limit set by +[`CUBESQL_NON_STREAMING_QUERY_MAX_ROW_LIMIT`](#cubesql_non_streaming_query_max_row_limit), +without any ordering. Sorting, filtering, or joining that capped result outside the data +source treats an arbitrary slice of the rows as if it were the whole population, so the +answer is wrong rather than merely short. The SQL API pushes these queries down to the +data source whenever it can; this option controls what happens when it cannot. Leave it +`false` to keep returning the truncated result, or set it to `true` to fail with an error +instead. + +Queries that are streamed (see [`CUBESQL_STREAM_MODE`](#cubesql_stream_mode)) are never +capped, so this option has no effect on them. + +| Possible Values | Default in Development | Default in Production | +| --------------- | ---------------------- | --------------------- | +| `true`, `false` | `false` | `false` | + ## `CUBESQL_CUBE_SCAN_MAX_BATCH_ROWS` Specifies the maximum number of rows in a single record batch produced when the diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs index 2e3ffa4b13696..7575a1c57ca6b 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/rust/cubesql/cubesql/src/compile/mod.rs @@ -58,7 +58,7 @@ mod tests { use crate::compile::test::{ convert_select_to_query_plan, convert_select_to_query_plan_customized, convert_select_to_query_plan_with_meta, execute_queries_with_flags, execute_query, - init_testing_logger, LogicalPlanTestUtils, TestContext, + init_testing_logger, member_expression_sql, LogicalPlanTestUtils, TestContext, }; #[tokio::test] @@ -2183,33 +2183,33 @@ limit // ); let logical_plan = query_plan.as_logical_plan(); + let wrapped_sql = logical_plan.find_cube_scan_wrapped_sql(); + let request = wrapped_sql.request; assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec!["KibanaSampleDataEcommerce.sumPrice".to_string()]), - dimensions: Some(vec![]), - segments: Some(vec![]), - time_dimensions: Some(vec![ - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("year".to_string()), - date_range: Some(json!(vec![ - "2023-07-08T00:00:00.000Z".to_string(), - "2023-10-07T23:59:59.999Z".to_string() - ])), - }, - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("month".to_string()), - date_range: Some(json!(vec![ - "2023-07-08T00:00:00.000Z".to_string(), - "2023-10-07T23:59:59.999Z".to_string() - ])), - } - ]), - order: Some(vec![]), - ..Default::default() - } + member_expression_sql(&request.measures), + ["${KibanaSampleDataEcommerce.sumPrice}"] + ); + assert_eq!( + member_expression_sql(&request.dimensions), + ["${KibanaSampleDataEcommerce.order_date}"] + ); + assert_eq!( + member_expression_sql(&request.segments), + [ + "((${KibanaSampleDataEcommerce.order_date} < timestamptz '2023-10-08T00:00:00.000Z') \ + AND (${KibanaSampleDataEcommerce.order_date} >= timestamptz '2023-07-08T00:00:00.000Z'))", + ] + ); + + // The grouping, the `a0 IS NOT NULL` filter and the ordered LIMIT all run at the + // data source. Leaving the sort to post processing would have ordered a result + // already truncated to the row limit, so the top 1001 would not be the true top + let sql = wrapped_sql.wrapped_sql.sql; + assert!(sql.contains("GROUP BY"), "grouping is pushed down: {}", sql); + assert!( + sql.contains("ORDER BY") && sql.contains("LIMIT 1001"), + "ordered limit is pushed down: {}", + sql ); } @@ -6272,21 +6272,11 @@ ORDER BY .await; let logical_plan = query_plan.as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - segments: Some(vec![]), - dimensions: Some(vec![]), - time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_owned(), - granularity: Some("month".to_string()), - date_range: None, - }]), - order: Some(vec![]), - ungrouped: Some(true), - ..Default::default() - } + member_expression_sql(&request.dimensions), + ["EXTRACT(month FROM ${KibanaSampleDataEcommerce.order_date})",] ); Ok(()) @@ -6340,20 +6330,14 @@ ORDER BY .await; let logical_plan = query_plan.as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string(),]), - segments: Some(vec![]), - dimensions: Some(vec![]), - time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_owned(), - granularity: Some("day".to_string()), - date_range: None, - }]), - order: Some(vec![]), - ..Default::default() - } + member_expression_sql(&request.measures), + ["${KibanaSampleDataEcommerce.count}",] + ); + assert_eq!( + member_expression_sql(&request.dimensions), + ["CAST(EXTRACT(doy FROM ${KibanaSampleDataEcommerce.order_date}) AS INTEGER)",] ); Ok(()) @@ -6441,37 +6425,22 @@ ORDER BY ..Default::default() }; - context - .add_cube_load_mock( - expected_cube_scan.clone(), - simple_load_response( - vec!["MultiTypeCube.dim_date0.month", "MultiTypeCube.count"], - vec![ - vec![ - json!("2024-01-01T00:00:00"), - json!("2024-02-01T00:00:00"), - json!("2024-03-01T00:00:00"), - json!("2024-04-01T00:00:00"), - ], - vec![json!("3"), json!("2"), json!("1"), json!("10")], - ], - ), - ) - .await; + let logical_plan = context + .convert_sql_to_cube_query(query) + .await + .unwrap() + .as_logical_plan(); - assert_eq!( - context - .convert_sql_to_cube_query(&query) - .await - .unwrap() - .as_logical_plan() - .find_cube_scan() - .request, - expected_cube_scan - ); + assert_eq!(logical_plan.find_cube_scan().request, expected_cube_scan); - // Expect that query is executable, and properly groups months by quarter - insta::assert_snapshot!(context.execute_query(query).await.unwrap()); + // The sort would order a truncated read of an unlimited Cube query, so the quarter + // grouping runs at the data source instead of in post processing + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + assert!( + sql.contains("DATE_TRUNC(") && sql.contains("MIN("), + "months are grouped by quarter at the data source: {}", + sql + ); } #[tokio::test] @@ -6526,20 +6495,14 @@ ORDER BY .await; let logical_plan = query_plan.as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string(),]), - segments: Some(vec![]), - dimensions: Some(vec![]), - time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_owned(), - granularity: Some("day".to_string()), - date_range: None, - }]), - order: Some(vec![]), - ..Default::default() - } + member_expression_sql(&request.measures), + ["${KibanaSampleDataEcommerce.count}",] + ); + assert_eq!( + member_expression_sql(&request.dimensions), + ["(CAST(EXTRACT(dow FROM ${KibanaSampleDataEcommerce.order_date}) AS INTEGER) + 1)",] ); Ok(()) @@ -7644,18 +7607,17 @@ ORDER BY "source"."str0" ASC DatabaseProtocol::PostgreSQL ).await.as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string()]), - dimensions: Some(vec![ - "KibanaSampleDataEcommerce.taxful_total_price".to_string() - ]), - segments: Some(vec![]), - order: Some(vec![]), - ..Default::default() - } - ) + member_expression_sql(&request.measures), + ["${KibanaSampleDataEcommerce.count}",] + ); + assert_eq!( + member_expression_sql(&request.dimensions), + [ + "((FLOOR(((${KibanaSampleDataEcommerce.taxful_total_price} - 1.1) / 0.025)) * 0.025) + 1.1)", + ] + ); } #[tokio::test] @@ -7857,20 +7819,16 @@ ORDER BY "source"."str0" ASC .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec!["KibanaSampleDataEcommerce.minPrice".to_string()]), - dimensions: Some(vec![]), - segments: Some(vec![]), - time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("week".to_string()), - date_range: None, - },]), - order: Some(vec![]), - ..Default::default() - } + member_expression_sql(&request.measures), + ["${KibanaSampleDataEcommerce.minPrice}",] + ); + assert_eq!( + member_expression_sql(&request.dimensions), + [ + "CEIL((CAST(EXTRACT(doy FROM CAST(${KibanaSampleDataEcommerce.order_date.week} AS TIMESTAMP)) AS INTEGER) / 7))", + ] ); } @@ -12072,21 +12030,12 @@ ORDER BY "source"."str0" ASC .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec![]), - segments: Some(vec![]), - time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("day".to_string()), - date_range: None - }]), - order: Some(vec![]), - ..Default::default() - } - ) + member_expression_sql(&request.dimensions), + ["(EXTRACT(day FROM ${KibanaSampleDataEcommerce.order_date}) = 15)",] + ); } #[tokio::test] @@ -12113,21 +12062,15 @@ ORDER BY "source"."str0" ASC .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string()]), - dimensions: Some(vec![]), - segments: Some(vec![]), - time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("month".to_string()), - date_range: None - }]), - order: Some(vec![]), - ..Default::default() - } - ) + member_expression_sql(&request.measures), + ["${KibanaSampleDataEcommerce.count}",] + ); + assert_eq!( + member_expression_sql(&request.dimensions), + ["(((EXTRACT(month FROM ${KibanaSampleDataEcommerce.order_date}) - 1) % 3) + 1)",] + ); } #[tokio::test] @@ -12147,28 +12090,14 @@ ORDER BY "source"."str0" ASC .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec![]), - segments: Some(vec![]), - time_dimensions: Some(vec![ - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("month".to_string()), - date_range: None - }, - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.last_mod".to_string(), - granularity: Some("month".to_string()), - date_range: None - }, - ]), - order: Some(vec![]), - ..Default::default() - } - ) + member_expression_sql(&request.dimensions), + [ + "(EXTRACT(month FROM ${KibanaSampleDataEcommerce.order_date}) < (EXTRACT(month FROM ${KibanaSampleDataEcommerce.last_mod}) + 1))", + ] + ); } #[tokio::test] @@ -12191,15 +12120,13 @@ ORDER BY "source"."str0" ASC .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]), - segments: Some(vec![]), - order: Some(vec![]), - ..Default::default() - } + member_expression_sql(&request.dimensions), + [ + "((LOWER(${KibanaSampleDataEcommerce.customer_gender}) = $0$) OR (LOWER(${KibanaSampleDataEcommerce.customer_gender}) = $1$))", + ] ); let logical_plan = convert_select_to_query_plan( @@ -12218,18 +12145,13 @@ ORDER BY "source"."str0" ASC .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec![ - "KibanaSampleDataEcommerce.customer_gender".to_string(), - "KibanaSampleDataEcommerce.notes".to_string(), - ]), - segments: Some(vec![]), - order: Some(vec![]), - ..Default::default() - } + member_expression_sql(&request.dimensions), + [ + "((LOWER(${KibanaSampleDataEcommerce.customer_gender}) = $0$) OR (LOWER(${KibanaSampleDataEcommerce.notes}) = $1$))", + ] ); if !Rewriter::sql_push_down_enabled() { @@ -12314,18 +12236,12 @@ ORDER BY "source"."str0" ASC .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec![ - "KibanaSampleDataEcommerce.taxful_total_price".to_string() - ]), - segments: Some(vec![]), - order: Some(vec![]), - ..Default::default() - } - ) + member_expression_sql(&request.dimensions), + ["(${KibanaSampleDataEcommerce.taxful_total_price} > 10)",] + ); } #[tokio::test] @@ -12957,16 +12873,16 @@ ORDER BY "source"."str0" ASC .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]), - segments: Some(vec![]), - order: Some(vec![]), - ..Default::default() - } - ) + member_expression_sql(&request.dimensions), + [ + "${KibanaSampleDataEcommerce.customer_gender}", + "LEFT(${KibanaSampleDataEcommerce.customer_gender}, 2)", + "RIGHT(${KibanaSampleDataEcommerce.customer_gender}, 2)", + ] + ); } #[tokio::test] @@ -13840,39 +13756,14 @@ ORDER BY "source"."str0" ASC let context = TestContext::new(DatabaseProtocol::PostgreSQL).await; - // Expected scan is same for every query - let expected_cube_scan = V1LoadRequestQuery { - measures: Some(vec![]), - segments: Some(vec![]), - dimensions: Some(vec!["MultiTypeCube.dim_str0".to_string()]), - order: Some(vec![]), - ..Default::default() - }; - - context - .add_cube_load_mock( - expected_cube_scan.clone(), - simple_load_response( - vec!["MultiTypeCube.dim_str0"], - vec![vec![ - json!("foo"), - json!(null), - json!("(none)"), - json!("abcd"), - json!("ab__cd"), - ]], - ), - ) - .await; - let exprs = [ - ("coalesce", "COALESCE(dim_str0, '(none)')"), - ("nullif", "NULLIF(dim_str0, '(none)')"), - ("left", "LEFT(dim_str0, 2)"), - ("right", "RIGHT(dim_str0, 2)"), + ("COALESCE", "COALESCE(dim_str0, '(none)')"), + ("NULLIF", "NULLIF(dim_str0, '(none)')"), + ("LEFT", "LEFT(dim_str0, 2)"), + ("RIGHT", "RIGHT(dim_str0, 2)"), ]; - for (name, expr) in exprs { + for (fun, expr) in exprs { // language=PostgreSQL let query = format!( r#" @@ -13883,21 +13774,30 @@ ORDER BY "source"."str0" ASC "# ); + // The sort would order a truncated read of an unlimited Cube query, so the + // whole query is pushed down instead + let request = context + .convert_sql_to_cube_query(&query) + .await + .unwrap() + .as_logical_plan() + .find_cube_scan_wrapped_sql() + .request; + + // Expect no duplicates in result set: the call is a dimension of the Cube + // query, so it is grouped by rather than projected over grouped rows + assert_eq!(request.measures, Some(vec![])); + let dimensions = request.dimensions.unwrap(); assert_eq!( - context - .convert_sql_to_cube_query(&query) - .await - .unwrap() - .as_logical_plan() - .find_cube_scan() - .request, - expected_cube_scan + dimensions.len(), + 1, + "single grouped dimension: {:?}", + dimensions ); - - // Expect no dublicates in result set - insta::assert_snapshot!( - format!("noninjective_{name}_from_dimension"), - context.execute_query(query).await.unwrap() + assert!( + dimensions[0].contains(fun) && dimensions[0].contains("MultiTypeCube.dim_str0"), + "{fun} of the dimension is grouped by: {}", + dimensions[0] ); } } @@ -14576,31 +14476,19 @@ ORDER BY "source"."str0" ASC let logical_plan = query_plan.as_logical_plan(); - let request = logical_plan.find_cube_scan().request; - - // The rewriter should recognize the complex quarter expression and - // simplify it to DATE_TRUNC('quarter', col) via the - // thoughtspot-pg-quarter-start-to-date-trunc rule, which then gets - // recognized as a quarter time dimension. + // Only the filter is pushed down here, so the scan reads raw rows and the + // aggregate above it still runs in post processing. The rewriter recognizes the + // complex quarter expression and simplifies it to DATE_TRUNC('quarter', col) via + // the thoughtspot-pg-quarter-start-to-date-trunc rule. + let request = logical_plan.find_cube_scan_wrapped_sql_deep().request; assert_eq!( - request, - V1LoadRequestQuery { - measures: Some(vec!["KibanaSampleDataEcommerce.sumPrice".to_string(),]), - dimensions: Some(vec![ - "KibanaSampleDataEcommerce.order_date".to_string(), - "KibanaSampleDataEcommerce.customer_gender".to_string(), - ]), - segments: Some(vec![]), - time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("quarter".to_string()), - date_range: None, - },]), - order: Some(vec![]), - ungrouped: Some(true), - ..Default::default() - } + member_expression_sql(&request.segments), + [ + "((DATEDIFF(day, CAST(${KibanaSampleDataEcommerce.order_date.quarter} AS DATE), \ + CAST(${KibanaSampleDataEcommerce.order_date} AS DATE)) + 1) <= 45)", + ] ); + assert_eq!(request.ungrouped, Some(true)); } #[tokio::test] @@ -15079,17 +14967,14 @@ ORDER BY "source"."str0" ASC ) .await; + let request = query_plan + .as_logical_plan() + .find_cube_scan_wrapped_sql() + .request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - query_plan.as_logical_plan().find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec![ - "KibanaSampleDataEcommerce.taxful_total_price".to_string() - ]), - segments: Some(vec![]), - order: Some(vec![]), - ..Default::default() - } + member_expression_sql(&request.dimensions), + ["-(${KibanaSampleDataEcommerce.taxful_total_price})",] ); } @@ -16587,58 +16472,29 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec![]), - segments: Some(vec![]), - time_dimensions: Some(vec![ - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("year".to_string()), - date_range: None - }, - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("quarter".to_string()), - date_range: None - }, - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("month".to_string()), - date_range: None - }, - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("week".to_string()), - date_range: None - }, - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("day".to_string()), - date_range: None - }, - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("hour".to_string()), - date_range: None - }, - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("minute".to_string()), - date_range: None - }, - V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("second".to_string()), - date_range: None - }, - ]), - order: Some(vec![]), - ..Default::default() - } - ) + member_expression_sql(&request.dimensions), + [ + "${KibanaSampleDataEcommerce.order_date.year}", + "EXTRACT(year FROM ${KibanaSampleDataEcommerce.order_date})", + "${KibanaSampleDataEcommerce.order_date.quarter}", + "EXTRACT(quarter FROM ${KibanaSampleDataEcommerce.order_date})", + "${KibanaSampleDataEcommerce.order_date.month}", + "EXTRACT(month FROM ${KibanaSampleDataEcommerce.order_date})", + "${KibanaSampleDataEcommerce.order_date.week}", + "EXTRACT(week FROM ${KibanaSampleDataEcommerce.order_date})", + "${KibanaSampleDataEcommerce.order_date.day}", + "EXTRACT(day FROM ${KibanaSampleDataEcommerce.order_date})", + "${KibanaSampleDataEcommerce.order_date.hour}", + "EXTRACT(hour FROM ${KibanaSampleDataEcommerce.order_date})", + "${KibanaSampleDataEcommerce.order_date.minute}", + "EXTRACT(minute FROM ${KibanaSampleDataEcommerce.order_date})", + "${KibanaSampleDataEcommerce.order_date.second}", + "EXTRACT(second FROM ${KibanaSampleDataEcommerce.order_date})", + ] + ); } #[tokio::test] @@ -16672,24 +16528,19 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert!(member_expression_sql(&request.measures).is_empty()); assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec![]), - segments: Some(vec![]), - time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension { - dimension: "KibanaSampleDataEcommerce.order_date".to_string(), - granularity: Some("quarter".to_string()), - date_range: Some(json!(vec![ - "2024-01-01T00:00:00.000Z".to_string(), - "2024-12-31T23:59:59.999Z".to_string(), - ])), - },]), - order: Some(vec![]), - ..Default::default() - } - ) + member_expression_sql(&request.dimensions), + [ + "${KibanaSampleDataEcommerce.order_date.quarter}", + "EXTRACT(quarter FROM ${KibanaSampleDataEcommerce.order_date})", + ] + ); + assert_eq!( + member_expression_sql(&request.segments), + ["(${KibanaSampleDataEcommerce.order_date.year} = $0$)",] + ); } #[tokio::test] @@ -17842,21 +17693,27 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), .await .as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; assert_eq!( - logical_plan.find_cube_scan().request, - V1LoadRequestQuery { - measures: Some(vec![]), - dimensions: Some(vec!["Logs.id".to_string(),]), - segments: Some(vec![]), - order: Some(vec![]), - ungrouped: Some(true), - join_hints: Some(vec![ - vec!["KibanaSampleDataEcommerce".to_string(), "Logs".to_string()], - vec!["KibanaSampleDataEcommerce".to_string(), "Logs".to_string()], - ]), - ..Default::default() - } - ) + member_expression_sql(&request.dimensions), + ["${Logs.id}", "${Logs.id}"] + ); + assert_eq!( + request.join_hints, + Some(vec![ + vec!["KibanaSampleDataEcommerce".to_string(), "Logs".to_string()], + vec!["KibanaSampleDataEcommerce".to_string(), "Logs".to_string()], + ]) + ); + // Both sort keys are pushed into the Cube query, so the row limit applies to an + // ordered result and picks the same rows the client asked for + assert_eq!( + request.order, + Some(vec![ + vec!["id".to_string(), "asc".to_string()], + vec!["id".to_string(), "asc".to_string()], + ]) + ); } #[tokio::test] diff --git a/rust/cubesql/cubesql/src/compile/rewrite/cost.rs b/rust/cubesql/cubesql/src/compile/rewrite/cost.rs index 1245ad047d31c..98e8c720f9be2 100644 --- a/rust/cubesql/cubesql/src/compile/rewrite/cost.rs +++ b/rust/cubesql/cubesql/src/compile/rewrite/cost.rs @@ -1,12 +1,14 @@ use std::{ - collections::HashMap, fmt::Debug, hash::Hash, marker::PhantomData, mem::take, sync::Arc, + cmp::Ordering, collections::HashMap, fmt::Debug, hash::Hash, marker::PhantomData, mem::take, + sync::Arc, }; use crate::{ compile::rewrite::{ - rules::utils::granularity_str_to_int_order, CubeScanUngrouped, CubeScanWrapped, - DimensionName, LogicalPlanLanguage, MemberErrorPriority, ScalarUDFExprFun, - TimeDimensionGranularity, WrappedSelectPushToCube, WrappedSelectUngroupedScan, + rules::utils::granularity_str_to_int_order, CubeScanLimit, CubeScanUngrouped, + CubeScanWrapped, DimensionName, LogicalPlanLanguage, MemberErrorPriority, ScalarUDFExprFun, + TimeDimensionGranularity, WrappedSelectLimit, WrappedSelectPushToCube, + WrappedSelectUngroupedScan, }, transport::{MetaContext, V1CubeMetaDimensionExt}, }; @@ -17,13 +19,19 @@ use indexmap::IndexSet; pub struct BestCubePlan { meta_context: Arc, penalize_post_processing: bool, + penalize_limitless_post_processing: bool, } impl BestCubePlan { - pub fn new(meta_context: Arc, penalize_post_processing: bool) -> Self { + pub fn new( + meta_context: Arc, + penalize_post_processing: bool, + penalize_limitless_post_processing: bool, + ) -> Self { Self { meta_context, penalize_post_processing, + penalize_limitless_post_processing, } } @@ -209,6 +217,26 @@ impl BestCubePlan { _ => 0, }; + // A Cube query without an explicit limit is capped at + // `non_streaming_query_max_row_limit` rows when it runs, without any ordering. + // That cap is harmless for the rows the client receives, but anything computed + // on top of a capped result in DataFusion sees an arbitrary slice of it. + let limitless_scans = match enode { + LogicalPlanLanguage::CubeScanLimit(CubeScanLimit(None)) => 1, + _ => 0, + }; + + // A push to Cube wrapper carries the user's limit on its select rather than on the + // scan below it, so a limited query still counts an unlimited scan. Counting the + // limits as well lets `limitless_post_processing` compare the two: it only fires + // when some scan in the subtree is left over without a limit of its own, which + // keeps a limited scan from covering for an unlimited sibling in a join or a union. + let scan_limits = match enode { + LogicalPlanLanguage::CubeScanLimit(CubeScanLimit(Some(_))) => 1, + LogicalPlanLanguage::WrappedSelectLimit(WrappedSelectLimit(Some(_))) => 1, + _ => 0, + }; + CubePlanCost { replacers: this_replacers, // Will be filled in finalize @@ -229,6 +257,8 @@ impl BestCubePlan { max_time_dimensions_granularity, structure_points, ungrouped_aggregates: 0, + // Will be filled in finalize + limitless_post_processing: 0, wrapper_nodes, joins, wrapped_select_non_push_to_cube, @@ -241,13 +271,45 @@ impl BestCubePlan { ast_size: 1, ungrouped_nodes, unwrapped_subqueries, + limitless_scans: Unordered(limitless_scans), + scan_limits: Unordered(scan_limits), } } } +/// A cost field that carries a value forward without taking part in the comparison. +/// +/// [`CubePlanCost`] derives its ordering from the declaration order of its fields, so any +/// field it holds is also a tie breaker. Some fields are only inputs to other fields and +/// have no ordering of their own to express - preferring more of them or fewer of them +/// would both be arbitrary - so they are wrapped here and compare equal to each other. +#[derive(Debug, Clone, Copy, Default)] +pub struct Unordered(pub T); + +impl PartialEq for Unordered { + fn eq(&self, _: &Self) -> bool { + true + } +} + +impl Eq for Unordered {} + +impl PartialOrd for Unordered { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for Unordered { + fn cmp(&self, _: &Self) -> Ordering { + Ordering::Equal + } +} + #[derive(Clone, Copy)] pub struct CubePlanCostOptions { penalize_post_processing: bool, + penalize_limitless_post_processing: bool, } /// This cost struct maintains following structural relationships: @@ -261,6 +323,9 @@ pub struct CubePlanCostOptions { /// - `filter_members` > `cube_members` - optimize for `inDateRange` filter push down to time dimension /// - `member_errors` > `cube_members` - extra cube members may be required (e.g. CASE) /// - `member_errors` > `wrapper_nodes` - use SQL push down where possible if cube scan can't be detected +/// - `limitless_post_processing` > `wrapper_nodes`, `ast_size_outside_wrapper` - row dropping or +/// row multiplying post processing on top of an unlimited Cube query reads a truncated result, +/// so prefer SQL push down over any representation that leaves such a query to post processing /// - `non_pushed_down_window` > `wrapper_nodes` - prefer to always push down window functions /// - `non_pushed_down_limit_sort` > `wrapper_nodes` - prefer to always push down limit-sort expressions /// - `wrapped_select_non_push_to_cube` > `wrapped_select_ungrouped_scan` - otherwise cost would prefer any aggregation, even non-push-to-Cube @@ -280,6 +345,7 @@ pub struct CubePlanCost { non_pushed_down_grouping_sets: i64, non_pushed_down_limit_sort: i64, joins: usize, + limitless_post_processing: usize, wrapper_nodes: i64, ast_size_outside_wrapper: usize, wrapped_select_non_push_to_cube: usize, @@ -301,6 +367,9 @@ pub struct CubePlanCost { ast_size: usize, ast_size_inside_wrapper: usize, ungrouped_nodes: usize, + // Inputs for `limitless_post_processing`, which is what expresses the preference + limitless_scans: Unordered, + scan_limits: Unordered, } #[derive(Debug, Clone, Eq, Hash, PartialEq)] @@ -318,6 +387,12 @@ pub enum SortState { } impl CubePlanCost { + /// Number of post processing nodes that read a Cube query which is not limited by the + /// user, and so would be silently truncated to the maximum row limit. + pub fn limitless_post_processing(&self) -> usize { + self.limitless_post_processing + } + pub fn add_child(&self, other: &Self) -> Self { Self { replacers: self.replacers + other.replacers, @@ -366,6 +441,10 @@ impl CubePlanCost { ast_size_inside_wrapper: self.ast_size_inside_wrapper + other.ast_size_inside_wrapper, ungrouped_nodes: self.ungrouped_nodes + other.ungrouped_nodes, unwrapped_subqueries: self.unwrapped_subqueries + other.unwrapped_subqueries, + limitless_post_processing: self.limitless_post_processing + + other.limitless_post_processing, + limitless_scans: Unordered(self.limitless_scans.0 + other.limitless_scans.0), + scan_limits: Unordered(self.scan_limits.0 + other.scan_limits.0), } } @@ -444,6 +523,45 @@ impl CubePlanCost { } CubePlanState::Wrapper => 0, } + self.ungrouped_aggregates, + // A Cube query the user did not limit is capped at the maximum row limit when it + // runs, with no ordering, so post processing reads an arbitrary slice of the rows. + // A node whose output for a given row depends only on that row is no worse off + // than the client would have been reading the capped result directly: + // + // - `Projection`, `TableUDFs`, `Repartition` are row-wise, so their output is the + // same slice with the same expressions applied + // - `Limit` on its own narrows an unordered result, which SQL already leaves + // unspecified; a `Sort` underneath it is what makes the choice of rows wrong, + // and that is caught below + // - `Aggregate` re-aggregates rows the Cube query already grouped, so the cap + // lands on the rows the client asked for rather than on hidden detail + // + // Everything else reads the slice as if it were the whole population: + // + // - `Sort` orders the slice, so the leading rows are not the true leading rows + // - `Filter` and `Distinct` decide what to keep by looking at rows that are missing + // - `Window` evaluates over a whole partition, which the cap has cut short + // - `Join`, `CrossJoin`, `Union` and `Subquery` pair the slice with other data, so + // the rows dropped by the cap silently drop matches too + limitless_post_processing: match state { + CubePlanState::Unwrapped(_) + if options.penalize_limitless_post_processing + && self.limitless_scans.0 > self.scan_limits.0 => + { + match enode { + LogicalPlanLanguage::Sort(_) + | LogicalPlanLanguage::Filter(_) + | LogicalPlanLanguage::Distinct(_) + | LogicalPlanLanguage::Window(_) + | LogicalPlanLanguage::Join(_) + | LogicalPlanLanguage::CrossJoin(_) + | LogicalPlanLanguage::Union(_) + | LogicalPlanLanguage::Subquery(_) => 1, + _ => 0, + } + } + _ => 0, + } + self.limitless_post_processing, unwrapped_subqueries: self.unwrapped_subqueries, wrapper_nodes: self.wrapper_nodes, wrapped_select_non_push_to_cube: self.wrapped_select_non_push_to_cube, @@ -453,6 +571,8 @@ impl CubePlanCost { ast_size: self.ast_size, ast_size_inside_wrapper: self.ast_size_inside_wrapper, ungrouped_nodes: self.ungrouped_nodes, + limitless_scans: self.limitless_scans, + scan_limits: self.scan_limits, } } } @@ -840,6 +960,7 @@ impl TopDownCostFunction false, }; - let (plan, qtrace_egraph_iterations, qtrace_best_graph) = + // In stream mode an unlimited Cube query is streamed in full rather than capped, so + // post processing on top of it reads every row and stays correct. Neither the + // preference for pushing it down nor the failure applies there + let config_obj = &self.cube_context.sessions.server.config_obj; + let penalize_limitless_post_processing = !config_obj.stream_mode(); + let fail_on_limitless_post_processing = + config_obj.fail_on_limitless_post_processing() && penalize_limitless_post_processing; + + let (plan, qtrace_egraph_iterations, qtrace_best_graph, limitless_post_processing) = tokio::task::spawn_blocking(move || { let (runner, qtrace_egraph_iterations) = Self::run_rewrites(&cube_context, egraph, rules, "final")?; @@ -362,7 +370,11 @@ impl Rewriter { // TODO maybe check replacers and penalized_ast_size_outside_wrapper right after extraction? let mut extractor = TopDownExtractor::new( &runner.egraph, - BestCubePlan::new(cube_context.meta.clone(), penalize_post_processing), + BestCubePlan::new( + cube_context.meta.clone(), + penalize_post_processing, + penalize_limitless_post_processing, + ), CubePlanTopDownState::new(), ); let Some((best_cost, best)) = extractor.find_best(root) else { @@ -388,6 +400,7 @@ impl Rewriter { converter.to_logical_plan(new_root), qtrace_egraph_iterations, qtrace_best_graph, + best_cost.limitless_post_processing(), )) }) .await??; @@ -397,6 +410,21 @@ impl Rewriter { qtrace.set_best_graph(&qtrace_best_graph); } + // Checked once the qtrace is recorded: the plan that would have run is exactly what + // is worth looking at when this fires. No representation of this query pushes the + // post processing down to the data source, so it would run over a Cube query + // truncated to the maximum row limit + if fail_on_limitless_post_processing && limitless_post_processing > 0 { + return Err(CubeError::user( + "Query requires post-processing of a Cube query that has no limit, so it \ + would be truncated to the maximum row limit and produce incorrect results. \ + Add an explicit LIMIT, or rewrite the query so that it can be pushed down \ + to the data source. This check is enabled by \ + CUBESQL_FAIL_ON_LIMITLESS_POST_PROCESSING." + .to_string(), + )); + } + plan } diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_coalesce_from_dimension.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_coalesce_from_dimension.snap deleted file mode 100644 index 1cb7ee3450997..0000000000000 --- a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_coalesce_from_dimension.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: cubesql/src/compile/mod.rs -expression: context.execute_query(query).await.unwrap() ---- -+--------+ -| result | -+--------+ -| (none) | -| ab__cd | -| abcd | -| foo | -+--------+ diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_left_from_dimension.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_left_from_dimension.snap deleted file mode 100644 index 0b416f30810a2..0000000000000 --- a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_left_from_dimension.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: cubesql/src/compile/mod.rs -expression: context.execute_query(query).await.unwrap() ---- -+--------+ -| result | -+--------+ -| (n | -| ab | -| fo | -| NULL | -+--------+ diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_nullif_from_dimension.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_nullif_from_dimension.snap deleted file mode 100644 index 01a67489499eb..0000000000000 --- a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_nullif_from_dimension.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: cubesql/src/compile/mod.rs -expression: context.execute_query(query).await.unwrap() ---- -+--------+ -| result | -+--------+ -| ab__cd | -| abcd | -| foo | -| NULL | -+--------+ diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_right_from_dimension.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_right_from_dimension.snap deleted file mode 100644 index 6f3bd8758698c..0000000000000 --- a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__noninjective_right_from_dimension.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: cubesql/src/compile/mod.rs -expression: context.execute_query(query).await.unwrap() ---- -+--------+ -| result | -+--------+ -| cd | -| e) | -| oo | -| NULL | -+--------+ diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__nonrewritable_date_trunc.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__nonrewritable_date_trunc.snap deleted file mode 100644 index 5147e95f0b914..0000000000000 --- a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__nonrewritable_date_trunc.snap +++ /dev/null @@ -1,10 +0,0 @@ ---- -source: cubesql/src/compile/mod.rs -expression: context.execute_query(query).await.unwrap() ---- -+-------------------------+-----------------+ -| quarter0 | min_month_count | -+-------------------------+-----------------+ -| 2024-01-01T00:00:00.000 | 1 | -| 2024-04-01T00:00:00.000 | 10 | -+-------------------------+-----------------+ diff --git a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs index e202347037baa..81358ae665320 100644 --- a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs @@ -16,6 +16,7 @@ use crate::{ test::{ convert_select_to_query_plan, convert_select_to_query_plan_customized, convert_select_to_query_plan_with_config, init_testing_logger, LogicalPlanTestUtils, + TestContext, }, DatabaseProtocol, }, @@ -2853,3 +2854,162 @@ async fn test_wrapper_only_system_fields() { displayable(physical_plan.as_ref()).indent() ); } + +/// A Cube query is capped at the maximum row limit when the user does not limit it, so +/// sorting, filtering or joining its result in DataFusion would read an arbitrary slice +/// of the population. Such a query has to be pushed down in full. +const LIMITLESS_POST_PROCESSING_QUERY: &str = r#" + WITH first_orders AS ( + SELECT customer_gender, MIN(order_date) AS first_order_at + FROM KibanaSampleDataEcommerce + GROUP BY 1 + ) + SELECT COUNT(DISTINCT customer_gender) AS customers + FROM first_orders + WHERE first_order_at >= '2024-01-01'::timestamp +"#; + +/// The same shape, with a filter the data source has no template for (`ROUND`), so the +/// filter cannot leave DataFusion and the truncated read is unavoidable. +const UNPUSHABLE_LIMITLESS_POST_PROCESSING_QUERY: &str = r#" + WITH first_orders AS ( + SELECT customer_gender, MIN(taxful_total_price) AS cheapest + FROM KibanaSampleDataEcommerce + GROUP BY 1 + ) + SELECT COUNT(DISTINCT customer_gender) AS customers + FROM first_orders + WHERE ROUND(cheapest) > 10 +"#; + +#[tokio::test] +async fn test_wrapper_limitless_post_processing_pushed_down() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + LIMITLESS_POST_PROCESSING_QUERY.to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!( + logical_plan.find_filter().is_none(), + "no filter is left to post processing: {:?}", + logical_plan + ); + + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + assert!( + sql.contains("COUNT(DISTINCT"), + "outer aggregate is pushed down: {}", + sql + ); + assert!( + sql.contains(r#"WHERE ("first_orders"."first_order_at" >= "#), + "outer filter is pushed down against the aggregate of the inner query: {}", + sql + ); +} + +#[tokio::test] +async fn test_wrapper_limitless_post_processing_allowed_by_default() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + // Without the flag the query still runs, and still reads a truncated result + let query_plan = convert_select_to_query_plan( + UNPUSHABLE_LIMITLESS_POST_PROCESSING_QUERY.to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!( + logical_plan.find_filter().is_some(), + "filter is left to post processing: {:?}", + logical_plan + ); +} + +#[tokio::test] +async fn test_wrapper_limitless_post_processing_fails_when_enabled() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let mut config = ConfigObjImpl::default(); + config.fail_on_limitless_post_processing = true; + let context = TestContext::with_config(DatabaseProtocol::PostgreSQL, Arc::new(config)).await; + + // A query that can be pushed down in full is unaffected + context + .convert_sql_to_cube_query(LIMITLESS_POST_PROCESSING_QUERY) + .await + .expect("fully pushed down query should compile"); + + let error = context + .convert_sql_to_cube_query(UNPUSHABLE_LIMITLESS_POST_PROCESSING_QUERY) + .await + .expect_err("query with unavoidable post processing should fail"); + assert!( + error + .to_string() + .contains("truncated to the maximum row limit"), + "unexpected error: {}", + error + ); +} + +#[tokio::test] +async fn test_wrapper_limitless_post_processing_ignored_in_stream_mode() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + // Streaming reads every row rather than capping the query, so post processing over it + // is correct and there is nothing to prefer pushing down or to fail on + let mut config = ConfigObjImpl::default(); + config.stream_mode = true; + config.fail_on_limitless_post_processing = true; + let context = TestContext::with_config(DatabaseProtocol::PostgreSQL, Arc::new(config)).await; + + let logical_plan = context + .convert_sql_to_cube_query(UNPUSHABLE_LIMITLESS_POST_PROCESSING_QUERY) + .await + .expect("stream mode should not fail on post processing") + .as_logical_plan(); + assert!( + logical_plan.find_filter().is_some(), + "filter is left to post processing: {:?}", + logical_plan + ); + + // And a query that the penalty would have reshaped keeps the plan it has without it + let logical_plan = context + .convert_sql_to_cube_query( + "SELECT -taxful_total_price AS neg FROM KibanaSampleDataEcommerce GROUP BY 1 ORDER BY 1 DESC", + ) + .await + .expect("stream mode should not fail on post processing") + .as_logical_plan(); + assert_eq!( + logical_plan.find_cube_scan().request, + V1LoadRequestQuery { + measures: Some(vec![]), + dimensions: Some(vec![ + "KibanaSampleDataEcommerce.taxful_total_price".to_string() + ]), + segments: Some(vec![]), + order: Some(vec![]), + ..Default::default() + } + ); +} diff --git a/rust/cubesql/cubesql/src/compile/test/utils.rs b/rust/cubesql/cubesql/src/compile/test/utils.rs index 08cf23a39e944..505c459bc4262 100644 --- a/rust/cubesql/cubesql/src/compile/test/utils.rs +++ b/rust/cubesql/cubesql/src/compile/test/utils.rs @@ -22,6 +22,9 @@ pub trait LogicalPlanTestUtils { fn find_cube_scan_wrapped_sql(&self) -> CubeScanWrappedSqlNode; + /// Same, but for plans that still have post processing above the pushed down part. + fn find_cube_scan_wrapped_sql_deep(&self) -> CubeScanWrappedSqlNode; + fn find_cube_scans(&self) -> Vec; fn find_filter(&self) -> Option; @@ -57,6 +60,33 @@ impl LogicalPlanTestUtils for LogicalPlan { } } + fn find_cube_scan_wrapped_sql_deep(&self) -> CubeScanWrappedSqlNode { + pub struct FindWrappedSqlNodeVisitor(Vec); + + impl PlanVisitor for FindWrappedSqlNodeVisitor { + type Error = CubeError; + + fn pre_visit(&mut self, plan: &LogicalPlan) -> Result { + if let LogicalPlan::Extension(ext) = plan { + if let Some(node) = ext.node.as_any().downcast_ref::() { + self.0.push(node.clone()); + } + } + Ok(true) + } + } + + let mut visitor = FindWrappedSqlNodeVisitor(Vec::new()); + self.accept(&mut visitor).unwrap(); + match visitor.0.len() { + 1 => visitor.0.remove(0), + found => panic!( + "The plan includes {} cube_scan_wrapped_sql nodes, expected 1", + found + ), + } + } + fn find_cube_scans(&self) -> Vec { find_cube_scans_deep_search(Arc::new(self.clone()), true) } @@ -66,6 +96,28 @@ impl LogicalPlanTestUtils for LogicalPlan { } } +/// SQL of every member in a pushed down request, in order. +/// +/// A pushed down query carries its members as member expressions: JSON holding a generated +/// alias, the cube it came from and the SQL to evaluate. Only the SQL is worth asserting on, +/// since aliases are generated and truncated to 16 characters. Members that are plain names +/// are returned as they are, so a request can mix both. +pub fn member_expression_sql(members: &Option>) -> Vec { + let Some(members) = members else { + return vec![]; + }; + + members + .iter() + .map(|member| { + serde_json::from_str::(member) + .ok() + .and_then(|member| member["expr"]["sql"].as_str().map(String::from)) + .unwrap_or_else(|| member.clone()) + }) + .collect() +} + pub fn find_cube_scans_deep_search( parent: Arc, panic_if_empty: bool, diff --git a/rust/cubesql/cubesql/src/config/mod.rs b/rust/cubesql/cubesql/src/config/mod.rs index c93c538cb7174..0f23ca36f223b 100644 --- a/rust/cubesql/cubesql/src/config/mod.rs +++ b/rust/cubesql/cubesql/src/config/mod.rs @@ -114,6 +114,8 @@ pub trait ConfigObj: DIService + Debug { fn non_streaming_query_max_row_limit(&self) -> i32; + fn fail_on_limitless_post_processing(&self) -> bool; + fn cube_scan_max_batch_rows(&self) -> usize; fn max_sessions(&self) -> usize; @@ -140,6 +142,7 @@ pub struct ConfigObjImpl { pub push_down_pull_up_split: bool, pub stream_mode: bool, pub non_streaming_query_max_row_limit: i32, + pub fail_on_limitless_post_processing: bool, pub cube_scan_max_batch_rows: usize, pub max_sessions: usize, pub no_implicit_order: bool, @@ -201,6 +204,10 @@ impl ConfigObjImpl { .unwrap_or(sql_push_down), stream_mode: env_parse("CUBESQL_STREAM_MODE", false), non_streaming_query_max_row_limit, + fail_on_limitless_post_processing: env_parse( + "CUBESQL_FAIL_ON_LIMITLESS_POST_PROCESSING", + false, + ), cube_scan_max_batch_rows: env_parse("CUBESQL_CUBE_SCAN_MAX_BATCH_ROWS", 65536), max_sessions: env_parse("CUBEJS_MAX_SESSIONS", 1024), no_implicit_order: env_parse("CUBESQL_SQL_NO_IMPLICIT_ORDER", true), @@ -268,6 +275,10 @@ impl ConfigObj for ConfigObjImpl { self.non_streaming_query_max_row_limit } + fn fail_on_limitless_post_processing(&self) -> bool { + self.fail_on_limitless_post_processing + } + fn cube_scan_max_batch_rows(&self) -> usize { self.cube_scan_max_batch_rows } @@ -314,6 +325,7 @@ impl Config { push_down_pull_up_split: true, stream_mode: false, non_streaming_query_max_row_limit: 50000, + fail_on_limitless_post_processing: false, cube_scan_max_batch_rows: 65536, max_sessions: 1024, no_implicit_order: true,