[arrow-datafusion-python] branch main updated: add bit_and,bit_or,bit_xor,bool_add,bool_or (#496)

[agrove]

This is an automated email from the ASF dual-hosted git repository.

agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion-python.git


The following commit(s) were added to refs/heads/main by this push:
     new 4c7b14c  add bit_and,bit_or,bit_xor,bool_add,bool_or (#496)
4c7b14c is described below

commit 4c7b14c6e4da99a7e52fddaf6fed9ae51ca2c66e
Author: zhenxing jiang <jiang...@gmail.com>
AuthorDate: Thu Oct 12 10:22:27 2023 -0500

    add bit_and,bit_or,bit_xor,bool_add,bool_or (#496)
    
    * add bit_and,bit_or,bit_xor,bool_add,bool_or
    
    * Update datafusion/tests/test_aggregation.py
    
    Co-authored-by: Liang-Chi Hsieh <vii...@gmail.com>
    
    ---------
    
    Co-authored-by: Liang-Chi Hsieh <vii...@gmail.com>
---
 datafusion/tests/test_aggregation.py | 38 ++++++++++++++++++++++++++++++++++--
 src/functions.rs                     | 10 ++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/datafusion/tests/test_aggregation.py 
b/datafusion/tests/test_aggregation.py
index 2c8c064..0a6c90c 100644
--- a/datafusion/tests/test_aggregation.py
+++ b/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]])
 
@@ -73,7 +74,7 @@ def test_built_in_aggregation(df):
         ],
     )
     result = agg_df.collect()[0]
-    values_a, values_b, values_c = df.collect()[0]
+    values_a, values_b, values_c, values_d = df.collect()[0]
 
     assert result.column(0) == pa.array([2], type=pa.uint64())
     assert result.column(1) == pa.array([4])
@@ -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])
diff --git a/src/functions.rs b/src/functions.rs
index ef26240..eed2815 100644
--- a/src/functions.rs
+++ b/src/functions.rs
@@ -362,6 +362,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))?;
@@ -489,6 +494,11 @@ 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))?;
 
     //Binary String Functions
     m.add_wrapped(wrap_pyfunction!(encode))?;