feat(cube): Coerce numeric types to Boolean in comparisons

Author: srh

datafusion/expr-common/src/type_coercion/binary.rs

@@ -496,6 +496,7 @@ pub fn comparison_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<D
         return Some(lhs_type.clone());
     }
     binary_numeric_coercion(lhs_type, rhs_type)
+        .or_else(|| number_boolean_coercion(lhs_type, rhs_type))
         .or_else(|| dictionary_comparison_coercion(lhs_type, rhs_type, true))
         .or_else(|| temporal_coercion_nonstrict_timezone(lhs_type, rhs_type))
         .or_else(|| string_coercion(lhs_type, rhs_type))
@@ -524,6 +525,16 @@ pub fn values_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataT
         .or_else(|| binary_coercion(lhs_type, rhs_type))
 }
 
+/// Coerce `lhs_type` and `rhs_type` to Boolean when doing a number-vs.-bool comparison.
+fn number_boolean_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataType> {
+    use arrow::datatypes::DataType::*;
+    match (lhs_type, rhs_type) {
+        (Boolean, _) if rhs_type.is_numeric() => Some(Boolean),
+        (_, Boolean) if lhs_type.is_numeric() => Some(Boolean),
+        _ => None,
+    }
+}
+
 /// Coerce `lhs_type` and `rhs_type` to a common type for the purposes of a comparison operation
 /// where one is numeric and one is `Utf8`/`LargeUtf8`.
 fn string_numeric_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataType> {

datafusion/optimizer/src/analyzer/type_coercion.rs

@@ -1473,10 +1473,21 @@ mod test {
         assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
 
         let empty = empty_with_type(DataType::Int64);
+        let plan =
+            LogicalPlan::Projection(Projection::try_new(vec![expr.clone()], empty)?);
+        // Cube: We support type conversion for boolean-int comparisons.
+        let expected = "Projection: CAST(a AS Boolean) IS TRUE\n  EmptyRelation";
+        assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, expected)?;
+        // let ret = assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, "");
+        // let err = ret.unwrap_err().to_string();
+        // assert!(err.contains("Cannot infer common argument type for comparison operation Int64 IS DISTINCT FROM Boolean"), "{err}");
+
+        // Cube: Add original test case for non-coercion but with a non-numeric type.
+        let empty = empty_with_type(DataType::Binary);
         let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
         let ret = assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), plan, "");
         let err = ret.unwrap_err().to_string();
-        assert!(err.contains("Cannot infer common argument type for comparison operation Int64 IS DISTINCT FROM Boolean"), "{err}");
+        assert!(err.contains("Cannot infer common argument type for comparison operation Binary IS DISTINCT FROM Boolean"), "{err}");
 
         // is not true
         let expr = col("a").is_not_true();