apache · andrioni · Feb 22, 2018 · Feb 22, 2018 · Feb 22, 2018 · Feb 22, 2018
diff --git a/cpp/src/arrow/compute/compute-test.cc b/cpp/src/arrow/compute/compute-test.cc
@@ -876,6 +876,23 @@ void CheckDictEncode(FunctionContext* ctx, const shared_ptr<DataType>& type,
   ASSERT_ARRAYS_EQUAL(expected, *result);
 }
 
+template <typename Type, typename T>
+void CheckCountValues(FunctionContext* ctx, const shared_ptr<DataType>& type,
+                      const vector<T>& in_values, const vector<bool>& in_is_valid,
+                      const vector<T>& out_values, const vector<bool>& out_is_valid,
+                      const vector<int64_t>& out_counts) {
+  shared_ptr<Array> input = _MakeArray<Type, T>(type, in_values, in_is_valid);
+  shared_ptr<Array> ex_values = _MakeArray<Type, T>(type, out_values, out_is_valid);
+  shared_ptr<Array> ex_counts =
+      _MakeArray<Int64Type, int64_t>(int64(), out_counts, out_is_valid);
+
+  shared_ptr<Array> result_values;
+  shared_ptr<Array> result_counts;
+  ASSERT_OK(CountValues(ctx, Datum(input), &result_values, &result_counts));
+  ASSERT_ARRAYS_EQUAL(*ex_values, *result_values);
+  ASSERT_ARRAYS_EQUAL(*ex_counts, *result_counts);
+}
+
 class TestHashKernel : public ComputeFixture, public TestBase {};
 
 template <typename Type>
@@ -903,6 +920,14 @@ TYPED_TEST(TestHashKernelPrimitive, DictEncode) {
                                 {0, 0, 0, 1, 0, 2});
 }
 
+TYPED_TEST(TestHashKernelPrimitive, CountValues) {
+  using T = typename TypeParam::c_type;
+  auto type = TypeTraits<TypeParam>::type_singleton();
+  CheckCountValues<TypeParam, T>(&this->ctx_, type, {2, 1, 2, 1, 2, 3, 4},
+                                 {true, false, true, true, true, true, false}, {2, 1, 3},
+                                 {}, {3, 1, 1});
+}
+
 TYPED_TEST(TestHashKernelPrimitive, PrimitiveResizeTable) {
   using T = typename TypeParam::c_type;
   // Skip this test for (u)int8
@@ -916,12 +941,14 @@ TYPED_TEST(TestHashKernelPrimitive, PrimitiveResizeTable) {
   vector<T> values;
   vector<T> uniques;
   vector<int32_t> indices;
+  vector<int64_t> counts;
   for (int64_t i = 0; i < kTotalValues * kRepeats; i++) {
     const auto val = static_cast<T>(i % kTotalValues);
     values.push_back(val);
 
     if (i < kTotalValues) {
       uniques.push_back(val);
+      counts.push_back(kRepeats);
     }
     indices.push_back(static_cast<int32_t>(i % kTotalValues));
   }
@@ -930,6 +957,8 @@ TYPED_TEST(TestHashKernelPrimitive, PrimitiveResizeTable) {
   CheckUnique<TypeParam, T>(&this->ctx_, type, values, {}, uniques, {});
 
   CheckDictEncode<TypeParam, T>(&this->ctx_, type, values, {}, uniques, {}, indices);
+
+  CheckCountValues<TypeParam, T>(&this->ctx_, type, values, {}, uniques, {}, counts);
 }
 
 TEST_F(TestHashKernel, UniqueTimeTimestamp) {
@@ -944,6 +973,19 @@ TEST_F(TestHashKernel, UniqueTimeTimestamp) {
                                       {});
 }
 
+TEST_F(TestHashKernel, CountValuesTimeTimestamp) {
+  CheckCountValues<Time32Type, int32_t>(&this->ctx_, time32(TimeUnit::SECOND),
+                                        {2, 1, 2, 1}, {true, false, true, true}, {2, 1},
+                                        {}, {2, 1});
+
+  CheckCountValues<Time64Type, int64_t>(&this->ctx_, time64(TimeUnit::NANO), {2, 1, 2, 1},
+                                        {true, false, true, true}, {2, 1}, {}, {2, 1});
+
+  CheckCountValues<TimestampType, int64_t>(&this->ctx_, timestamp(TimeUnit::NANO),
+                                           {2, 1, 2, 1}, {true, false, true, true},
+                                           {2, 1}, {}, {2, 1});
+}
+
 TEST_F(TestHashKernel, UniqueBoolean) {
   CheckUnique<BooleanType, bool>(&this->ctx_, boolean(), {true, true, false, true},
                                  {true, false, true, true}, {true, false}, {});
@@ -978,6 +1020,23 @@ TEST_F(TestHashKernel, DictEncodeBoolean) {
                                      {}, {0, 1, 0, 1, 0});
 }
 
+TEST_F(TestHashKernel, CountValuesBoolean) {
+  CheckCountValues<BooleanType, bool>(&this->ctx_, boolean(), {true, true, false, true},
+                                      {true, false, true, true}, {true, false}, {},
+                                      {2, 1});
+
+  CheckCountValues<BooleanType, bool>(&this->ctx_, boolean(), {false, true, false, true},
+                                      {true, false, true, true}, {false, true}, {},
+                                      {2, 1});
+
+  // No nulls
+  CheckCountValues<BooleanType, bool>(&this->ctx_, boolean(), {true, true, false, true},
+                                      {}, {true, false}, {}, {3, 1});
+
+  CheckCountValues<BooleanType, bool>(&this->ctx_, boolean(), {false, true, false, true},
+                                      {}, {false, true}, {}, {2, 2});
+}
+
 TEST_F(TestHashKernel, UniqueBinary) {
   CheckUnique<BinaryType, std::string>(&this->ctx_, binary(),
                                        {"test", "", "test2", "test"},
@@ -997,6 +1056,16 @@ TEST_F(TestHashKernel, DictEncodeBinary) {
       {true, false, true, true, true}, {"test", "test2", "baz"}, {}, {0, 0, 1, 0, 2});
 }
 
+TEST_F(TestHashKernel, CountValuesBinary) {
+  CheckCountValues<BinaryType, std::string>(
+      &this->ctx_, binary(), {"test", "", "test2", "test"}, {true, false, true, true},
+      {"test", "test2"}, {}, {2, 1});
+
+  CheckCountValues<StringType, std::string>(
+      &this->ctx_, utf8(), {"test", "", "test2", "test"}, {true, false, true, true},
+      {"test", "test2"}, {}, {2, 1});
+}
+
 TEST_F(TestHashKernel, BinaryResizeTable) {
   const int64_t kTotalValues = 10000;
   const int64_t kRepeats = 10;
@@ -1046,6 +1115,7 @@ TEST_F(TestHashKernel, FixedSizeBinaryResizeTable) {
   vector<std::string> values;
   vector<std::string> uniques;
   vector<int32_t> indices;
+  vector<int64_t> counts;
   for (int64_t i = 0; i < kTotalValues * kRepeats; i++) {
     int64_t index = i % kTotalValues;
     std::stringstream ss;
@@ -1056,6 +1126,7 @@ TEST_F(TestHashKernel, FixedSizeBinaryResizeTable) {
 
     if (i < kTotalValues) {
       uniques.push_back(val);
+      counts.push_back(kRepeats);
     }
     indices.push_back(static_cast<int32_t>(i % kTotalValues));
   }
@@ -1065,6 +1136,8 @@ TEST_F(TestHashKernel, FixedSizeBinaryResizeTable) {
                                                 {});
   CheckDictEncode<FixedSizeBinaryType, std::string>(&this->ctx_, type, values, {},
                                                     uniques, {}, indices);
+  CheckCountValues<FixedSizeBinaryType, std::string>(&this->ctx_, type, values, {},
+                                                     uniques, {}, counts);
 }
 
 TEST_F(TestHashKernel, UniqueDecimal) {
@@ -1084,6 +1157,15 @@ TEST_F(TestHashKernel, DictEncodeDecimal) {
                                               {}, {0, 0, 1, 0, 2});
 }
 
+TEST_F(TestHashKernel, CountValuesDecimal) {
+  vector<Decimal128> values{12, 12, 11, 12};
+  vector<Decimal128> expected{12, 11};
+
+  CheckCountValues<Decimal128Type, Decimal128>(&this->ctx_, decimal(2, 0), values,
+                                               {true, false, true, true}, expected, {},
+                                               {2, 1});
+}
+
 TEST_F(TestHashKernel, ChunkedArrayInvoke) {
   vector<std::string> values1 = {"foo", "bar", "foo"};
   vector<std::string> values2 = {"bar", "baz", "quuux", "foo"};
@@ -1095,6 +1177,9 @@ TEST_F(TestHashKernel, ChunkedArrayInvoke) {
   vector<std::string> dict_values = {"foo", "bar", "baz", "quuux"};
   auto ex_dict = _MakeArray<StringType, std::string>(type, dict_values, {});
 
+  vector<int64_t> counts = {3, 2, 1, 1};
+  auto ex_counts = _MakeArray<Int64Type, int64_t>(int64(), counts, {});
+
   ArrayVector arrays = {a1, a2};
   auto carr = std::make_shared<ChunkedArray>(arrays);
 
@@ -1103,6 +1188,13 @@ TEST_F(TestHashKernel, ChunkedArrayInvoke) {
   ASSERT_OK(Unique(&this->ctx_, Datum(carr), &result));
   ASSERT_ARRAYS_EQUAL(*ex_dict, *result);
 
+  // Count values
+  shared_ptr<Array> cv_uniques;
+  shared_ptr<Array> cv_counts;
+  ASSERT_OK(CountValues(&this->ctx_, Datum(carr), &cv_uniques, &cv_counts));
+  ASSERT_ARRAYS_EQUAL(*ex_dict, *cv_uniques);
+  ASSERT_ARRAYS_EQUAL(*ex_counts, *cv_counts);
+
   // Dictionary encode
   auto dict_type = dictionary(int32(), ex_dict);
 

diff --git a/cpp/src/arrow/compute/kernels/hash.cc b/cpp/src/arrow/compute/kernels/hash.cc
@@ -749,6 +749,50 @@ class DictEncodeImpl : public HashTableKernel<Type, DictEncodeImpl<Type>> {
   Int32Builder indices_builder_;
 };
 
+// ----------------------------------------------------------------------
+// Count values implementation
+
+template <typename Type>
+class CountValuesImpl : public HashTableKernel<Type, CountValuesImpl<Type>> {
+ public:
+  static constexpr bool allow_expand = true;
+  using Base = HashTableKernel<Type, CountValuesImpl>;
+
+  CountValuesImpl(const std::shared_ptr<DataType>& type, MemoryPool* pool)
+      : Base(type, pool) {}
+
+  Status Reserve(const int64_t length) {
+    counts_.reserve(length);
+    return Status::OK();
+  }
+
+  void ObserveNull() {}
+
+  void ObserveFound(const hash_slot_t slot) { counts_[slot]++; }
+
+  void ObserveNotFound(const hash_slot_t slot) { counts_.emplace_back(1); }
+
+  Status DoubleSize() { return Base::DoubleTableSize(); }
+
+  Status Flush(Datum* out) override {
+    Int64Builder builder(Base::pool_);
+    std::shared_ptr<ArrayData> result;
+
+    for (const int64_t value : counts_) {
+      RETURN_NOT_OK(builder.Append(value));
+    }
+
+    RETURN_NOT_OK(builder.FinishInternal(&result));
+    out->value = std::move(result);
+    return Status::OK();
+  }
+
+  using Base::Append;
+
+ private:
+  std::vector<int64_t> counts_;
+};
+
 // ----------------------------------------------------------------------
 // Kernel wrapper for generic hash table kernels
 
@@ -871,6 +915,48 @@ Status GetDictionaryEncodeKernel(FunctionContext* ctx,
   return Status::OK();
 }
 
+Status GetCountValuesKernel(FunctionContext* ctx, const std::shared_ptr<DataType>& type,
+                            std::unique_ptr<HashKernel>* out) {
+  std::unique_ptr<HashTable> hasher;
+
+#define COUNT_VALUES_CASE(InType)                                        \
+  case InType::type_id:                                                  \
+    hasher.reset(new CountValuesImpl<InType>(type, ctx->memory_pool())); \
+    break
+
+  switch (type->id()) {
+    COUNT_VALUES_CASE(NullType);
+    COUNT_VALUES_CASE(BooleanType);
+    COUNT_VALUES_CASE(UInt8Type);
+    COUNT_VALUES_CASE(Int8Type);
+    COUNT_VALUES_CASE(UInt16Type);
+    COUNT_VALUES_CASE(Int16Type);
+    COUNT_VALUES_CASE(UInt32Type);
+    COUNT_VALUES_CASE(Int32Type);
+    COUNT_VALUES_CASE(UInt64Type);
+    COUNT_VALUES_CASE(Int64Type);
+    COUNT_VALUES_CASE(FloatType);
+    COUNT_VALUES_CASE(DoubleType);
+    COUNT_VALUES_CASE(Date32Type);
+    COUNT_VALUES_CASE(Date64Type);
+    COUNT_VALUES_CASE(Time32Type);
+    COUNT_VALUES_CASE(Time64Type);
+    COUNT_VALUES_CASE(TimestampType);
+    COUNT_VALUES_CASE(BinaryType);
+    COUNT_VALUES_CASE(StringType);
+    COUNT_VALUES_CASE(FixedSizeBinaryType);
+    COUNT_VALUES_CASE(Decimal128Type);
+    default:
+      break;
+  }
+
+#undef COUNT_VALUES_CASE
+
+  CHECK_IMPLEMENTED(hasher, "count-values", type);
+  out->reset(new HashKernelImpl(std::move(hasher)));
+  return Status::OK();
+}
+
 namespace {
 
 Status InvokeHash(FunctionContext* ctx, HashKernel* func, const Datum& value,
@@ -918,5 +1004,18 @@ Status DictionaryEncode(FunctionContext* ctx, const Datum& value, Datum* out) {
   return Status::OK();
 }
 
+Status CountValues(FunctionContext* ctx, const Datum& value,
+                   std::shared_ptr<Array>* out_uniques,
+                   std::shared_ptr<Array>* out_counts) {
+  std::unique_ptr<HashKernel> func;
+  RETURN_NOT_OK(GetCountValuesKernel(ctx, value.type(), &func));
+
+  std::vector<Datum> counts_datum;
+  RETURN_NOT_OK(InvokeHash(ctx, func.get(), value, &counts_datum, out_uniques));
+
+  *out_counts = MakeArray(counts_datum.back().array());
+  return Status::OK();
+}
+
 }  // namespace compute
 }  // namespace arrow
diff --git a/cpp/src/arrow/compute/kernels/hash.h b/cpp/src/arrow/compute/kernels/hash.h
@@ -51,6 +51,10 @@ Status GetDictionaryEncodeKernel(FunctionContext* ctx,
                                  const std::shared_ptr<DataType>& type,
                                  std::unique_ptr<HashKernel>* kernel);
 
+ARROW_EXPORT
+Status GetCountValuesKernel(FunctionContext* ctx, const std::shared_ptr<DataType>& type,
+                            std::unique_ptr<HashKernel>* kernel);
+
 /// \brief Compute unique elements from an array-like object
 /// \param[in] context the FunctionContext
 /// \param[in] datum array-like input
@@ -71,6 +75,19 @@ Status Unique(FunctionContext* context, const Datum& datum, std::shared_ptr<Arra
 ARROW_EXPORT
 Status DictionaryEncode(FunctionContext* context, const Datum& data, Datum* out);
 
+/// \brief Return counts of unique elements from an array-like object
+/// \param[in] context the FunctionContext
+/// \param[in] value array-like input
+/// \param[out] out_uniques unique elements as Array
+/// \param[out] out_counts counts per element as Array, same shape as out_uniques
+///
+/// \since 0.10.0
+/// \note API not yet finalized
+ARROW_EXPORT
+Status CountValues(FunctionContext* context, const Datum& value,
+                   std::shared_ptr<Array>* out_uniques,
+                   std::shared_ptr<Array>* out_counts);
+
 // TODO(wesm): Define API for incremental dictionary encoding
 
 // TODO(wesm): Define API for regularizing DictionaryArray objects with
@@ -95,11 +112,6 @@ Status DictionaryEncode(FunctionContext* context, const Datum& data, Datum* out)
 // Status IsIn(FunctionContext* context, const Datum& values, const Datum& member_set,
 //             Datum* out);
 
-// ARROW_EXPORT
-// Status CountValues(FunctionContext* context, const Datum& values,
-//                    std::shared_ptr<Array>* out_uniques,
-//                    std::shared_ptr<Array>* out_counts);
-
 }  // namespace compute
 }  // namespace arrow