Skip to content

Commit 9b3c27d

Browse files
paveltiunovpaveltiunov
committed
feat(cubesql): SQL push down support for IS NULL and IS NOT NULL expressions
1 parent 0f8de97 commit 9b3c27d

File tree

7 files changed

+263
-109
lines changed

7 files changed

+263
-109
lines changed

‎rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs‎

Lines changed: 114 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -713,6 +713,8 @@ impl CubeScanWrapperNode {
713713
ungrouped_scan_node.clone(),
714714
)
715715
.await?;
716+
let expr_sql =
717+
Self::escape_interpolation_quotes(expr_sql, ungrouped_scan_node.is_some());
716718
sql = new_sql_query;
717719

718720
let original_alias = expr_name(&original_expr, &schema)?;
@@ -889,27 +891,62 @@ impl CubeScanWrapperNode {
889891
ungrouped_scan_node.clone(),
890892
)
891893
.await?;
892-
let resulting_sql = Self::escape_interpolation_quotes(
893-
sql_generator
894-
.get_sql_templates()
895-
.binary_expr(left, op.to_string(), right)
896-
.map_err(|e| {
897-
DataFusionError::Internal(format!(
898-
"Can't generate SQL for binary expr: {}",
899-
e
900-
))
901-
})?,
902-
ungrouped_scan_node.is_some(),
903-
);
894+
let resulting_sql = sql_generator
895+
.get_sql_templates()
896+
.binary_expr(left, op.to_string(), right)
897+
.map_err(|e| {
898+
DataFusionError::Internal(format!(
899+
"Can't generate SQL for binary expr: {}",
900+
e
901+
))
902+
})?;
904903
Ok((resulting_sql, sql_query))
905904
}
906905
// Expr::AnyExpr { .. } => {}
907906
// Expr::Like(_) => {}-=
908907
// Expr::ILike(_) => {}
909908
// Expr::SimilarTo(_) => {}
910909
// Expr::Not(_) => {}
911-
// Expr::IsNotNull(_) => {}
912-
// Expr::IsNull(_) => {}
910+
Expr::IsNotNull(expr) => {
911+
let (expr, sql_query) = Self::generate_sql_for_expr(
912+
plan.clone(),
913+
sql_query,
914+
sql_generator.clone(),
915+
*expr,
916+
ungrouped_scan_node.clone(),
917+
)
918+
.await?;
919+
let resulting_sql = sql_generator
920+
.get_sql_templates()
921+
.is_null_expr(expr, true)
922+
.map_err(|e| {
923+
DataFusionError::Internal(format!(
924+
"Can't generate SQL for is not null expr: {}",
925+
e
926+
))
927+
})?;
928+
Ok((resulting_sql, sql_query))
929+
}
930+
Expr::IsNull(expr) => {
931+
let (expr, sql_query) = Self::generate_sql_for_expr(
932+
plan.clone(),
933+
sql_query,
934+
sql_generator.clone(),
935+
*expr,
936+
ungrouped_scan_node.clone(),
937+
)
938+
.await?;
939+
let resulting_sql = sql_generator
940+
.get_sql_templates()
941+
.is_null_expr(expr, false)
942+
.map_err(|e| {
943+
DataFusionError::Internal(format!(
944+
"Can't generate SQL for is null expr: {}",
945+
e
946+
))
947+
})?;
948+
Ok((resulting_sql, sql_query))
949+
}
913950
// Expr::Negative(_) => {}
914951
// Expr::GetIndexedField { .. } => {}
915952
// Expr::Between { .. } => {}
@@ -967,18 +1004,12 @@ impl CubeScanWrapperNode {
9671004
} else {
9681005
None
9691006
};
970-
let resulting_sql = Self::escape_interpolation_quotes(
971-
sql_generator
972-
.get_sql_templates()
973-
.case(expr, when_then_expr_sql, else_expr)
974-
.map_err(|e| {
975-
DataFusionError::Internal(format!(
976-
"Can't generate SQL for case: {}",
977-
e
978-
))
979-
})?,
980-
ungrouped_scan_node.is_some(),
981-
);
1007+
let resulting_sql = sql_generator
1008+
.get_sql_templates()
1009+
.case(expr, when_then_expr_sql, else_expr)
1010+
.map_err(|e| {
1011+
DataFusionError::Internal(format!("Can't generate SQL for case: {}", e))
1012+
})?;
9821013
Ok((resulting_sql, sql_query))
9831014
}
9841015
Expr::Cast { expr, data_type } => {
@@ -1022,18 +1053,12 @@ impl CubeScanWrapperNode {
10221053
)));
10231054
}
10241055
};
1025-
let resulting_sql = Self::escape_interpolation_quotes(
1026-
sql_generator
1027-
.get_sql_templates()
1028-
.cast_expr(expr, data_type.to_string())
1029-
.map_err(|e| {
1030-
DataFusionError::Internal(format!(
1031-
"Can't generate SQL for cast: {}",
1032-
e
1033-
))
1034-
})?,
1035-
ungrouped_scan_node.is_some(),
1036-
);
1056+
let resulting_sql = sql_generator
1057+
.get_sql_templates()
1058+
.cast_expr(expr, data_type.to_string())
1059+
.map_err(|e| {
1060+
DataFusionError::Internal(format!("Can't generate SQL for cast: {}", e))
1061+
})?;
10371062
Ok((resulting_sql, sql_query))
10381063
}
10391064
// Expr::TryCast { .. } => {}
@@ -1050,18 +1075,15 @@ impl CubeScanWrapperNode {
10501075
ungrouped_scan_node.clone(),
10511076
)
10521077
.await?;
1053-
let resulting_sql = Self::escape_interpolation_quotes(
1054-
sql_generator
1055-
.get_sql_templates()
1056-
.sort_expr(expr, asc, nulls_first)
1057-
.map_err(|e| {
1058-
DataFusionError::Internal(format!(
1059-
"Can't generate SQL for sort expr: {}",
1060-
e
1061-
))
1062-
})?,
1063-
ungrouped_scan_node.is_some(),
1064-
);
1078+
let resulting_sql = sql_generator
1079+
.get_sql_templates()
1080+
.sort_expr(expr, asc, nulls_first)
1081+
.map_err(|e| {
1082+
DataFusionError::Internal(format!(
1083+
"Can't generate SQL for sort expr: {}",
1084+
e
1085+
))
1086+
})?;
10651087
Ok((resulting_sql, sql_query))
10661088
}
10671089

@@ -1142,18 +1164,15 @@ impl CubeScanWrapperNode {
11421164
};
11431165
let interval = format!("{} {}", num, date_part);
11441166
(
1145-
Self::escape_interpolation_quotes(
1146-
sql_generator
1147-
.get_sql_templates()
1148-
.interval_expr(interval, num, date_part.to_string())
1149-
.map_err(|e| {
1150-
DataFusionError::Internal(format!(
1151-
"Can't generate SQL for interval: {}",
1152-
e
1153-
))
1154-
})?,
1155-
ungrouped_scan_node.is_some(),
1156-
),
1167+
sql_generator
1168+
.get_sql_templates()
1169+
.interval_expr(interval, num, date_part.to_string())
1170+
.map_err(|e| {
1171+
DataFusionError::Internal(format!(
1172+
"Can't generate SQL for interval: {}",
1173+
e
1174+
))
1175+
})?,
11571176
sql_query,
11581177
)
11591178
} else {
@@ -1185,18 +1204,15 @@ impl CubeScanWrapperNode {
11851204
sql_args.push(sql);
11861205
}
11871206
Ok((
1188-
Self::escape_interpolation_quotes(
1189-
sql_generator
1190-
.get_sql_templates()
1191-
.scalar_function(fun.name.to_string(), sql_args, None)
1192-
.map_err(|e| {
1193-
DataFusionError::Internal(format!(
1194-
"Can't generate SQL for scalar function: {}",
1195-
e
1196-
))
1197-
})?,
1198-
ungrouped_scan_node.is_some(),
1199-
),
1207+
sql_generator
1208+
.get_sql_templates()
1209+
.scalar_function(fun.name.to_string(), sql_args, None)
1210+
.map_err(|e| {
1211+
DataFusionError::Internal(format!(
1212+
"Can't generate SQL for scalar function: {}",
1213+
e
1214+
))
1215+
})?,
12001216
sql_query,
12011217
))
12021218
}
@@ -1221,18 +1237,15 @@ impl CubeScanWrapperNode {
12211237
)
12221238
.await?;
12231239
return Ok((
1224-
Self::escape_interpolation_quotes(
1225-
sql_generator
1226-
.get_sql_templates()
1227-
.extract_expr(date_part.to_string(), arg_sql)
1228-
.map_err(|e| {
1229-
DataFusionError::Internal(format!(
1240+
sql_generator
1241+
.get_sql_templates()
1242+
.extract_expr(date_part.to_string(), arg_sql)
1243+
.map_err(|e| {
1244+
DataFusionError::Internal(format!(
12301245
"Can't generate SQL for scalar function: {}",
12311246
e
12321247
))
1233-
})?,
1234-
ungrouped_scan_node.is_some(),
1235-
),
1248+
})?,
12361249
query,
12371250
));
12381251
}
@@ -1269,18 +1282,15 @@ impl CubeScanWrapperNode {
12691282
sql_args.push(sql);
12701283
}
12711284
Ok((
1272-
Self::escape_interpolation_quotes(
1273-
sql_generator
1274-
.get_sql_templates()
1275-
.scalar_function(fun.to_string(), sql_args, date_part)
1276-
.map_err(|e| {
1277-
DataFusionError::Internal(format!(
1278-
"Can't generate SQL for scalar function: {}",
1279-
e
1280-
))
1281-
})?,
1282-
ungrouped_scan_node.is_some(),
1283-
),
1285+
sql_generator
1286+
.get_sql_templates()
1287+
.scalar_function(fun.to_string(), sql_args, date_part)
1288+
.map_err(|e| {
1289+
DataFusionError::Internal(format!(
1290+
"Can't generate SQL for scalar function: {}",
1291+
e
1292+
))
1293+
})?,
12841294
sql_query,
12851295
))
12861296
}
@@ -1311,18 +1321,15 @@ impl CubeScanWrapperNode {
13111321
sql_args.push(sql);
13121322
}
13131323
Ok((
1314-
Self::escape_interpolation_quotes(
1315-
sql_generator
1316-
.get_sql_templates()
1317-
.aggregate_function(fun, sql_args, distinct)
1318-
.map_err(|e| {
1319-
DataFusionError::Internal(format!(
1320-
"Can't generate SQL for aggregate function: {}",
1321-
e
1322-
))
1323-
})?,
1324-
ungrouped_scan_node.is_some(),
1325-
),
1324+
sql_generator
1325+
.get_sql_templates()
1326+
.aggregate_function(fun, sql_args, distinct)
1327+
.map_err(|e| {
1328+
DataFusionError::Internal(format!(
1329+
"Can't generate SQL for aggregate function: {}",
1330+
e
1331+
))
1332+
})?,
13261333
sql_query,
13271334
))
13281335
}

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18897,6 +18897,35 @@ ORDER BY \"COUNT(count)\" DESC"
1889718897
);
1889818898
}
1889918899

18900+
#[tokio::test]
18901+
async fn test_case_wrapper_with_null() {
18902+
if !Rewriter::sql_push_down_enabled() {
18903+
return;
18904+
}
18905+
init_logger();
18906+
18907+
let query_plan = convert_select_to_query_plan(
18908+
"SELECT CASE WHEN taxful_total_price IS NULL THEN NULL WHEN taxful_total_price < taxful_total_price * 2 THEN taxful_total_price END FROM KibanaSampleDataEcommerce GROUP BY 1"
18909+
.to_string(),
18910+
DatabaseProtocol::PostgreSQL,
18911+
)
18912+
.await;
18913+
18914+
let logical_plan = query_plan.as_logical_plan();
18915+
assert!(logical_plan
18916+
.find_cube_scan_wrapper()
18917+
.wrapped_sql
18918+
.unwrap()
18919+
.sql
18920+
.contains("CASE WHEN"));
18921+
18922+
let physical_plan = query_plan.as_physical_plan().await.unwrap();
18923+
println!(
18924+
"Physical plan: {}",
18925+
displayable(physical_plan.as_ref()).indent()
18926+
);
18927+
}
18928+
1890018929
#[tokio::test]
1890118930
async fn test_case_wrapper_ungrouped_on_dimension() {
1890218931
if !Rewriter::sql_push_down_enabled() {

‎rust/cubesql/cubesql/src/compile/rewrite/cost.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,11 @@ pub struct BestCubePlan;
1818
pub struct CubePlanCost {
1919
replacers: i64,
2020
table_scans: i64,
21+
empty_wrappers: i64,
2122
non_detected_cube_scans: i64,
2223
filters: i64,
2324
structure_points: i64,
2425
filter_members: i64,
25-
empty_wrappers: i64,
2626
member_errors: i64,
2727
wrapper_nodes: i64,
2828
ast_size_outside_wrapper: usize,

0 commit comments

Comments
 (0)