feat: Add bool_and, bool_or aggregate functions

Author: MazterQyou

datafusion/core/src/physical_plan/aggregates.rs

@@ -83,6 +83,7 @@ pub fn return_type(
             Ok(coerced_data_types[0].clone())
         }
         AggregateFunction::ApproxMedian => Ok(coerced_data_types[0].clone()),
+        AggregateFunction::BoolAnd | AggregateFunction::BoolOr => Ok(DataType::Boolean),
     }
 }
 
@@ -297,6 +298,13 @@ pub fn create_aggregate_expr(
                 "MEDIAN(DISTINCT) aggregations are not available".to_string(),
             ));
         }
+        (AggregateFunction::BoolAnd, _) => Arc::new(expressions::BoolAnd::new(
+            coerced_phy_exprs[0].clone(),
+            name,
+        )),
+        (AggregateFunction::BoolOr, _) => {
+            Arc::new(expressions::BoolOr::new(coerced_phy_exprs[0].clone(), name))
+        }
     })
 }
 
@@ -374,16 +382,19 @@ pub(super) fn signature(fun: &AggregateFunction) -> Signature {
                 .collect(),
             Volatility::Immutable,
         ),
+        AggregateFunction::BoolAnd | AggregateFunction::BoolOr => {
+            Signature::exact(vec![DataType::Boolean], Volatility::Immutable)
+        }
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::physical_plan::expressions::{
-        ApproxDistinct, ApproxMedian, ApproxPercentileCont, ArrayAgg, Avg, Correlation,
-        Count, Covariance, DistinctArrayAgg, DistinctCount, Max, Min, Stddev, Sum,
-        Variance,
+        ApproxDistinct, ApproxMedian, ApproxPercentileCont, ArrayAgg, Avg, BoolAnd,
+        BoolOr, Correlation, Count, Covariance, DistinctArrayAgg, DistinctCount, Max,
+        Min, Stddev, Sum, Variance,
     };
     use crate::{error::Result, scalar::ScalarValue};
 
@@ -995,6 +1006,45 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_bool_and_or_expr() -> Result<()> {
+        let funcs = vec![AggregateFunction::BoolAnd, AggregateFunction::BoolOr];
+        for fun in funcs {
+            let input_schema =
+                Schema::new(vec![Field::new("c1", DataType::Boolean, true)]);
+            let input_phy_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![Arc::new(
+                expressions::Column::new_with_schema("c1", &input_schema).unwrap(),
+            )];
+            let result_agg_phy_exprs = create_aggregate_expr(
+                &fun,
+                false,
+                &input_phy_exprs[0..1],
+                &input_schema,
+                "c1",
+            )?;
+            match fun {
+                AggregateFunction::BoolAnd => {
+                    assert!(result_agg_phy_exprs.as_any().is::<BoolAnd>());
+                    assert_eq!("c1", result_agg_phy_exprs.name());
+                    assert_eq!(
+                        Field::new("c1", DataType::Boolean, true),
+                        result_agg_phy_exprs.field().unwrap()
+                    );
+                }
+                AggregateFunction::BoolOr => {
+                    assert!(result_agg_phy_exprs.as_any().is::<BoolOr>());
+                    assert_eq!("c1", result_agg_phy_exprs.name());
+                    assert_eq!(
+                        Field::new("c1", DataType::Boolean, true),
+                        result_agg_phy_exprs.field().unwrap()
+                    );
+                }
+                _ => {}
+            };
+        }
+        Ok(())
+    }
+
     #[test]
     fn test_median() -> Result<()> {
         let observed = return_type(&AggregateFunction::ApproxMedian, &[DataType::Utf8]);
@@ -1158,4 +1208,32 @@ mod tests {
         let observed = return_type(&AggregateFunction::Stddev, &[DataType::Utf8]);
         assert!(observed.is_err());
     }
+
+    #[test]
+    fn test_bool_and_return_type() -> Result<()> {
+        let observed = return_type(&AggregateFunction::BoolAnd, &[DataType::Boolean])?;
+        assert_eq!(DataType::Boolean, observed);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_bool_and_no_utf8() {
+        let observed = return_type(&AggregateFunction::BoolAnd, &[DataType::Utf8]);
+        assert!(observed.is_err());
+    }
+
+    #[test]
+    fn test_bool_or_return_type() -> Result<()> {
+        let observed = return_type(&AggregateFunction::BoolOr, &[DataType::Boolean])?;
+        assert_eq!(DataType::Boolean, observed);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_bool_or_no_utf8() {
+        let observed = return_type(&AggregateFunction::BoolOr, &[DataType::Utf8]);
+        assert!(observed.is_err());
+    }
 }

datafusion/expr/src/aggregate_function.rs

@@ -57,6 +57,10 @@ pub enum AggregateFunction {
     ApproxPercentileContWithWeight,
     /// ApproxMedian
     ApproxMedian,
+    /// BoolAnd
+    BoolAnd,
+    /// BoolOr
+    BoolOr,
 }
 
 impl fmt::Display for AggregateFunction {
@@ -92,6 +96,8 @@ impl FromStr for AggregateFunction {
                 AggregateFunction::ApproxPercentileContWithWeight
             }
             "approx_median" => AggregateFunction::ApproxMedian,
+            "bool_and" => AggregateFunction::BoolAnd,
+            "bool_or" => AggregateFunction::BoolOr,
             _ => {
                 return Err(DataFusionError::Plan(format!(
                     "There is no built-in function named {}",

datafusion/expr/src/expr_fn.rs

@@ -179,6 +179,24 @@ pub fn approx_percentile_cont_with_weight(
     }
 }
 
+/// Create an expression to represent the bool_and() aggregate function
+pub fn bool_and(expr: Expr) -> Expr {
+    Expr::AggregateFunction {
+        fun: aggregate_function::AggregateFunction::BoolAnd,
+        distinct: false,
+        args: vec![expr],
+    }
+}
+
+/// Create an expression to represent the bool_or() aggregate function
+pub fn bool_or(expr: Expr) -> Expr {
+    Expr::AggregateFunction {
+        fun: aggregate_function::AggregateFunction::BoolOr,
+        distinct: false,
+        args: vec![expr],
+    }
+}
+
 // TODO(kszucs): this seems buggy, unary_scalar_expr! is used for many
 // varying arity functions
 /// Create an convenience function representing a unary scalar function

datafusion/physical-expr/src/coercion_rule/aggregate_rule.rs

@@ -182,6 +182,15 @@ pub fn coerce_types(
             }
             Ok(input_types.to_vec())
         }
+        AggregateFunction::BoolAnd | AggregateFunction::BoolOr => {
+            if input_types[0] != DataType::Boolean {
+                return Err(DataFusionError::Plan(format!(
+                    "The function {:?} does not support inputs of type {:?}.",
+                    agg_fun, input_types[0]
+                )));
+            }
+            Ok(input_types.to_vec())
+        }
     }
 }