Review

Author: xhochy

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

@@ -50,31 +50,29 @@ class ManagedAggregateState {
 };
 
 Status AggregateUnaryKernel::Call(FunctionContext* ctx, const Datum& input, Datum* out) {
-  if (input.is_arraylike()) {
-    auto state = ManagedAggregateState::Make(aggregate_function_, ctx->memory_pool());
-    if (!state) return Status::OutOfMemory("AggregateState allocation failed");
+  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 (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_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 {
-    return Status::Invalid("AggregateKernel expects Array or ChunkedArray datum");
+    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

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

python/pyarrow/_compute.pyx

@@ -48,6 +48,14 @@ cdef _sum_chunked_array(array: ChunkedArray):
 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)