adriangb commented on code in PR #19321:
URL: https://github.com/apache/datafusion/pull/19321#discussion_r2644611809
##########
datafusion/expr/src/expr.rs:
##########
@@ -3929,6 +4056,366 @@ mod test {
}
}
+ #[test]
+ fn infer_placeholder_simple_case() {
Review Comment:
```suggestion
fn infer_case_expr_placeholder_simple() {
```
The name was a bit confusing to me
##########
datafusion/expr/src/expr.rs:
##########
@@ -3929,6 +4056,366 @@ mod test {
}
}
+ #[test]
+ fn infer_placeholder_simple_case() {
+ // CASE department_id WHEN $1 THEN $2 WHEN $3 THEN 'Engineering' ELSE
$4 END
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("department_id", DataType::Int32, true),
+ Field::new("name", DataType::Utf8, true),
+ ]));
+ let df_schema = DFSchema::try_from(schema).unwrap();
+
+ let case_expr = case(col("department_id"))
+ .when(placeholder("$1"), placeholder("$2"))
+ .when(placeholder("$3"), lit("Engineering"))
+ .otherwise(placeholder("$4"))
+ .unwrap();
+
+ let (inferred_expr, contains_placeholder) =
+ case_expr.infer_placeholder_types(&df_schema).unwrap();
+
+ assert!(contains_placeholder);
+
+ match inferred_expr {
+ Expr::Case(Case {
+ expr,
+ when_then_expr,
+ else_expr,
+ }) => {
+ // Base expression should remain unchanged
+ assert!(matches!(**expr.as_ref().unwrap(), Expr::Column(_)));
+
+ // WHEN expressions should infer Int32 (to match department_id
column)
+ for (when_expr, then_expr) in &when_then_expr {
+ if let Expr::Placeholder(placeholder) = when_expr.as_ref()
{
+ assert_eq!(
+ placeholder.field.as_ref().unwrap().data_type(),
+ &DataType::Int32,
+ "WHEN placeholder {} should infer Int32",
+ placeholder.id
+ );
+ }
+
+ // THEN expressions should infer Utf8 (to match the known
literal)
+ if let Expr::Placeholder(placeholder) = then_expr.as_ref()
{
+ assert_eq!(
+ placeholder.field.as_ref().unwrap().data_type(),
+ &DataType::Utf8,
+ "THEN placeholder {} should infer Utf8",
+ placeholder.id
+ );
+ }
+ }
+
+ // ELSE expression should infer Utf8 (to match THEN
expressions)
+ if let Some(Expr::Placeholder(placeholder)) =
+ else_expr.as_ref().map(|e| e.as_ref())
+ {
+ assert_eq!(
+ placeholder.field.as_ref().unwrap().data_type(),
+ &DataType::Utf8,
+ "ELSE placeholder {} should infer Utf8",
+ placeholder.id
+ );
+ }
+ }
+ _ => panic!("Expected Case expression"),
+ }
+ }
+
+ #[test]
+ fn infer_placeholder_searched_case() {
Review Comment:
Is there any way we could make these slt tests?
##########
datafusion/expr/src/expr.rs:
##########
@@ -2024,6 +2025,110 @@ impl Expr {
| Expr::SimilarTo(Like { expr, pattern, .. }) => {
rewrite_placeholder(pattern.as_mut(), expr.as_ref(),
schema)?;
}
+ Expr::Case(Case {
+ expr,
+ when_then_expr,
+ else_expr,
+ }) => {
+ // If `expr` is present, then it must match the types of
each WHEN expression.
+ // If `expr` is not present, then the type of each WHEN
expression must evaluate to a boolean.
+ // The types of the THEN and ELSE expressions must match
the result type of the CASE expression.
+
+ fn concrete_type(expr: &Expr, schema: &DFSchema) ->
Option<DataType> {
Review Comment:
This seems like it could be a broader helper function? Maybe on `Expr`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
