Update python wrapper for arguments appropriate to bool operators

Author: timsaucer

python/datafusion/functions.py

@@ -1901,14 +1901,38 @@ def bit_xor(
     return Expr(f.bit_xor(expression.expr, distinct=distinct, filter=filter_raw))
 
 
-def bool_and(arg: Expr, distinct: bool = False) -> Expr:
-    """Computes the boolean AND of the argument."""
-    return Expr(f.bool_and(arg.expr, distinct=distinct))
+def bool_and(expression: Expr, filter: Optional[Expr] = None) -> Expr:
+    """Computes the boolean AND of the argument.
 
+    This aggregate function will compare every value in the input partition. These are
+    expected to be boolean values.
 
-def bool_or(arg: Expr, distinct: bool = False) -> Expr:
-    """Computes the boolean OR of the argument."""
-    return Expr(f.bool_or(arg.expr, distinct=distinct))
+    If using the builder functions described in ref:`_aggregation` this function ignores
+    the options ``order_by``, ``null_treatment``, and ``distinct``.
+
+    Args:
+        expression: Argument to perform calculation on
+        filter: If provided, only compute against rows for which the filter is true
+    """
+    filter_raw = filter.expr if filter is not None else None
+    return Expr(f.bool_and(expression.expr, filter=filter_raw))
+
+
+def bool_or(expression: Expr, filter: Optional[Expr] = None) -> Expr:
+    """Computes the boolean OR of the argument.
+
+    This aggregate function will compare every value in the input partition. These are
+    expected to be boolean values.
+
+    If using the builder functions described in ref:`_aggregation` this function ignores
+    the options ``order_by``, ``null_treatment``, and ``distinct``.
+
+    Args:
+        expression: Argument to perform calculation on
+        filter: If provided, only compute against rows for which the filter is true
+    """
+    filter_raw = filter.expr if filter is not None else None
+    return Expr(f.bool_or(expression.expr, filter=filter_raw))
 
 
 def lead(

python/datafusion/tests/test_aggregation.py

@@ -179,7 +179,7 @@ def test_aggregate_100(df_aggregate_100):
     assert result.column("c5") == pa.array([83, 68, 122, 124, 117])
 
 
-data_test_bitwise_functions = [
+data_test_bitwise_and_boolean_functions = [
     ("bit_and", f.bit_and(column("a")), [0]),
     ("bit_and_filter", f.bit_and(column("a"), filter=column("a") != lit(2)), [1]),
     ("bit_or", f.bit_or(column("b")), [6]),
@@ -192,29 +192,19 @@ def test_aggregate_100(df_aggregate_100):
         f.bit_xor(column("b"), distinct=True, filter=column("a") != lit(3)),
         [4],
     ),
+    ("bool_and", f.bool_and(column("d")), [False]),
+    ("bool_and_filter", f.bool_and(column("d"), filter=column("a") != lit(3)), [True]),
+    ("bool_or", f.bool_or(column("d")), [True]),
+    ("bool_or_filter", f.bool_or(column("d"), filter=column("a") == lit(3)), [False]),
 ]
 
 
-@pytest.mark.parametrize("name,expr,result", data_test_bitwise_functions)
-def test_bit_add_or_xor(df, name, expr, result):
+@pytest.mark.parametrize("name,expr,result", data_test_bitwise_and_boolean_functions)
+def test_bit_and_bool_fns(df, name, expr, result):
     df = df.aggregate([], [expr.alias(name)])
 
     expected = {
         name: result,
     }
 
     assert df.collect()[0].to_pydict() == expected
-
-
-def test_bool_and_or(df):
-    df = df.aggregate(
-        [],
-        [
-            f.bool_and(column("d")),
-            f.bool_or(column("d")),
-        ],
-    )
-    result = df.collect()
-    result = result[0]
-    assert result.column(0) == pa.array([False])
-    assert result.column(1) == pa.array([True])