diff --git a/docs-mintlify/reference/configuration/environment-variables.mdx b/docs-mintlify/reference/configuration/environment-variables.mdx index b7b2053a04d0c..979325bf9ef2f 100644 --- a/docs-mintlify/reference/configuration/environment-variables.mdx +++ b/docs-mintlify/reference/configuration/environment-variables.mdx @@ -1563,6 +1563,28 @@ 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 with no `LIMIT` clause, or with one above the maximum +row limit set by +[`CUBESQL_NON_STREAMING_QUERY_MAX_ROW_LIMIT`](#cubesql_non_streaming_query_max_row_limit). + +Such a query is capped at that limit, without any ordering, and a `LIMIT` above it is +clamped back down to it. 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 cd3e3389daaa6..19e2d39b878fb 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/rust/cubesql/cubesql/src/compile/mod.rs @@ -59,7 +59,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] @@ -2184,33 +2184,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..362be5f7699ea 100644 --- a/rust/cubesql/cubesql/src/compile/rewrite/cost.rs +++ b/rust/cubesql/cubesql/src/compile/rewrite/cost.rs @@ -1,29 +1,49 @@ 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}, }; use egg::{Analysis, EGraph, Id, Language, RecExpr}; use indexmap::IndexSet; +/// Whether a query carrying this limit still gets truncated by the row cap. +/// +/// No limit is capped at `max_row_limit`, and so is any limit above it, since +/// [`CubeScanExecutionPlan::execute`] clamps the request back down to the cap before +/// sending it. +fn is_limitless(limit: &Option, max_row_limit: usize) -> bool { + limit.is_none_or(|limit| limit > max_row_limit) +} + #[derive(Debug)] pub struct BestCubePlan { meta_context: Arc, penalize_post_processing: bool, + penalize_limitless_post_processing: bool, + max_row_limit: usize, } 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, + max_row_limit: usize, + ) -> Self { Self { meta_context, penalize_post_processing, + penalize_limitless_post_processing, + max_row_limit, } } @@ -209,6 +229,23 @@ impl BestCubePlan { _ => 0, }; + // A Cube query 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. + // + // A limit of its own only helps while it is at or below the cap, since anything + // above is clamped back down to it and truncates just the same. BI tools often send + // a large defensive limit, so those queries are read as unlimited here. + let limitless_scans = match enode { + LogicalPlanLanguage::CubeScanLimit(CubeScanLimit(limit)) + if is_limitless(limit, self.max_row_limit) => + { + 1 + } + _ => 0, + }; + CubePlanCost { replacers: this_replacers, // Will be filled in finalize @@ -229,6 +266,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 +280,46 @@ impl BestCubePlan { ast_size: 1, ungrouped_nodes, unwrapped_subqueries, + limitless_scans: Unordered(limitless_scans), + // Will be filled in finalize + limitless_ungrouped_aggregates: Unordered(0), } } } +/// 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 +333,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 +355,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 +377,11 @@ pub struct CubePlanCost { ast_size: usize, ast_size_inside_wrapper: usize, ungrouped_nodes: usize, + // Input for `limitless_post_processing`, which is what expresses the preference + limitless_scans: Unordered, + // Reported rather than preferred: extraction prices these through + // `ungrouped_aggregates`, this only records the ones that read a truncated scan + limitless_ungrouped_aggregates: Unordered, } #[derive(Debug, Clone, Eq, Hash, PartialEq)] @@ -318,6 +399,20 @@ pub enum SortState { } impl CubePlanCost { + /// Whether this plan computes something in DataFusion over a Cube query that the row + /// cap truncates, and so would answer with a number that is wrong rather than short. + /// + /// Two shapes qualify. `limitless_post_processing` counts the post processing that + /// extraction already prefers to push down. An `Aggregate` over an ungrouped scan is + /// the other: extraction prices it through `ungrouped_aggregates` instead, but the + /// rows it reads are raw and capped, so its result is just as wrong. + /// + /// Both are counted at the node that does the reading, so neither can pair an + /// aggregate in one branch with an unlimited scan in another. + pub fn truncates_post_processing(&self) -> bool { + self.limitless_post_processing > 0 || self.limitless_ungrouped_aggregates.0 > 0 + } + pub fn add_child(&self, other: &Self) -> Self { Self { replacers: self.replacers + other.replacers, @@ -366,6 +461,12 @@ 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), + limitless_ungrouped_aggregates: Unordered( + self.limitless_ungrouped_aggregates.0 + other.limitless_ungrouped_aggregates.0, + ), } } @@ -373,6 +474,7 @@ impl CubePlanCost { &self, state: &CubePlanState, sort_state: &SortState, + under_limit: bool, enode: &LogicalPlanLanguage, options: CubePlanCostOptions, ) -> Self { @@ -444,6 +546,48 @@ 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. That holds + // while its grouping matches the scan's, where the cap lands on the rows the + // client asked for rather than on hidden detail. A coarser rollup, such as + // `SUM(cnt)` over a scan grouped by city, does read a capped slice and is + // wrong - it is exempt here because pushing harder proved counterproductive, + // full push down being ungrouped itself, not because it is safe + // + // 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 > 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 +597,28 @@ impl CubePlanCost { ast_size: self.ast_size, ast_size_inside_wrapper: self.ast_size_inside_wrapper, ungrouped_nodes: self.ungrouped_nodes, + // A limit above bounds every scan below it, so nothing there is left to be + // truncated by the row limit + limitless_scans: if under_limit { + Unordered(0) + } else { + self.limitless_scans + }, + // Counted at the `Aggregate` itself, so the scan it reads is the one in its + // own subtree rather than any unlimited scan elsewhere in the plan + limitless_ungrouped_aggregates: Unordered( + match state { + CubePlanState::Unwrapped(_) + if self.ungrouped_nodes > 0 && self.limitless_scans.0 > 0 => + { + match enode { + LogicalPlanLanguage::Aggregate(_) => 1, + _ => 0, + } + } + _ => 0, + } + self.limitless_ungrouped_aggregates.0, + ), } } } @@ -742,16 +908,54 @@ impl TopDownCost for CubePlanCost { pub struct CubePlanTopDownState { wrapped: CubePlanState, limit: SortState, + /// Whether a select above bounds the rows every scan below it can return + under_limit: bool, + max_row_limit: usize, } impl CubePlanTopDownState { - pub fn new() -> Self { + pub fn new(max_row_limit: usize) -> Self { Self { wrapped: CubePlanState::Unwrapped(0), limit: SortState::None, + under_limit: false, + max_row_limit, } } + /// Whether this node bounds the rows every scan below it can return. + /// + /// A push to Cube wrapper carries the user's limit on its select rather than on the + /// scan below it, so the scan itself still looks unlimited. Reading the limit here and + /// carrying it down the plan means a scan is judged by the query it actually ends up + /// in, and a limited scan can never stand in for an unlimited one in a sibling branch. + pub fn introduces_limit( + &self, + node: &LogicalPlanLanguage, + egraph: &EGraph, + ) -> bool + where + A: Analysis, + { + let limit_id = match node { + LogicalPlanLanguage::WrappedSelect(params) => params[10], + LogicalPlanLanguage::CubeScan(params) => params[4], + _ => return false, + }; + // Anything but a single limit at or below the row cap leaves the scans below + // unbounded, which is the safe reading: the penalty applies rather than being + // silently skipped + let nodes = &egraph[limit_id].nodes; + !nodes.is_empty() + && nodes.iter().all(|node| match node { + LogicalPlanLanguage::WrappedSelectLimit(WrappedSelectLimit(limit)) + | LogicalPlanLanguage::CubeScanLimit(CubeScanLimit(limit)) => { + !is_limitless(limit, self.max_row_limit) + } + _ => false, + }) + } + pub fn is_wrapped( &self, node: &LogicalPlanLanguage, @@ -818,7 +1022,14 @@ impl TopDownState for CubePlanTopDownState { _ => SortState::None, }; - Self { wrapped, limit } + let under_limit = self.under_limit || self.introduces_limit(node, egraph); + + Self { + wrapped, + limit, + under_limit, + max_row_limit: self.max_row_limit, + } } } @@ -837,9 +1048,11 @@ 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 max_row_limit = config_obj.non_streaming_query_max_row_limit().max(0) as usize; + 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, truncates_post_processing) = tokio::task::spawn_blocking(move || { let (runner, qtrace_egraph_iterations) = Self::run_rewrites(&cube_context, egraph, rules, "final")?; @@ -362,8 +371,13 @@ 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), - CubePlanTopDownState::new(), + BestCubePlan::new( + cube_context.meta.clone(), + penalize_post_processing, + penalize_limitless_post_processing, + max_row_limit, + ), + CubePlanTopDownState::new(max_row_limit), ); let Some((best_cost, best)) = extractor.find_best(root) else { return Err(CubeError::rewrite("Unable to find best plan".to_string())); @@ -388,6 +402,7 @@ impl Rewriter { converter.to_logical_plan(new_root), qtrace_egraph_iterations, qtrace_best_graph, + best_cost.truncates_post_processing(), )) }) .await??; @@ -397,6 +412,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 && truncates_post_processing { + return Err(CubeError::user( + "Query requires post-processing of a Cube query with no LIMIT, or with one \ + above the maximum row limit, so it would be truncated to that limit and \ + produce incorrect results. Add a LIMIT at or below it, 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 edd4aa3ff37fe..53df62aea1afc 100644 --- a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs @@ -3084,3 +3084,316 @@ async fn test_wrapper_date_add_negative_and_out_of_range_intervals() { ); } } + +/// 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("CUBESQL_FAIL_ON_LIMITLESS_POST_PROCESSING"), + "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() + } + ); +} + +/// A limit bounds the rows of the query it sits on, not of a query beside it. Summing +/// limits and limitless scans over the whole plan would let the limited branch here stand +/// in for the unlimited one, and the union would read a truncated half without saying so. +#[tokio::test] +async fn test_wrapper_limitless_post_processing_sibling_without_limit() { + 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; + + let error = context + .convert_sql_to_cube_query( + "(SELECT customer_gender AS g FROM KibanaSampleDataEcommerce GROUP BY 1 LIMIT 10) \ + UNION ALL \ + (SELECT customer_gender AS g FROM KibanaSampleDataEcommerce GROUP BY 1)", + ) + .await + .expect_err("union with an unlimited branch should fail"); + assert!( + error + .to_string() + .contains("CUBESQL_FAIL_ON_LIMITLESS_POST_PROCESSING"), + "unexpected error: {}", + error + ); +} + +/// A limit only bounds the query while it is at or below the row cap. Anything above is +/// clamped back down to the cap before the request is sent, so it truncates exactly like +/// no limit at all - and a large defensive limit is a shape BI tools emit routinely. +#[tokio::test] +async fn test_wrapper_limitless_post_processing_limit_above_row_cap() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + // Matches `non_streaming_query_max_row_limit` for the default config under test + let max_row_limit = ConfigObjImpl::default().non_streaming_query_max_row_limit as i64; + + for (limit, bounded) in [ + (10, true), + (max_row_limit, true), + (max_row_limit + 1, false), + (max_row_limit * 20, false), + ] { + let mut config = ConfigObjImpl::default(); + config.fail_on_limitless_post_processing = true; + let context = + TestContext::with_config(DatabaseProtocol::PostgreSQL, Arc::new(config)).await; + + // `ROUND` has no template here, so the filter cannot leave post processing + let result = context + .convert_sql_to_cube_query(&format!( + "SELECT COUNT(*) \ + FROM ( \ + SELECT customer_gender AS gender, MIN(taxful_total_price) AS cheapest \ + FROM KibanaSampleDataEcommerce \ + GROUP BY 1 \ + LIMIT {limit} \ + ) first_orders \ + WHERE ROUND(cheapest) > 10" + )) + .await; + + if bounded { + result.unwrap_or_else(|error| { + panic!( + "LIMIT {} is within the row cap and should compile: {}", + limit, error + ) + }); + } else { + let error = result.err().unwrap_or_else(|| { + panic!("LIMIT {} is clamped to the row cap and should fail", limit) + }); + assert!( + error + .to_string() + .contains("CUBESQL_FAIL_ON_LIMITLESS_POST_PROCESSING"), + "unexpected error for LIMIT {}: {}", + limit, + error + ); + } + } +} + +/// An `Aggregate` over an ungrouped scan is priced by `ungrouped_aggregates` rather than +/// by `limitless_post_processing`, so extraction already accounts for it. The rows it +/// reads are still raw and capped, though, so the operator who turned the check on to be +/// told about truncated results has to hear about this one too. +#[tokio::test] +async fn test_wrapper_limitless_post_processing_ungrouped_aggregate() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + // The day of quarter grouping cannot be pushed down, so the aggregate stays in post + // processing over an ungrouped scan, and nothing above it drops or reorders rows + let query = r#" + SELECT + CAST("inner_query"."order_date" AS date) + - CAST("inner_query"."quarter_start" AS date) + + 1 AS "day_of_quarter", + MEASURE("inner_query"."sumPrice") AS "revenue" + FROM ( + SELECT + "ta_1"."order_date" AS "order_date", + CAST( + EXTRACT(YEAR FROM "ta_1"."order_date") || '-' + || EXTRACT(MONTH FROM "ta_1"."order_date") || '-01' + AS DATE) + + (((MOD(CAST((EXTRACT(MONTH FROM "ta_1"."order_date") - 1) + AS numeric), 3) + 1) - 1) * -1) + * INTERVAL '1 month' + AS "quarter_start", + CASE WHEN "ta_1"."customer_gender" = 'female' + THEN "ta_1"."sumPrice" END AS "sumPrice" + FROM "db"."public"."KibanaSampleDataEcommerce" AS "ta_1" + ) "inner_query" + GROUP BY 1 + "#; + + // Without the check the query still runs, reading a capped slice of raw rows + let logical_plan = + convert_select_to_query_plan(query.to_string(), DatabaseProtocol::PostgreSQL) + .await + .as_logical_plan(); + let request = logical_plan.find_cube_scan().request; + assert_eq!(request.ungrouped, Some(true)); + assert_eq!(request.limit, None); + + let mut config = ConfigObjImpl::default(); + config.fail_on_limitless_post_processing = true; + let context = TestContext::with_config(DatabaseProtocol::PostgreSQL, Arc::new(config)).await; + let error = context + .convert_sql_to_cube_query(query) + .await + .expect_err("aggregate over an unlimited ungrouped scan should fail"); + assert!( + error + .to_string() + .contains("CUBESQL_FAIL_ON_LIMITLESS_POST_PROCESSING"), + "unexpected error: {}", + error + ); +} 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,