ARROW-1572: [C++] Implement "value counts" kernels for tabulating value frequencies

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);
 

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

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