You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2021/04/14 02:18:48 UTC
[GitHub] [arrow] cyb70289 commented on a change in pull request #10009: ARROW-11568: [C++][Compute] Rewrite mode kernel
cyb70289 commented on a change in pull request #10009:
URL: https://github.com/apache/arrow/pull/10009#discussion_r612891573
##########
File path: cpp/src/arrow/compute/kernels/aggregate_mode.cc
##########
@@ -31,340 +33,359 @@ namespace internal {
namespace {
+using arrow::internal::checked_pointer_cast;
+using arrow::internal::VisitSetBitRunsVoid;
+
+using ModeState = OptionsWrapper<ModeOptions>;
+
constexpr char kModeFieldName[] = "mode";
constexpr char kCountFieldName[] = "count";
-// {value:count} map
-template <typename CType>
-using CounterMap = std::unordered_map<CType, int64_t>;
-
-// map based counter for floating points
-template <typename ArrayType, typename CType = typename ArrayType::TypeClass::c_type>
-enable_if_t<std::is_floating_point<CType>::value, CounterMap<CType>> CountValuesByMap(
- const ArrayType& array, int64_t& nan_count) {
- CounterMap<CType> value_counts_map;
- const ArrayData& data = *array.data();
- const CType* values = data.GetValues<CType>(1);
-
- nan_count = 0;
- if (array.length() > array.null_count()) {
- arrow::internal::VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length,
- [&](int64_t pos, int64_t len) {
- for (int64_t i = 0; i < len; ++i) {
- const auto value = values[pos + i];
- if (std::isnan(value)) {
- ++nan_count;
- } else {
- ++value_counts_map[value];
- }
- }
- });
+template <typename InType, typename CType = typename InType::c_type>
+Result<std::pair<CType*, int64_t*>> PrepareOutput(int64_t n, KernelContext* ctx,
+ Datum* out) {
+ const auto& mode_type = TypeTraits<InType>::type_singleton();
+ const auto& count_type = int64();
+
+ auto mode_data = ArrayData::Make(mode_type, /*length=*/n, /*null_count=*/0);
+ mode_data->buffers.resize(2, nullptr);
+ auto count_data = ArrayData::Make(count_type, n, 0);
+ count_data->buffers.resize(2, nullptr);
+
+ CType* mode_buffer = nullptr;
+ int64_t* count_buffer = nullptr;
+
+ if (n > 0) {
+ ARROW_ASSIGN_OR_RAISE(mode_data->buffers[1], ctx->Allocate(n * sizeof(CType)));
+ ARROW_ASSIGN_OR_RAISE(count_data->buffers[1], ctx->Allocate(n * sizeof(int64_t)));
+ mode_buffer = mode_data->template GetMutableValues<CType>(1);
+ count_buffer = count_data->template GetMutableValues<int64_t>(1);
}
- return value_counts_map;
-}
-
-// map base counter for non floating points
-template <typename ArrayType, typename CType = typename ArrayType::TypeClass::c_type>
-enable_if_t<!std::is_floating_point<CType>::value, CounterMap<CType>> CountValuesByMap(
- const ArrayType& array) {
- CounterMap<CType> value_counts_map;
- const ArrayData& data = *array.data();
- const CType* values = data.GetValues<CType>(1);
-
- if (array.length() > array.null_count()) {
- arrow::internal::VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length,
- [&](int64_t pos, int64_t len) {
- for (int64_t i = 0; i < len; ++i) {
- ++value_counts_map[values[pos + i]];
- }
- });
- }
+ const auto& out_type =
+ struct_({field(kModeFieldName, mode_type), field(kCountFieldName, count_type)});
+ *out = Datum(ArrayData::Make(out_type, n, {nullptr}, {mode_data, count_data}, 0));
- return value_counts_map;
+ return std::make_pair(mode_buffer, count_buffer);
}
-// vector based counter for int8 or integers with small value range
-template <typename ArrayType, typename CType = typename ArrayType::TypeClass::c_type>
-CounterMap<CType> CountValuesByVector(const ArrayType& array, CType min, CType max) {
- const int range = static_cast<int>(max - min);
- DCHECK(range >= 0 && range < 64 * 1024 * 1024);
- const ArrayData& data = *array.data();
- const CType* values = data.GetValues<CType>(1);
-
- std::vector<int64_t> value_counts_vector(range + 1);
- if (array.length() > array.null_count()) {
- arrow::internal::VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length,
- [&](int64_t pos, int64_t len) {
- for (int64_t i = 0; i < len; ++i) {
- ++value_counts_vector[values[pos + i] - min];
- }
- });
- }
-
- // Transfer value counts to a map to be consistent with other chunks
- CounterMap<CType> value_counts_map(range + 1);
- for (int i = 0; i <= range; ++i) {
- CType value = static_cast<CType>(i + min);
- int64_t count = value_counts_vector[i];
- if (count) {
- value_counts_map[value] = count;
+// find top-n value:count pairs with minimal heap
+// suboptimal for tiny or large n, possibly okay as we're not in hot path
+template <typename InType, typename Generator>
+void Finalize(KernelContext* ctx, Datum* out, Generator&& gen) {
+ using CType = typename InType::c_type;
+
+ using ValueCountPair = std::pair<CType, int64_t>;
+ auto gt = [](const ValueCountPair& lhs, const ValueCountPair& rhs) {
+ const bool rhs_is_nan = rhs.first != rhs.first; // nan as largest value
+ return lhs.second > rhs.second ||
+ (lhs.second == rhs.second && (lhs.first < rhs.first || rhs_is_nan));
+ };
+
+ std::priority_queue<ValueCountPair, std::vector<ValueCountPair>, decltype(gt)> min_heap(
+ std::move(gt));
+
+ const ModeOptions& options = ModeState::Get(ctx);
+ while (true) {
+ const ValueCountPair& value_count = gen();
+ DCHECK_NE(value_count.second, 0);
+ if (value_count.second < 0) break; // EOF reached
+ if (static_cast<int64_t>(min_heap.size()) < options.n) {
+ min_heap.push(value_count);
+ } else if (gt(value_count, min_heap.top())) {
+ min_heap.pop();
+ min_heap.push(value_count);
}
}
+ const int64_t n = min_heap.size();
- return value_counts_map;
-}
-
-// map or vector based counter for int16/32/64 per value range
-template <typename ArrayType, typename CType = typename ArrayType::TypeClass::c_type>
-CounterMap<CType> CountValuesByMapOrVector(const ArrayType& array) {
- // see https://issues.apache.org/jira/browse/ARROW-9873
- static constexpr int kMinArraySize = 8192 / sizeof(CType);
- static constexpr int kMaxValueRange = 16384;
- const ArrayData& data = *array.data();
- const CType* values = data.GetValues<CType>(1);
-
- if ((array.length() - array.null_count()) >= kMinArraySize) {
- CType min = std::numeric_limits<CType>::max();
- CType max = std::numeric_limits<CType>::min();
-
- arrow::internal::VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length,
- [&](int64_t pos, int64_t len) {
- for (int64_t i = 0; i < len; ++i) {
- const auto value = values[pos + i];
- min = std::min(min, value);
- max = std::max(max, value);
- }
- });
-
- if (static_cast<uint64_t>(max) - static_cast<uint64_t>(min) <= kMaxValueRange) {
- return CountValuesByVector(array, min, max);
- }
- }
- return CountValuesByMap(array);
-}
+ CType* mode_buffer;
+ int64_t* count_buffer;
+ KERNEL_ASSIGN_OR_RAISE(std::tie(mode_buffer, count_buffer), ctx,
+ PrepareOutput<InType>(n, ctx, out));
-// bool
-template <typename ArrayType, typename CType = typename ArrayType::TypeClass::c_type>
-enable_if_t<is_boolean_type<typename ArrayType::TypeClass>::value, CounterMap<CType>>
-CountValues(const ArrayType& array, int64_t& nan_count) {
- // we need just count ones and zeros
- CounterMap<CType> map;
- if (array.length() > array.null_count()) {
- map[true] = array.true_count();
- map[false] = array.length() - array.null_count() - map[true];
+ for (int64_t i = n - 1; i >= 0; --i) {
+ std::tie(mode_buffer[i], count_buffer[i]) = min_heap.top();
+ min_heap.pop();
}
- nan_count = 0;
- return map;
}
-// int8
-template <typename ArrayType, typename CType = typename ArrayType::TypeClass::c_type>
-enable_if_t<is_integer_type<typename ArrayType::TypeClass>::value && sizeof(CType) == 1,
- CounterMap<CType>>
-CountValues(const ArrayType& array, int64_t& nan_count) {
- using Limits = std::numeric_limits<CType>;
- nan_count = 0;
- return CountValuesByVector(array, Limits::min(), Limits::max());
-}
+// count value occurances for integers with narrow value range
+// O(1) space, O(n) time
+template <typename T>
+struct CountModer {
+ using CType = typename T::c_type;
-// int16/32/64
-template <typename ArrayType, typename CType = typename ArrayType::TypeClass::c_type>
-enable_if_t<is_integer_type<typename ArrayType::TypeClass>::value && (sizeof(CType) > 1),
- CounterMap<CType>>
-CountValues(const ArrayType& array, int64_t& nan_count) {
- nan_count = 0;
- return CountValuesByMapOrVector(array);
-}
+ CType min;
+ std::vector<int64_t> counts;
-// float/double
-template <typename ArrayType, typename CType = typename ArrayType::TypeClass::c_type>
-enable_if_t<(std::is_floating_point<CType>::value), CounterMap<CType>> // NOLINT format
-CountValues(const ArrayType& array, int64_t& nan_count) {
- nan_count = 0;
- return CountValuesByMap(array, nan_count);
-}
+ CountModer(CType min, CType max) {
+ uint32_t value_range = static_cast<uint32_t>(max - min) + 1;
+ DCHECK_LT(value_range, 1 << 20);
+ this->min = min;
+ this->counts.resize(value_range, 0);
+ }
-template <typename ArrowType>
-struct ModeState {
- using ThisType = ModeState<ArrowType>;
- using CType = typename ArrowType::c_type;
-
- void MergeFrom(ThisType&& state) {
- if (this->value_counts.empty()) {
- this->value_counts = std::move(state.value_counts);
- } else {
- for (const auto& value_count : state.value_counts) {
- auto value = value_count.first;
- auto count = value_count.second;
- this->value_counts[value] += count;
+ void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+ // count values in all chunks, ignore nulls
+ const Datum& datum = batch[0];
+ const int64_t in_length = datum.length() - datum.null_count();
+ if (in_length > 0) {
+ for (const auto& array : datum.chunks()) {
+ const ArrayData& data = *array->data();
+ const CType* values = data.GetValues<CType>(1);
+ VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length,
+ [&](int64_t pos, int64_t len) {
+ for (int64_t i = 0; i < len; ++i) {
+ ++this->counts[values[pos + i] - this->min];
+ }
+ });
}
}
- if (is_floating_type<ArrowType>::value) {
- this->nan_count += state.nan_count;
- }
- }
-
- // find top-n value/count pairs with min-heap (priority queue with '>' comparator)
- void Finalize(CType* modes, int64_t* counts, const int64_t n) {
- DCHECK(n >= 1 && n <= this->DistinctValues());
- // mode 'greater than' comparator: larger count or same count with smaller value
- using ValueCountPair = std::pair<CType, int64_t>;
- auto mode_gt = [](const ValueCountPair& lhs, const ValueCountPair& rhs) {
- const bool rhs_is_nan = rhs.first != rhs.first; // nan as largest value
- return lhs.second > rhs.second ||
- (lhs.second == rhs.second && (lhs.first < rhs.first || rhs_is_nan));
+ // generator to emit next value:count pair
+ int index = 0;
+ auto gen = [&]() {
+ for (; index < static_cast<int>(counts.size()); ++index) {
+ if (counts[index] != 0) {
+ auto value_count =
+ std::make_pair(static_cast<CType>(index + this->min), counts[index]);
+ ++index;
+ return value_count;
+ }
+ }
+ return std::make_pair<CType, int64_t>(0, -1); // EOF
};
- // initialize min-heap with first n modes
- std::vector<ValueCountPair> vector(n);
- // push nan if exists
- const bool has_nan = is_floating_type<ArrowType>::value && this->nan_count > 0;
- if (has_nan) {
- vector[0] = std::make_pair(static_cast<CType>(NAN), this->nan_count);
- }
- // push n or n-1 modes
- auto it = this->value_counts.cbegin();
- for (int i = has_nan; i < n; ++i) {
- vector[i] = *it++;
- }
- // turn to min-heap
- std::priority_queue<ValueCountPair, std::vector<ValueCountPair>, decltype(mode_gt)>
- min_heap(std::move(mode_gt), std::move(vector));
-
- // iterate and insert modes into min-heap
- // - mode < heap top: ignore mode
- // - mode > heap top: discard heap top, insert mode
- for (; it != this->value_counts.cend(); ++it) {
- if (mode_gt(*it, min_heap.top())) {
- min_heap.pop();
- min_heap.push(*it);
- }
+ Finalize<T>(ctx, out, std::move(gen));
+ }
+};
+
+// booleans can be handled more straightforward
+template <>
+struct CountModer<BooleanType> {
+ void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+ int64_t counts[2]{};
+
+ const Datum& datum = batch[0];
+ for (const auto& array : datum.chunks()) {
+ if (array->length() > array->null_count()) {
+ const int64_t true_count =
+ checked_pointer_cast<BooleanArray>(array)->true_count();
+ const int64_t false_count = array->length() - array->null_count() - true_count;
+ counts[true] += true_count;
+ counts[false] += false_count;
+ };
}
- // pop modes from min-heap and insert into output array (in reverse order)
- DCHECK_EQ(min_heap.size(), static_cast<size_t>(n));
- for (int64_t i = n - 1; i >= 0; --i) {
- std::tie(modes[i], counts[i]) = min_heap.top();
- min_heap.pop();
+ const ModeOptions& options = ModeState::Get(ctx);
+ const int64_t distinct_values = (counts[0] != 0) + (counts[1] != 0);
+ const int64_t n = std::min(options.n, distinct_values);
+
+ bool* mode_buffer;
+ int64_t* count_buffer;
+ KERNEL_ASSIGN_OR_RAISE(std::tie(mode_buffer, count_buffer), ctx,
+ PrepareOutput<BooleanType>(n, ctx, out));
+
+ if (n >= 1) {
+ const bool index = counts[1] > counts[0];
+ mode_buffer[0] = index;
+ count_buffer[0] = counts[index];
+ if (n == 2) {
+ mode_buffer[1] = !index;
+ count_buffer[1] = counts[!index];
+ }
}
}
+};
- int64_t DistinctValues() const {
- return this->value_counts.size() +
- (is_floating_type<ArrowType>::value && this->nan_count > 0);
- }
+// copy and sort approach for floating points or integers with wide value range
+// O(n) space, O(nlogn) time
+template <typename T>
+struct SortModer {
+ using CType = typename T::c_type;
+ using Allocator = arrow::stl::allocator<CType>;
- int64_t nan_count = 0; // only make sense to floating types
- CounterMap<CType> value_counts;
-};
+ int64_t nan_count = 0;
-template <typename ArrowType>
-struct ModeImpl : public ScalarAggregator {
- using ThisType = ModeImpl<ArrowType>;
- using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
- using CType = typename ArrowType::c_type;
+ void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+ // copy all chunks to a buffer, ignore nulls and nans
+ std::vector<CType, Allocator> in_buffer(Allocator(ctx->memory_pool()));
- ModeImpl(const std::shared_ptr<DataType>& out_type, const ModeOptions& options)
- : out_type(out_type), options(options) {}
+ const Datum& datum = batch[0];
+ const int64_t in_length = datum.length() - datum.null_count();
+ if (in_length > 0) {
+ in_buffer.resize(in_length);
- void Consume(KernelContext*, const ExecBatch& batch) override {
- ArrayType array(batch[0].array());
- this->state.value_counts = CountValues(array, this->state.nan_count);
- }
+ int64_t index = 0;
+ for (const auto& array : datum.chunks()) {
+ index += CopyArray(in_buffer.data() + index, *array);
+ }
- void MergeFrom(KernelContext*, KernelState&& src) override {
- auto& other = checked_cast<ThisType&>(src);
- this->state.MergeFrom(std::move(other.state));
- }
+ // drop nan
+ if (is_floating_type<T>::value) {
+ const auto& it = std::remove_if(in_buffer.begin(), in_buffer.end(),
+ [](CType v) { return v != v; });
+ this->nan_count = in_buffer.end() - it;
+ in_buffer.resize(it - in_buffer.begin());
+ }
+ }
- static std::shared_ptr<ArrayData> MakeArrayData(
- const std::shared_ptr<DataType>& data_type, int64_t n) {
- auto data = ArrayData::Make(data_type, n, 0);
- data->buffers.resize(2);
- data->buffers[0] = nullptr;
- data->buffers[1] = nullptr;
- return data;
+ // sort the input data to count same values
+ std::sort(in_buffer.begin(), in_buffer.end());
+
+ // generator to emit next value:count pair
+ auto it = in_buffer.cbegin();
+ int64_t nan_count_copy = this->nan_count;
+ auto gen = [&]() {
+ if (ARROW_PREDICT_FALSE(it == in_buffer.cend())) {
+ // handle NAN at last
+ if (nan_count_copy > 0) {
+ auto value_count = std::make_pair(static_cast<CType>(NAN), nan_count_copy);
+ nan_count_copy = 0;
+ return value_count;
+ }
+ return std::make_pair<CType, int64_t>(0, -1); // EOF
+ }
+ // count same values
+ const CType value = *it;
+ int64_t count = 0;
+ do {
+ ++it;
+ ++count;
+ } while (it != in_buffer.cend() && *it == value);
+ return std::make_pair(value, count);
+ };
+
+ Finalize<T>(ctx, out, std::move(gen));
}
- void Finalize(KernelContext* ctx, Datum* out) override {
- const auto& mode_type = TypeTraits<ArrowType>::type_singleton();
- const auto& count_type = int64();
- const auto& out_type =
- struct_({field(kModeFieldName, mode_type), field(kCountFieldName, count_type)});
-
- int64_t n = this->options.n;
- if (n > state.DistinctValues()) {
- n = state.DistinctValues();
- } else if (n < 0) {
- n = 0;
+ static int64_t CopyArray(CType* buffer, const Array& array) {
+ const int64_t n = array.length() - array.null_count();
+ if (n > 0) {
+ int64_t index = 0;
+ const ArrayData& data = *array.data();
+ const CType* values = data.GetValues<CType>(1);
+ VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length,
+ [&](int64_t pos, int64_t len) {
+ memcpy(buffer + index, values + pos, len * sizeof(CType));
+ index += len;
+ });
+ DCHECK_EQ(index, n);
}
+ return n;
+ }
+};
- auto mode_data = this->MakeArrayData(mode_type, n);
- auto count_data = this->MakeArrayData(count_type, n);
- if (n > 0) {
- KERNEL_ASSIGN_OR_RAISE(mode_data->buffers[1], ctx,
- ctx->Allocate(n * sizeof(CType)));
- KERNEL_ASSIGN_OR_RAISE(count_data->buffers[1], ctx,
- ctx->Allocate(n * sizeof(int64_t)));
- CType* mode_buffer = mode_data->template GetMutableValues<CType>(1);
- int64_t* count_buffer = count_data->template GetMutableValues<int64_t>(1);
- this->state.Finalize(mode_buffer, count_buffer, n);
+// pick counting or sorting approach per integers value range
+template <typename T>
+struct CountOrSortModer {
+ using CType = typename T::c_type;
+
+ void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+ // cross point to benefit from counting approach
+ // about 2x improvement for int32/64 from micro-benchmarking
+ static constexpr int kMinArraySize = 8192;
+ static constexpr int kMaxValueRange = 32768;
+
+ const Datum& datum = batch[0];
+ if (datum.length() - datum.null_count() >= kMinArraySize) {
+ CType min = std::numeric_limits<CType>::max();
+ CType max = std::numeric_limits<CType>::min();
+
+ for (const auto& array : datum.chunks()) {
+ const ArrayData& data = *array->data();
+ const CType* values = data.GetValues<CType>(1);
+ VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length,
+ [&](int64_t pos, int64_t len) {
+ for (int64_t i = 0; i < len; ++i) {
+ min = std::min(min, values[pos + i]);
+ max = std::max(max, values[pos + i]);
+ }
+ });
+ }
+
+ if (static_cast<uint64_t>(max) - static_cast<uint64_t>(min) <= kMaxValueRange) {
+ CountModer<T>(min, max).Exec(ctx, batch, out);
+ return;
+ }
}
- *out = Datum(ArrayData::Make(out_type, n, {nullptr}, {mode_data, count_data}, 0));
+ SortModer<T>().Exec(ctx, batch, out);
}
+};
- std::shared_ptr<DataType> out_type;
- ModeState<ArrowType> state;
- ModeOptions options;
+template <typename InType, typename Enable = void>
+struct Moder;
+
+template <>
+struct Moder<Int8Type> {
+ CountModer<Int8Type> impl;
+ Moder() : impl(-128, 127) {}
};
-struct ModeInitState {
- std::unique_ptr<KernelState> state;
- KernelContext* ctx;
- const DataType& in_type;
- const std::shared_ptr<DataType>& out_type;
- const ModeOptions& options;
+template <>
+struct Moder<UInt8Type> {
+ CountModer<UInt8Type> impl;
+ Moder() : impl(0, 255) {}
+};
- ModeInitState(KernelContext* ctx, const DataType& in_type,
- const std::shared_ptr<DataType>& out_type, const ModeOptions& options)
- : ctx(ctx), in_type(in_type), out_type(out_type), options(options) {}
+template <>
+struct Moder<BooleanType> {
+ CountModer<BooleanType> impl;
+};
- Status Visit(const DataType&) { return Status::NotImplemented("No mode implemented"); }
+template <typename InType>
+struct Moder<InType, enable_if_t<(is_integer_type<InType>::value &&
+ (sizeof(typename InType::c_type) > 1))>> {
+ CountOrSortModer<InType> impl;
+};
- Status Visit(const HalfFloatType&) {
- return Status::NotImplemented("No mode implemented");
- }
+template <typename InType>
+struct Moder<InType, enable_if_t<is_floating_type<InType>::value>> {
+ SortModer<InType> impl;
+};
- template <typename Type>
- enable_if_t<is_number_type<Type>::value || is_boolean_type<Type>::value, Status> Visit(
- const Type&) {
- state.reset(new ModeImpl<Type>(out_type, options));
- return Status::OK();
- }
+template <typename _, typename InType>
+struct ModeExecutor {
+ static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+ if (ctx->state() == nullptr) {
+ ctx->SetStatus(Status::Invalid("Mode requires ModeOptions"));
+ return;
+ }
+ const ModeOptions& options = ModeState::Get(ctx);
+ if (options.n <= 0) {
+ ctx->SetStatus(Status::Invalid("ModeOption::n must be positive"));
Review comment:
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org