migrate bitwise functions to UDAF

The functions now take a single expression instead of a Vec<_>.

Ref: apache/datafusion#10930

Author: Michael-J-Ward

python/datafusion/functions.py

@@ -1482,31 +1482,26 @@ def last_value(
     )
 
 
-def bit_and(*args: Expr, distinct: bool = False) -> Expr:
+def bit_and(arg: Expr, distinct: bool = False) -> Expr:
     """Computes the bitwise AND of the argument."""
-    args = [arg.expr for arg in args]
-    return Expr(f.bit_and(*args, distinct=distinct))
+    return Expr(f.bit_and(arg.expr, distinct=distinct))
 
 
-def bit_or(*args: Expr, distinct: bool = False) -> Expr:
+def bit_or(arg: Expr, distinct: bool = False) -> Expr:
     """Computes the bitwise OR of the argument."""
-    args = [arg.expr for arg in args]
-    return Expr(f.bit_or(*args, distinct=distinct))
+    return Expr(f.bit_or(arg.expr, distinct=distinct))
 
 
-def bit_xor(*args: Expr, distinct: bool = False) -> Expr:
+def bit_xor(arg: Expr, distinct: bool = False) -> Expr:
     """Computes the bitwise XOR of the argument."""
-    args = [arg.expr for arg in args]
-    return Expr(f.bit_xor(*args, distinct=distinct))
+    return Expr(f.bit_xor(arg.expr, distinct=distinct))
 
 
-def bool_and(*args: Expr, distinct: bool = False) -> Expr:
+def bool_and(arg: Expr, distinct: bool = False) -> Expr:
     """Computes the boolean AND of the arugment."""
-    args = [arg.expr for arg in args]
-    return Expr(f.bool_and(*args, distinct=distinct))
+    return Expr(f.bool_and(arg.expr, distinct=distinct))
 
 
-def bool_or(*args: Expr, distinct: bool = False) -> Expr:
+def bool_or(arg: Expr, distinct: bool = False) -> Expr:
     """Computes the boolean OR of the arguement."""
-    args = [arg.expr for arg in args]
-    return Expr(f.bool_or(*args, distinct=distinct))
+    return Expr(f.bool_or(arg.expr, distinct=distinct))

src/functions.rs

@@ -98,6 +98,56 @@ pub fn avg(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
     }
 }
 
+#[pyfunction]
+pub fn bit_and(expr_x: PyExpr, distinct: bool) -> PyResult<PyExpr> {
+    let expr = functions_aggregate::expr_fn::bit_and(expr_x.expr);
+    if distinct {
+        Ok(expr.distinct().build()?.into())
+    } else {
+        Ok(expr.into())
+    }
+}
+
+#[pyfunction]
+pub fn bit_or(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
+    let expr = functions_aggregate::expr_fn::bit_or(expression.expr);
+    if distinct {
+        Ok(expr.distinct().build()?.into())
+    } else {
+        Ok(expr.into())
+    }
+}
+
+#[pyfunction]
+pub fn bit_xor(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
+    let expr = functions_aggregate::expr_fn::bit_xor(expression.expr);
+    if distinct {
+        Ok(expr.distinct().build()?.into())
+    } else {
+        Ok(expr.into())
+    }
+}
+
+#[pyfunction]
+pub fn bool_and(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
+    let expr = functions_aggregate::expr_fn::bool_and(expression.expr);
+    if distinct {
+        Ok(expr.distinct().build()?.into())
+    } else {
+        Ok(expr.into())
+    }
+}
+
+#[pyfunction]
+pub fn bool_or(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
+    let expr = functions_aggregate::expr_fn::bool_or(expression.expr);
+    if distinct {
+        Ok(expr.distinct().build()?.into())
+    } else {
+        Ok(expr.into())
+    }
+}
+
 #[pyfunction]
 pub fn mean(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
     // alias for avg
@@ -937,11 +987,6 @@ array_fn!(range, start stop step);
 aggregate_function!(array_agg, ArrayAgg);
 aggregate_function!(max, Max);
 aggregate_function!(min, Min);
-aggregate_function!(bit_and, BitAnd);
-aggregate_function!(bit_or, BitOr);
-aggregate_function!(bit_xor, BitXor);
-aggregate_function!(bool_and, BoolAnd);
-aggregate_function!(bool_or, BoolOr);
 
 pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(abs))?;