Skip to content

Commit

Permalink
feat(cubesql): SQL push down support for IS NULL and IS NOT NULL …
Browse files Browse the repository at this point in the history 
…expressions
  • Loading branch information
@paveltiunov
paveltiunov committed Nov 1, 2023
1 parent 0f8de97 commit 9b3c27d
Show file tree
Hide file tree
Showing 7 changed files with 263 additions and 109 deletions.
221 changes: 114 additions & 107 deletions rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -713,6 +713,8 @@ impl CubeScanWrapperNode {
ungrouped_scan_node.clone(),
)
.await?;
let expr_sql =
Self::escape_interpolation_quotes(expr_sql, ungrouped_scan_node.is_some());
sql = new_sql_query;

let original_alias = expr_name(&original_expr, &schema)?;
Expand Down Expand Up @@ -889,27 +891,62 @@ impl CubeScanWrapperNode {
ungrouped_scan_node.clone(),
)
.await?;
let resulting_sql = Self::escape_interpolation_quotes(
sql_generator
.get_sql_templates()
.binary_expr(left, op.to_string(), right)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for binary expr: {}",
e
))
})?,
ungrouped_scan_node.is_some(),
);
let resulting_sql = sql_generator
.get_sql_templates()
.binary_expr(left, op.to_string(), right)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for binary expr: {}",
e
))
})?;
Ok((resulting_sql, sql_query))
}
// Expr::AnyExpr { .. } => {}
// Expr::Like(_) => {}-=
// Expr::ILike(_) => {}
// Expr::SimilarTo(_) => {}
// Expr::Not(_) => {}
// Expr::IsNotNull(_) => {}
// Expr::IsNull(_) => {}
Expr::IsNotNull(expr) => {
let (expr, sql_query) = Self::generate_sql_for_expr(
plan.clone(),
sql_query,
sql_generator.clone(),
*expr,
ungrouped_scan_node.clone(),
)
.await?;
let resulting_sql = sql_generator
.get_sql_templates()
.is_null_expr(expr, true)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for is not null expr: {}",
e
))
})?;
Ok((resulting_sql, sql_query))
}
Expr::IsNull(expr) => {
let (expr, sql_query) = Self::generate_sql_for_expr(
plan.clone(),
sql_query,
sql_generator.clone(),
*expr,
ungrouped_scan_node.clone(),
)
.await?;
let resulting_sql = sql_generator
.get_sql_templates()
.is_null_expr(expr, false)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for is null expr: {}",
e
))
})?;
Ok((resulting_sql, sql_query))
}
// Expr::Negative(_) => {}
// Expr::GetIndexedField { .. } => {}
// Expr::Between { .. } => {}
Expand Down Expand Up @@ -967,18 +1004,12 @@ impl CubeScanWrapperNode {
} else {
None
};
let resulting_sql = Self::escape_interpolation_quotes(
sql_generator
.get_sql_templates()
.case(expr, when_then_expr_sql, else_expr)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for case: {}",
e
))
})?,
ungrouped_scan_node.is_some(),
);
let resulting_sql = sql_generator
.get_sql_templates()
.case(expr, when_then_expr_sql, else_expr)
.map_err(|e| {
DataFusionError::Internal(format!("Can't generate SQL for case: {}", e))
})?;
Ok((resulting_sql, sql_query))
}
Expr::Cast { expr, data_type } => {
Expand Down Expand Up @@ -1022,18 +1053,12 @@ impl CubeScanWrapperNode {
)));
}
};
let resulting_sql = Self::escape_interpolation_quotes(
sql_generator
.get_sql_templates()
.cast_expr(expr, data_type.to_string())
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for cast: {}",
e
))
})?,
ungrouped_scan_node.is_some(),
);
let resulting_sql = sql_generator
.get_sql_templates()
.cast_expr(expr, data_type.to_string())
.map_err(|e| {
DataFusionError::Internal(format!("Can't generate SQL for cast: {}", e))
})?;
Ok((resulting_sql, sql_query))
}
// Expr::TryCast { .. } => {}
Expand All @@ -1050,18 +1075,15 @@ impl CubeScanWrapperNode {
ungrouped_scan_node.clone(),
)
.await?;
let resulting_sql = Self::escape_interpolation_quotes(
sql_generator
.get_sql_templates()
.sort_expr(expr, asc, nulls_first)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for sort expr: {}",
e
))
})?,
ungrouped_scan_node.is_some(),
);
let resulting_sql = sql_generator
.get_sql_templates()
.sort_expr(expr, asc, nulls_first)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for sort expr: {}",
e
))
})?;
Ok((resulting_sql, sql_query))
}

Expand Down Expand Up @@ -1142,18 +1164,15 @@ impl CubeScanWrapperNode {
};
let interval = format!("{} {}", num, date_part);
(
Self::escape_interpolation_quotes(
sql_generator
.get_sql_templates()
.interval_expr(interval, num, date_part.to_string())
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for interval: {}",
e
))
})?,
ungrouped_scan_node.is_some(),
),
sql_generator
.get_sql_templates()
.interval_expr(interval, num, date_part.to_string())
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for interval: {}",
e
))
})?,
sql_query,
)
} else {
Expand Down Expand Up @@ -1185,18 +1204,15 @@ impl CubeScanWrapperNode {
sql_args.push(sql);
}
Ok((
Self::escape_interpolation_quotes(
sql_generator
.get_sql_templates()
.scalar_function(fun.name.to_string(), sql_args, None)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for scalar function: {}",
e
))
})?,
ungrouped_scan_node.is_some(),
),
sql_generator
.get_sql_templates()
.scalar_function(fun.name.to_string(), sql_args, None)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for scalar function: {}",
e
))
})?,
sql_query,
))
}
Expand All @@ -1221,18 +1237,15 @@ impl CubeScanWrapperNode {
)
.await?;
return Ok((
Self::escape_interpolation_quotes(
sql_generator
.get_sql_templates()
.extract_expr(date_part.to_string(), arg_sql)
.map_err(|e| {
DataFusionError::Internal(format!(
sql_generator
.get_sql_templates()
.extract_expr(date_part.to_string(), arg_sql)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for scalar function: {}",
e
))
})?,
ungrouped_scan_node.is_some(),
),
})?,
query,
));
}
Expand Down Expand Up @@ -1269,18 +1282,15 @@ impl CubeScanWrapperNode {
sql_args.push(sql);
}
Ok((
Self::escape_interpolation_quotes(
sql_generator
.get_sql_templates()
.scalar_function(fun.to_string(), sql_args, date_part)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for scalar function: {}",
e
))
})?,
ungrouped_scan_node.is_some(),
),
sql_generator
.get_sql_templates()
.scalar_function(fun.to_string(), sql_args, date_part)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for scalar function: {}",
e
))
})?,
sql_query,
))
}
Expand Down Expand Up @@ -1311,18 +1321,15 @@ impl CubeScanWrapperNode {
sql_args.push(sql);
}
Ok((
Self::escape_interpolation_quotes(
sql_generator
.get_sql_templates()
.aggregate_function(fun, sql_args, distinct)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for aggregate function: {}",
e
))
})?,
ungrouped_scan_node.is_some(),
),
sql_generator
.get_sql_templates()
.aggregate_function(fun, sql_args, distinct)
.map_err(|e| {
DataFusionError::Internal(format!(
"Can't generate SQL for aggregate function: {}",
e
))
})?,
sql_query,
))
}
Expand Down
29 changes: 29 additions & 0 deletions rust/cubesql/cubesql/src/compile/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18897,6 +18897,35 @@ ORDER BY \"COUNT(count)\" DESC"
);
}

#[tokio::test]
async fn test_case_wrapper_with_null() {
if !Rewriter::sql_push_down_enabled() {
return;
}
init_logger();

let query_plan = convert_select_to_query_plan(
"SELECT CASE WHEN taxful_total_price IS NULL THEN NULL WHEN taxful_total_price < taxful_total_price * 2 THEN taxful_total_price END FROM KibanaSampleDataEcommerce GROUP BY 1"
.to_string(),
DatabaseProtocol::PostgreSQL,
)
.await;

let logical_plan = query_plan.as_logical_plan();
assert!(logical_plan
.find_cube_scan_wrapper()
.wrapped_sql
.unwrap()
.sql
.contains("CASE WHEN"));

let physical_plan = query_plan.as_physical_plan().await.unwrap();
println!(
"Physical plan: {}",
displayable(physical_plan.as_ref()).indent()
);
}

#[tokio::test]
async fn test_case_wrapper_ungrouped_on_dimension() {
if !Rewriter::sql_push_down_enabled() {
Expand Down
2 changes: 1 addition & 1 deletion rust/cubesql/cubesql/src/compile/rewrite/cost.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,11 @@ pub struct BestCubePlan;
pub struct CubePlanCost {
replacers: i64,
table_scans: i64,
empty_wrappers: i64,
non_detected_cube_scans: i64,
filters: i64,
structure_points: i64,
filter_members: i64,
empty_wrappers: i64,
member_errors: i64,
wrapper_nodes: i64,
ast_size_outside_wrapper: usize,
Expand Down
Loading

0 comments on commit 9b3c27d

Please sign in to comment.