ARROW-7493: [Python] Expose sum kernel in pyarrow.compute and support ChunkedArray inputs

Author: xhochy

cpp/src/arrow/compute/kernels/aggregate.cc

@@ -50,14 +50,29 @@ class ManagedAggregateState {
 };
 
 Status AggregateUnaryKernel::Call(FunctionContext* ctx, const Datum& input, Datum* out) {
-  if (!input.is_array()) return Status::Invalid("AggregateKernel expects Array datum");
+  if (input.is_arraylike()) {
+    auto state = ManagedAggregateState::Make(aggregate_function_, ctx->memory_pool());
+    if (!state) return Status::OutOfMemory("AggregateState allocation failed");
 
-  auto state = ManagedAggregateState::Make(aggregate_function_, ctx->memory_pool());
-  if (!state) return Status::OutOfMemory("AggregateState allocation failed");
-
-  auto array = input.make_array();
-  RETURN_NOT_OK(aggregate_function_->Consume(*array, state->mutable_data()));
-  RETURN_NOT_OK(aggregate_function_->Finalize(state->mutable_data(), out));
+    if (input.is_array()) {
+      auto array = input.make_array();
+      RETURN_NOT_OK(aggregate_function_->Consume(*array, state->mutable_data()));
+    } else {
+      auto chunked_array = input.chunked_array();
+      for (int64_t i = 0; i < chunked_array->num_chunks(); i++) {
+        auto tmp_state =
+            ManagedAggregateState::Make(aggregate_function_, ctx->memory_pool());
+        if (!tmp_state) return Status::OutOfMemory("AggregateState allocation failed");
+        RETURN_NOT_OK(aggregate_function_->Consume(*chunked_array->chunk(i),
+                                                   tmp_state->mutable_data()));
+        RETURN_NOT_OK(
+            aggregate_function_->Merge(tmp_state->mutable_data(), state->mutable_data()));
+      }
+    }
+    RETURN_NOT_OK(aggregate_function_->Finalize(state->mutable_data(), out));
+  } else {
+    return Status::Invalid("AggregateKernel expects Array or ChunkedArray datum");
+  }
 
   return Status::OK();
 }

cpp/src/arrow/compute/kernels/aggregate_test.cc

@@ -96,12 +96,29 @@ void ValidateSum(FunctionContext* ctx, const Array& input, Datum expected) {
   DatumEqual<OutputType>::EnsureEqual(result, expected);
 }
 
+template <typename ArrowType>
+void ValidateSum(FunctionContext* ctx, const std::shared_ptr<ChunkedArray>& input,
+                 Datum expected) {
+  using OutputType = typename FindAccumulatorType<ArrowType>::Type;
+
+  Datum result;
+  ASSERT_OK(Sum(ctx, input, &result));
+  DatumEqual<OutputType>::EnsureEqual(result, expected);
+}
+
 template <typename ArrowType>
 void ValidateSum(FunctionContext* ctx, const char* json, Datum expected) {
   auto array = ArrayFromJSON(TypeTraits<ArrowType>::type_singleton(), json);
   ValidateSum<ArrowType>(ctx, *array, expected);
 }
 
+template <typename ArrowType>
+void ValidateSum(FunctionContext* ctx, const std::vector<std::string>& json,
+                 Datum expected) {
+  auto array = ChunkedArrayFromJSON(TypeTraits<ArrowType>::type_singleton(), json);
+  ValidateSum<ArrowType>(ctx, array, expected);
+}
+
 template <typename ArrowType>
 void ValidateSum(FunctionContext* ctx, const Array& array) {
   ValidateSum<ArrowType>(ctx, array, NaiveSum<ArrowType>(array));
@@ -123,6 +140,18 @@ TYPED_TEST(TestNumericSumKernel, SimpleSum) {
   ValidateSum<TypeParam>(&this->ctx_, "[0, 1, 2, 3, 4, 5]",
                          Datum(std::make_shared<ScalarType>(static_cast<T>(5 * 6 / 2))));
 
+  std::vector<std::string> json{"[0, 1, 2, 3, 4, 5]"};
+  ValidateSum<TypeParam>(&this->ctx_, json,
+                         Datum(std::make_shared<ScalarType>(static_cast<T>(5 * 6 / 2))));
+
+  json = {"[0, 1, 2]", "[3, 4, 5]"};
+  ValidateSum<TypeParam>(&this->ctx_, json,
+                         Datum(std::make_shared<ScalarType>(static_cast<T>(5 * 6 / 2))));
+
+  json = {"[0, 1, 2]", "[]", "[3, 4, 5]"};
+  ValidateSum<TypeParam>(&this->ctx_, json,
+                         Datum(std::make_shared<ScalarType>(static_cast<T>(5 * 6 / 2))));
+
   const T expected_result = static_cast<T>(14);
   ValidateSum<TypeParam>(&this->ctx_, "[1, null, 3, null, 3, null, 7]",
                          Datum(std::make_shared<ScalarType>(expected_result)));

cpp/src/arrow/compute/kernels/sum.h

@@ -24,6 +24,7 @@
 namespace arrow {
 
 class Array;
+class ChunkedArray;
 class DataType;
 class Status;
 

python/CMakeLists.txt

@@ -370,7 +370,12 @@ if(UNIX)
   set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
 endif()
 
-set(CYTHON_EXTENSIONS lib _fs _csv _json)
+set(CYTHON_EXTENSIONS
+    lib
+    _fs
+    _csv
+    _json
+    _compute)
 
 set(LINK_LIBS arrow_shared arrow_python_shared)
 

python/pyarrow/_compute.pyx

@@ -0,0 +1,62 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# cython: language_level = 3
+
+from pyarrow.lib cimport (
+    Array,
+    wrap_datum,
+    _context,
+    check_status,
+    ChunkedArray
+)
+from pyarrow.includes.libarrow cimport CDatum, Sum
+
+from typing import Union
+
+
+cdef _sum_array(array: Array):
+    cdef CDatum out
+
+    with nogil:
+        check_status(Sum(_context(), CDatum(array.sp_array), &out))
+
+    return wrap_datum(out)
+
+
+cdef _sum_chunked_array(array: ChunkedArray):
+    cdef CDatum out
+
+    with nogil:
+        check_status(Sum(_context(), CDatum(array.sp_chunked_array), &out))
+
+    return wrap_datum(out)
+
+
+def sum(array: Union[Array, ChunkedArray]):
+    """
+    Sum the values in a numerical (chunked) array.
+    """
+    if isinstance(array, Array):
+        return _sum_array(array)
+    elif isinstance(array, ChunkedArray):
+        return _sum_chunked_array(array)
+    else:
+        raise ValueError(
+            "Only pyarrow.Array and pyarrow.ChunkedArray supported as"
+            " an input, passed {}".format(type(array))
+        )

python/pyarrow/compute.py

@@ -0,0 +1,22 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import absolute_import
+
+from pyarrow._compute import (  # noqa
+    sum
+)

python/pyarrow/lib.pxd

@@ -30,6 +30,7 @@ cdef extern from "Python.h":
     int PySlice_Check(object)
 
 
+cdef CFunctionContext* _context() nogil
 cdef int check_status(const CStatus& status) nogil except -1
 
 cdef class Message:
@@ -429,6 +430,7 @@ cdef class ExtensionArray(Array):
 
 
 cdef wrap_array_output(PyObject* output)
+cdef wrap_datum(const CDatum& datum)
 cdef object box_scalar(DataType type,
                        const shared_ptr[CArray]& sp_array,
                        int64_t index)

python/pyarrow/tests/test_compute.py

@@ -20,6 +20,7 @@
 import pytest
 
 import pyarrow as pa
+import pyarrow.compute
 
 
 all_array_types = [
@@ -44,7 +45,7 @@
 ]
 
 
-@pytest.mark.parametrize('arrow_type', [
+numerical_arrow_types = [
     pa.int8(),
     pa.int16(),
     pa.int64(),
@@ -53,10 +54,32 @@
     pa.uint64(),
     pa.float32(),
     pa.float64()
-])
-def test_sum(arrow_type):
+]
+
+
+@pytest.mark.parametrize('arrow_type', numerical_arrow_types)
+def test_sum_array(arrow_type):
     arr = pa.array([1, 2, 3, 4], type=arrow_type)
     assert arr.sum() == 10
+    assert pa.compute.sum(arr) == 10
+
+
+@pytest.mark.parametrize('arrow_type', numerical_arrow_types)
+def test_sum_chunked_array(arrow_type):
+    arr = pa.chunked_array([pa.array([1, 2, 3, 4], type=arrow_type)])
+    assert pa.compute.sum(arr) == 10
+
+    arr = pa.chunked_array([
+        pa.array([1, 2], type=arrow_type), pa.array([3, 4], type=arrow_type)
+    ])
+    assert pa.compute.sum(arr) == 10
+
+    arr = pa.chunked_array([
+        pa.array([1, 2], type=arrow_type),
+        pa.array([], type=arrow_type),
+        pa.array([3, 4], type=arrow_type)
+    ])
+    assert pa.compute.sum(arr) == 10
 
 
 @pytest.mark.parametrize(('ty', 'values'), all_array_types)

python/setup.py

@@ -179,6 +179,7 @@ def initialize_options(self):
         '_csv',
         '_json',
         '_cuda',
+        '_compute',
         '_flight',
         '_dataset',
         '_parquet',