add bit_and,bit_or,bit_xor,bool_add,bool_or

Author: jiangzhx

datafusion/tests/test_aggregation.py

@@ -33,8 +33,9 @@ def df():
             pa.array([1, 2, 3]),
             pa.array([4, 4, 6]),
             pa.array([9, 8, 5]),
+            pa.array([True, True, False]),
         ],
-        names=["a", "b", "c"],
+        names=["a", "b", "c", "d"],
     )
     return ctx.create_dataframe([[batch]])
 
@@ -125,3 +126,36 @@ def test_built_in_aggregation(df):
     np.testing.assert_array_almost_equal(
         result.column(21), np.var(values_c, ddof=1)
     )
+
+
+def test_bit_add_or_xor(df):
+
+    df = df.aggregate(
+        [],
+        [
+            f.bit_and(column("a")),
+            f.bit_or(column("b")),
+            f.bit_xor(column("c")),
+        ],
+    )
+
+    result = df.collect()
+    result = result[0]
+    assert result.column(0) == pa.array([0])
+    assert result.column(1) == pa.array([6])
+    assert result.column(2) == pa.array([4])
+
+
+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])

src/functions.rs

@@ -357,6 +357,11 @@ aggregate_function!(stddev_samp, Stddev);
 aggregate_function!(var, Variance);
 aggregate_function!(var_pop, VariancePop);
 aggregate_function!(var_samp, Variance);
+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: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(abs))?;
@@ -482,5 +487,10 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(var_pop))?;
     m.add_wrapped(wrap_pyfunction!(var_samp))?;
     m.add_wrapped(wrap_pyfunction!(window))?;
+    m.add_wrapped(wrap_pyfunction!(bit_and))?;
+    m.add_wrapped(wrap_pyfunction!(bit_or))?;
+    m.add_wrapped(wrap_pyfunction!(bit_xor))?;
+    m.add_wrapped(wrap_pyfunction!(bool_and))?;
+    m.add_wrapped(wrap_pyfunction!(bool_or))?;
     Ok(())
 }