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

This only exposes the `Sum` kernel, I will do more once this PR got review and is merged.

Closes #6123 from xhochy/ARROW-7493 and squashes the following commits:

f65ba06 <Antoine Pitrou> Test summing a ChunkedArray with 0 chunks
7c8d526 <Antoine Pitrou> Nits
8fa07b0 <Uwe L. Korn> Revert back to explicit json variable
1575a74 <Uwe L. Korn> Review
d03af5f <Uwe L. Korn> s/int64/int
ce2b974 <Uwe L. Korn> Don't use 3.6+ type annotations
d0b25ed <Uwe L. Korn> ARROW-7493:  Expose sum kernel in pyarrow.compute and support ChunkedArray inputs

Lead-authored-by: Uwe L. Korn <[email protected]>
Co-authored-by: Antoine Pitrou <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>

Author: xhochy

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

@@ -50,16 +50,33 @@ 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()) {
+    return Status::Invalid("AggregateKernel expects Array or ChunkedArray datum");
+  }
   auto state = ManagedAggregateState::Make(aggregate_function_, ctx->memory_pool());
-  if (!state) return Status::OutOfMemory("AggregateState allocation failed");
+  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 (int 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 Status::OK();
+  return aggregate_function_->Finalize(state->mutable_data(), out);
 }
 
 std::shared_ptr<DataType> AggregateUnaryKernel::out_type() const {

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,22 @@ 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> chunks = {"[0, 1, 2, 3, 4, 5]"};
+  ValidateSum<TypeParam>(&this->ctx_, chunks,
+                         Datum(std::make_shared<ScalarType>(static_cast<T>(5 * 6 / 2))));
+
+  chunks = {"[0, 1, 2]", "[3, 4, 5]"};
+  ValidateSum<TypeParam>(&this->ctx_, chunks,
+                         Datum(std::make_shared<ScalarType>(static_cast<T>(5 * 6 / 2))));
+
+  chunks = {"[0, 1, 2]", "[]", "[3, 4, 5]"};
+  ValidateSum<TypeParam>(&this->ctx_, chunks,
+                         Datum(std::make_shared<ScalarType>(static_cast<T>(5 * 6 / 2))));
+
+  chunks = {};
+  ValidateSum<TypeParam>(&this->ctx_, chunks,
+                         Datum(std::make_shared<ScalarType>()));  // null
+
   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;
 

cpp/src/arrow/testing/gtest_util.cc

@@ -231,7 +231,7 @@ std::shared_ptr<ChunkedArray> ChunkedArrayFromJSON(const std::shared_ptr<DataTyp
   for (const std::string& chunk_json : json) {
     out_chunks.push_back(ArrayFromJSON(type, chunk_json));
   }
-  return std::make_shared<ChunkedArray>(std::move(out_chunks));
+  return std::make_shared<ChunkedArray>(std::move(out_chunks), type);
 }
 
 std::shared_ptr<RecordBatch> RecordBatchFromJSON(const std::shared_ptr<Schema>& schema,

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,68 @@
+# 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
+
+
+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):
+    """
+    Sum the values in a numerical (chunked) array.
+
+    Parameters
+    ----------
+    array : pyarrow.Array or pyarrow.ChunkedArray
+
+    Returns
+    -------
+    sum : pyarrow.Scalar
+    """
+    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/includes/libarrow.pxd

@@ -762,6 +762,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
     cdef cppclass CScalar" arrow::Scalar":
         shared_ptr[CDataType] type
+        c_bool is_valid
 
     cdef cppclass CInt8Scalar" arrow::Int8Scalar"(CScalar):
         int8_t value

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/scalar.pxi

@@ -1031,6 +1031,7 @@ cdef class ScalarValue(Scalar):
     def __hash__(self):
         return hash(self.as_py())
 
+
 cdef class UInt8Scalar(ScalarValue):
     """
     Concrete class for uint8 scalars.
@@ -1041,7 +1042,7 @@ cdef class UInt8Scalar(ScalarValue):
         Return this value as a Python int.
         """
         cdef CUInt8Scalar* sp = <CUInt8Scalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef class Int8Scalar(ScalarValue):
@@ -1054,7 +1055,7 @@ cdef class Int8Scalar(ScalarValue):
         Return this value as a Python int.
         """
         cdef CInt8Scalar* sp = <CInt8Scalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef class UInt16Scalar(ScalarValue):
@@ -1067,7 +1068,7 @@ cdef class UInt16Scalar(ScalarValue):
         Return this value as a Python int.
         """
         cdef CUInt16Scalar* sp = <CUInt16Scalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef class Int16Scalar(ScalarValue):
@@ -1080,7 +1081,7 @@ cdef class Int16Scalar(ScalarValue):
         Return this value as a Python int.
         """
         cdef CInt16Scalar* sp = <CInt16Scalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef class UInt32Scalar(ScalarValue):
@@ -1093,7 +1094,7 @@ cdef class UInt32Scalar(ScalarValue):
         Return this value as a Python int.
         """
         cdef CUInt32Scalar* sp = <CUInt32Scalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef class Int32Scalar(ScalarValue):
@@ -1106,7 +1107,7 @@ cdef class Int32Scalar(ScalarValue):
         Return this value as a Python int.
         """
         cdef CInt32Scalar* sp = <CInt32Scalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef class UInt64Scalar(ScalarValue):
@@ -1119,7 +1120,7 @@ cdef class UInt64Scalar(ScalarValue):
         Return this value as a Python int.
         """
         cdef CUInt64Scalar* sp = <CUInt64Scalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef class Int64Scalar(ScalarValue):
@@ -1132,7 +1133,7 @@ cdef class Int64Scalar(ScalarValue):
         Return this value as a Python int.
         """
         cdef CInt64Scalar* sp = <CInt64Scalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef class FloatScalar(ScalarValue):
@@ -1145,7 +1146,7 @@ cdef class FloatScalar(ScalarValue):
         Return this value as a Python float.
         """
         cdef CFloatScalar* sp = <CFloatScalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef class DoubleScalar(ScalarValue):
@@ -1158,7 +1159,7 @@ cdef class DoubleScalar(ScalarValue):
         Return this value as a Python float.
         """
         cdef CDoubleScalar* sp = <CDoubleScalar*> self.sp_scalar.get()
-        return sp.value
+        return sp.value if sp.is_valid else None
 
 
 cdef dict _scalar_classes = {

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,41 @@
     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
+
+    arr = pa.array([], type=arrow_type)
+    assert arr.sum() == None  # noqa: E711
+    assert pa.compute.sum(arr) == None  # noqa: E711
+
+
+@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
+
+    arr = pa.chunked_array((), type=arrow_type)
+    print(arr, type(arr))
+    assert arr.num_chunks == 0
+    assert pa.compute.sum(arr) == None  # noqa: E711
 
 
 @pytest.mark.parametrize(('ty', 'values'), all_array_types)