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/07/06 06:27:32 UTC
[GitHub] [arrow] michalursa opened a new pull request #10660: ARROW-12759: [C++][Compute] Add ExecNode for group by
michalursa opened a new pull request #10660:
URL: https://github.com/apache/arrow/pull/10660
Adding ExecNode for grouped aggregation.
--
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.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] bkietz commented on a change in pull request #10660: ARROW-12759: [C++][Compute] Add ExecNode for group by
Posted by GitBox <gi...@apache.org>.
bkietz commented on a change in pull request #10660:
URL: https://github.com/apache/arrow/pull/10660#discussion_r668001720
##########
File path: cpp/src/arrow/compute/kernels/hash_aggregate.cc
##########
@@ -748,67 +752,128 @@ struct GrouperFastImpl : Grouper {
/// Implementations should be default constructible and perform initialization in
/// Init().
struct GroupedAggregator : KernelState {
- virtual Status Init(ExecContext*, const FunctionOptions*,
- const std::shared_ptr<DataType>&) = 0;
+ virtual Status Init(ExecContext*, const FunctionOptions*) = 0;
+
+ virtual Status Resize(int64_t new_num_groups) = 0;
virtual Status Consume(const ExecBatch& batch) = 0;
- virtual Result<Datum> Finalize() = 0;
+ virtual Status Merge(GroupedAggregator&& other, const ArrayData& group_id_mapping) = 0;
- template <typename Reserve>
- Status MaybeReserve(int64_t old_num_groups, const ExecBatch& batch,
- const Reserve& reserve) {
- int64_t new_num_groups = batch[2].scalar_as<UInt32Scalar>().value;
- if (new_num_groups <= old_num_groups) {
- return Status::OK();
- }
- return reserve(new_num_groups - old_num_groups);
- }
+ virtual Result<Datum> Finalize() = 0;
virtual std::shared_ptr<DataType> out_type() const = 0;
};
+template <typename Impl>
+Result<std::unique_ptr<KernelState>> HashAggregateInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ auto impl = ::arrow::internal::make_unique<Impl>();
+ RETURN_NOT_OK(impl->Init(ctx->exec_context(), args.options));
+ return std::move(impl);
+}
+
+HashAggregateKernel MakeKernel(InputType argument_type, KernelInit init) {
+ HashAggregateKernel kernel;
+
+ kernel.init = std::move(init);
+
+ kernel.signature = KernelSignature::Make(
+ {std::move(argument_type), InputType::Array(Type::UINT32)},
+ OutputType(
+ [](KernelContext* ctx, const std::vector<ValueDescr>&) -> Result<ValueDescr> {
+ return checked_cast<GroupedAggregator*>(ctx->state())->out_type();
+ }));
+
+ kernel.resize = [](KernelContext* ctx, int64_t num_groups) {
+ return checked_cast<GroupedAggregator*>(ctx->state())->Resize(num_groups);
+ };
+
+ kernel.consume = [](KernelContext* ctx, const ExecBatch& batch) {
+ return checked_cast<GroupedAggregator*>(ctx->state())->Consume(batch);
+ };
+
+ kernel.merge = [](KernelContext* ctx, KernelState&& other,
+ const ArrayData& group_id_mapping) {
+ return checked_cast<GroupedAggregator*>(ctx->state())
+ ->Merge(checked_cast<GroupedAggregator&&>(other), group_id_mapping);
+ };
+
+ kernel.finalize = [](KernelContext* ctx, Datum* out) {
+ ARROW_ASSIGN_OR_RAISE(*out,
+ checked_cast<GroupedAggregator*>(ctx->state())->Finalize());
+ return Status::OK();
+ };
+
+ return kernel;
+}
+
+Status AddHashAggKernels(
+ const std::vector<std::shared_ptr<DataType>>& types,
+ Result<HashAggregateKernel> make_kernel(const std::shared_ptr<DataType>&),
+ HashAggregateFunction* function) {
+ for (const auto& ty : types) {
+ ARROW_ASSIGN_OR_RAISE(auto kernel, make_kernel(ty));
+ RETURN_NOT_OK(function->AddKernel(std::move(kernel)));
+ }
+ return Status::OK();
+}
+
// ----------------------------------------------------------------------
// Count implementation
struct GroupedCountImpl : public GroupedAggregator {
- Status Init(ExecContext* ctx, const FunctionOptions* options,
- const std::shared_ptr<DataType>&) override {
+ Status Init(ExecContext* ctx, const FunctionOptions* options) override {
options_ = checked_cast<const ScalarAggregateOptions&>(*options);
counts_ = BufferBuilder(ctx->memory_pool());
return Status::OK();
}
- Status Consume(const ExecBatch& batch) override {
- RETURN_NOT_OK(MaybeReserve(num_groups_, batch, [&](int64_t added_groups) {
- num_groups_ += added_groups;
- return counts_.Append(added_groups * sizeof(int64_t), 0);
- }));
+ Status Resize(int64_t new_num_groups) override {
+ auto added_groups = new_num_groups - num_groups_;
+ num_groups_ = new_num_groups;
+ return counts_.Append(added_groups * sizeof(int64_t), 0);
+ }
+
+ Status Merge(GroupedAggregator&& raw_other,
+ const ArrayData& group_id_mapping) override {
+ auto other = checked_cast<GroupedCountImpl*>(&raw_other);
- auto group_ids = batch[1].array()->GetValues<uint32_t>(1);
- auto raw_counts = reinterpret_cast<int64_t*>(counts_.mutable_data());
+ auto counts = reinterpret_cast<int64_t*>(counts_.mutable_data());
+ auto other_counts = reinterpret_cast<const int64_t*>(other->counts_.mutable_data());
+
+ auto g = group_id_mapping.GetValues<uint32_t>(1);
+ for (int64_t other_g = 0; other_g < group_id_mapping.length; ++other_g, ++g) {
+ counts[*g] += other_counts[other_g];
+ }
+ return Status::OK();
+ }
+
+ Status Consume(const ExecBatch& batch) override {
+ auto counts = reinterpret_cast<int64_t*>(counts_.mutable_data());
const auto& input = batch[0].array();
- if (!options_.skip_nulls) {
- if (input->GetNullCount() != 0) {
- for (int64_t i = 0, input_i = input->offset; i < input->length; ++i, ++input_i) {
- auto g = group_ids[i];
- raw_counts[g] += !BitUtil::GetBit(input->buffers[0]->data(), input_i);
- }
+ if (options_.skip_nulls) {
+ auto g_begin =
+ reinterpret_cast<const uint32_t*>(batch[1].array()->buffers[1]->data());
+
+ arrow::internal::VisitSetBitRunsVoid(input->buffers[0], input->offset,
+ input->length,
+ [&](int64_t offset, int64_t length) {
+ auto g = g_begin + offset;
+ for (int64_t i = 0; i < length; ++i, ++g) {
+ counts[*g] += 1;
+ }
+ });
+ } else if (input->MayHaveNulls()) {
Review comment:
If `skip_nulls`, we want to count only valid slots. If `!skip_nulls`, we want to count only null slots- in that case if `!MayHaveNulls` then no counts will be incremented
--
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.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] bkietz commented on a change in pull request #10660: ARROW-12759: [C++][Compute] Add ExecNode for group by
Posted by GitBox <gi...@apache.org>.
bkietz commented on a change in pull request #10660:
URL: https://github.com/apache/arrow/pull/10660#discussion_r664575727
##########
File path: cpp/src/arrow/compute/exec/exec_plan.cc
##########
@@ -565,5 +565,253 @@ AsyncGenerator<util::optional<ExecBatch>> MakeSinkNode(ExecNode* input,
return out;
}
+struct GroupByNode : ExecNode {
+ GroupByNode(ExecNode* input, std::string label, std::shared_ptr<Schema> output_schema,
+ ExecContext* ctx, const std::vector<int>&& key_field_ids,
+ std::unique_ptr<internal::Grouper>&& grouper,
+ const std::vector<int>&& agg_src_field_ids,
+ const std::vector<const HashAggregateKernel*>&& agg_kernels,
+ std::vector<std::unique_ptr<KernelState>>&& agg_states)
+ : ExecNode(input->plan(), std::move(label), {input}, {"groupby"},
+ std::move(output_schema), /*num_outputs=*/1),
+ ctx_(ctx),
+ key_field_ids_(std::move(key_field_ids)),
+ grouper_(std::move(grouper)),
+ agg_src_field_ids_(std::move(agg_src_field_ids)),
+ agg_kernels_(std::move(agg_kernels)),
+ agg_states_(std::move(agg_states)) {}
+
+ const char* kind_name() override { return "GroupByNode"; }
+
+ Status ProcessInputBatch(const ExecBatch& batch) {
+ // Create a batch with key columns
+ std::vector<Datum> keys(key_field_ids_.size());
+ for (size_t i = 0; i < key_field_ids_.size(); ++i) {
+ keys[i] = batch.values[key_field_ids_[i]];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch key_batch, ExecBatch::Make(keys));
+
+ // Create a batch with group ids
+ ARROW_ASSIGN_OR_RAISE(Datum id_batch, grouper_->Consume(key_batch));
+
+ // Execute aggregate kernels
+ auto num_groups = grouper_->num_groups();
+ for (size_t i = 0; i < agg_kernels_.size(); ++i) {
+ KernelContext kernel_ctx{ctx_};
+ kernel_ctx.SetState(agg_states_[i].get());
+ ARROW_ASSIGN_OR_RAISE(
+ auto agg_batch, ExecBatch::Make({batch.values[agg_src_field_ids_[i]], id_batch,
+ Datum(num_groups)}));
+ RETURN_NOT_OK(agg_kernels_[i]->consume(&kernel_ctx, agg_batch));
+ }
+
+ return Status::OK();
+ }
+
+ Status OutputResult() {
+ // Finalize output
+ ArrayDataVector out_data(agg_kernels_.size() + key_field_ids_.size());
+ auto it = out_data.begin();
+
+ for (size_t i = 0; i < agg_kernels_.size(); ++i) {
+ KernelContext batch_ctx{ctx_};
+ batch_ctx.SetState(agg_states_[i].get());
+ Datum out;
+ RETURN_NOT_OK(agg_kernels_[i]->finalize(&batch_ctx, &out));
+ *it++ = out.array();
+ }
+
+ ARROW_ASSIGN_OR_RAISE(ExecBatch out_keys, grouper_->GetUniques());
+ for (const auto& key : out_keys.values) {
+ *it++ = key.array();
+ }
+
+ uint32_t num_groups = grouper_->num_groups();
+ int num_result_batches = (num_groups + output_batch_size_ - 1) / output_batch_size_;
+ outputs_[0]->InputFinished(this, num_result_batches);
+
+ for (int i = 0; i < num_result_batches; ++i) {
+ // Check finished flag
+ if (finished_) {
Review comment:
```suggestion
if (stop_requested_) {
```
##########
File path: cpp/src/arrow/compute/exec/exec_plan.cc
##########
@@ -565,5 +565,253 @@ AsyncGenerator<util::optional<ExecBatch>> MakeSinkNode(ExecNode* input,
return out;
}
+struct GroupByNode : ExecNode {
+ GroupByNode(ExecNode* input, std::string label, std::shared_ptr<Schema> output_schema,
+ ExecContext* ctx, const std::vector<int>&& key_field_ids,
+ std::unique_ptr<internal::Grouper>&& grouper,
+ const std::vector<int>&& agg_src_field_ids,
+ const std::vector<const HashAggregateKernel*>&& agg_kernels,
+ std::vector<std::unique_ptr<KernelState>>&& agg_states)
+ : ExecNode(input->plan(), std::move(label), {input}, {"groupby"},
+ std::move(output_schema), /*num_outputs=*/1),
+ ctx_(ctx),
+ key_field_ids_(std::move(key_field_ids)),
+ grouper_(std::move(grouper)),
+ agg_src_field_ids_(std::move(agg_src_field_ids)),
+ agg_kernels_(std::move(agg_kernels)),
+ agg_states_(std::move(agg_states)) {}
+
+ const char* kind_name() override { return "GroupByNode"; }
+
+ Status ProcessInputBatch(const ExecBatch& batch) {
+ // Create a batch with key columns
+ std::vector<Datum> keys(key_field_ids_.size());
+ for (size_t i = 0; i < key_field_ids_.size(); ++i) {
+ keys[i] = batch.values[key_field_ids_[i]];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch key_batch, ExecBatch::Make(keys));
+
+ // Create a batch with group ids
+ ARROW_ASSIGN_OR_RAISE(Datum id_batch, grouper_->Consume(key_batch));
+
+ // Execute aggregate kernels
+ auto num_groups = grouper_->num_groups();
+ for (size_t i = 0; i < agg_kernels_.size(); ++i) {
+ KernelContext kernel_ctx{ctx_};
+ kernel_ctx.SetState(agg_states_[i].get());
+ ARROW_ASSIGN_OR_RAISE(
+ auto agg_batch, ExecBatch::Make({batch.values[agg_src_field_ids_[i]], id_batch,
+ Datum(num_groups)}));
+ RETURN_NOT_OK(agg_kernels_[i]->consume(&kernel_ctx, agg_batch));
+ }
+
+ return Status::OK();
+ }
+
+ Status OutputResult() {
+ // Finalize output
+ ArrayDataVector out_data(agg_kernels_.size() + key_field_ids_.size());
+ auto it = out_data.begin();
+
+ for (size_t i = 0; i < agg_kernels_.size(); ++i) {
+ KernelContext batch_ctx{ctx_};
+ batch_ctx.SetState(agg_states_[i].get());
+ Datum out;
+ RETURN_NOT_OK(agg_kernels_[i]->finalize(&batch_ctx, &out));
+ *it++ = out.array();
+ }
+
+ ARROW_ASSIGN_OR_RAISE(ExecBatch out_keys, grouper_->GetUniques());
+ for (const auto& key : out_keys.values) {
+ *it++ = key.array();
+ }
+
+ uint32_t num_groups = grouper_->num_groups();
+ int num_result_batches = (num_groups + output_batch_size_ - 1) / output_batch_size_;
+ outputs_[0]->InputFinished(this, num_result_batches);
+
+ for (int i = 0; i < num_result_batches; ++i) {
Review comment:
Eventually this will need to re-seed parallelism by sumbitting the calls to `outputs_[0]->InputReceived` to the Executor
##########
File path: cpp/src/arrow/compute/exec/exec_plan.cc
##########
@@ -565,5 +565,253 @@ AsyncGenerator<util::optional<ExecBatch>> MakeSinkNode(ExecNode* input,
return out;
}
+struct GroupByNode : ExecNode {
+ GroupByNode(ExecNode* input, std::string label, std::shared_ptr<Schema> output_schema,
+ ExecContext* ctx, const std::vector<int>&& key_field_ids,
+ std::unique_ptr<internal::Grouper>&& grouper,
+ const std::vector<int>&& agg_src_field_ids,
+ const std::vector<const HashAggregateKernel*>&& agg_kernels,
+ std::vector<std::unique_ptr<KernelState>>&& agg_states)
+ : ExecNode(input->plan(), std::move(label), {input}, {"groupby"},
+ std::move(output_schema), /*num_outputs=*/1),
+ ctx_(ctx),
+ key_field_ids_(std::move(key_field_ids)),
+ grouper_(std::move(grouper)),
+ agg_src_field_ids_(std::move(agg_src_field_ids)),
+ agg_kernels_(std::move(agg_kernels)),
+ agg_states_(std::move(agg_states)) {}
+
+ const char* kind_name() override { return "GroupByNode"; }
+
+ Status ProcessInputBatch(const ExecBatch& batch) {
+ // Create a batch with key columns
+ std::vector<Datum> keys(key_field_ids_.size());
+ for (size_t i = 0; i < key_field_ids_.size(); ++i) {
+ keys[i] = batch.values[key_field_ids_[i]];
+ }
+ ARROW_ASSIGN_OR_RAISE(ExecBatch key_batch, ExecBatch::Make(keys));
+
+ // Create a batch with group ids
+ ARROW_ASSIGN_OR_RAISE(Datum id_batch, grouper_->Consume(key_batch));
+
+ // Execute aggregate kernels
+ auto num_groups = grouper_->num_groups();
+ for (size_t i = 0; i < agg_kernels_.size(); ++i) {
+ KernelContext kernel_ctx{ctx_};
+ kernel_ctx.SetState(agg_states_[i].get());
+ ARROW_ASSIGN_OR_RAISE(
+ auto agg_batch, ExecBatch::Make({batch.values[agg_src_field_ids_[i]], id_batch,
+ Datum(num_groups)}));
+ RETURN_NOT_OK(agg_kernels_[i]->consume(&kernel_ctx, agg_batch));
+ }
+
+ return Status::OK();
+ }
+
+ Status OutputResult() {
+ // Finalize output
+ ArrayDataVector out_data(agg_kernels_.size() + key_field_ids_.size());
+ auto it = out_data.begin();
+
+ for (size_t i = 0; i < agg_kernels_.size(); ++i) {
+ KernelContext batch_ctx{ctx_};
+ batch_ctx.SetState(agg_states_[i].get());
+ Datum out;
+ RETURN_NOT_OK(agg_kernels_[i]->finalize(&batch_ctx, &out));
+ *it++ = out.array();
+ }
+
+ ARROW_ASSIGN_OR_RAISE(ExecBatch out_keys, grouper_->GetUniques());
+ for (const auto& key : out_keys.values) {
+ *it++ = key.array();
+ }
+
+ uint32_t num_groups = grouper_->num_groups();
+ int num_result_batches = (num_groups + output_batch_size_ - 1) / output_batch_size_;
+ outputs_[0]->InputFinished(this, num_result_batches);
+
+ for (int i = 0; i < num_result_batches; ++i) {
+ // Check finished flag
+ if (finished_) {
+ break;
+ }
+
+ // Slice arrays
+ int64_t batch_start = i * output_batch_size_;
+ int64_t batch_length =
+ std::min(output_batch_size_, static_cast<int>(num_groups - batch_start));
+ std::vector<Datum> output_slices(out_data.size());
+ for (size_t out_field_id = 0; out_field_id < out_data.size(); ++out_field_id) {
+ output_slices[out_field_id] =
+ out_data[out_field_id]->Slice(batch_start, batch_length);
+ }
+
+ ARROW_ASSIGN_OR_RAISE(ExecBatch output_batch, ExecBatch::Make(output_slices));
+ outputs_[0]->InputReceived(this, i, output_batch);
+ }
+
+ return Status::OK();
+ }
+
+ void InputReceived(ExecNode* input, int seq, ExecBatch batch) override {
+ DCHECK_EQ(input, inputs_[0]);
+
+ std::unique_lock<std::mutex> lock(mutex_);
Review comment:
```suggestion
std::unique_lock<std::mutex> lock(mutex_);
if (stop_requested_) {
return;
}
```
##########
File path: cpp/src/arrow/compute/exec/exec_plan.h
##########
@@ -264,5 +265,13 @@ ARROW_EXPORT
Result<ExecNode*> MakeProjectNode(ExecNode* input, std::string label,
std::vector<Expression> exprs);
+/// \brief Make a node which groups input rows based on key fields and computes
+/// aggregates for each group
+Result<ExecNode*> MakeGroupByNode(ExecNode* input, std::string label,
+ std::vector<std::string> keys,
Review comment:
```suggestion
std::vector<FieldRef> keys,
```
--
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.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] bkietz commented on pull request #10660: ARROW-12759: [C++][Compute] Add ExecNode for group by
Posted by GitBox <gi...@apache.org>.
bkietz commented on pull request #10660:
URL: https://github.com/apache/arrow/pull/10660#issuecomment-885331436
@pitrou care to take a look?
--
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.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] github-actions[bot] commented on pull request #10660: ARROW-12759: [C++][Compute] Add ExecNode for group by
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #10660:
URL: https://github.com/apache/arrow/pull/10660#issuecomment-874495739
https://issues.apache.org/jira/browse/ARROW-12759
--
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.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] bkietz closed pull request #10660: ARROW-12759: [C++][Compute] Add ExecNode for group by
Posted by GitBox <gi...@apache.org>.
bkietz closed pull request #10660:
URL: https://github.com/apache/arrow/pull/10660
--
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.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] github-actions[bot] commented on pull request #10660: ARROW-12759: [C++][Compute] Add ExecNode for group by
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #10660:
URL: https://github.com/apache/arrow/pull/10660#issuecomment-874495739
https://issues.apache.org/jira/browse/ARROW-12759
--
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.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow] lidavidm commented on a change in pull request #10660: ARROW-12759: [C++][Compute] Add ExecNode for group by
Posted by GitBox <gi...@apache.org>.
lidavidm commented on a change in pull request #10660:
URL: https://github.com/apache/arrow/pull/10660#discussion_r667957128
##########
File path: cpp/src/arrow/compute/kernels/hash_aggregate.cc
##########
@@ -748,67 +752,128 @@ struct GrouperFastImpl : Grouper {
/// Implementations should be default constructible and perform initialization in
/// Init().
struct GroupedAggregator : KernelState {
- virtual Status Init(ExecContext*, const FunctionOptions*,
- const std::shared_ptr<DataType>&) = 0;
+ virtual Status Init(ExecContext*, const FunctionOptions*) = 0;
+
+ virtual Status Resize(int64_t new_num_groups) = 0;
virtual Status Consume(const ExecBatch& batch) = 0;
- virtual Result<Datum> Finalize() = 0;
+ virtual Status Merge(GroupedAggregator&& other, const ArrayData& group_id_mapping) = 0;
- template <typename Reserve>
- Status MaybeReserve(int64_t old_num_groups, const ExecBatch& batch,
- const Reserve& reserve) {
- int64_t new_num_groups = batch[2].scalar_as<UInt32Scalar>().value;
- if (new_num_groups <= old_num_groups) {
- return Status::OK();
- }
- return reserve(new_num_groups - old_num_groups);
- }
+ virtual Result<Datum> Finalize() = 0;
virtual std::shared_ptr<DataType> out_type() const = 0;
};
+template <typename Impl>
+Result<std::unique_ptr<KernelState>> HashAggregateInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ auto impl = ::arrow::internal::make_unique<Impl>();
+ RETURN_NOT_OK(impl->Init(ctx->exec_context(), args.options));
+ return std::move(impl);
+}
+
+HashAggregateKernel MakeKernel(InputType argument_type, KernelInit init) {
+ HashAggregateKernel kernel;
+
+ kernel.init = std::move(init);
+
+ kernel.signature = KernelSignature::Make(
+ {std::move(argument_type), InputType::Array(Type::UINT32)},
+ OutputType(
+ [](KernelContext* ctx, const std::vector<ValueDescr>&) -> Result<ValueDescr> {
+ return checked_cast<GroupedAggregator*>(ctx->state())->out_type();
+ }));
+
+ kernel.resize = [](KernelContext* ctx, int64_t num_groups) {
+ return checked_cast<GroupedAggregator*>(ctx->state())->Resize(num_groups);
+ };
+
+ kernel.consume = [](KernelContext* ctx, const ExecBatch& batch) {
+ return checked_cast<GroupedAggregator*>(ctx->state())->Consume(batch);
+ };
+
+ kernel.merge = [](KernelContext* ctx, KernelState&& other,
+ const ArrayData& group_id_mapping) {
+ return checked_cast<GroupedAggregator*>(ctx->state())
+ ->Merge(checked_cast<GroupedAggregator&&>(other), group_id_mapping);
+ };
+
+ kernel.finalize = [](KernelContext* ctx, Datum* out) {
+ ARROW_ASSIGN_OR_RAISE(*out,
+ checked_cast<GroupedAggregator*>(ctx->state())->Finalize());
+ return Status::OK();
+ };
+
+ return kernel;
+}
+
+Status AddHashAggKernels(
+ const std::vector<std::shared_ptr<DataType>>& types,
+ Result<HashAggregateKernel> make_kernel(const std::shared_ptr<DataType>&),
+ HashAggregateFunction* function) {
+ for (const auto& ty : types) {
+ ARROW_ASSIGN_OR_RAISE(auto kernel, make_kernel(ty));
+ RETURN_NOT_OK(function->AddKernel(std::move(kernel)));
+ }
+ return Status::OK();
+}
+
// ----------------------------------------------------------------------
// Count implementation
struct GroupedCountImpl : public GroupedAggregator {
- Status Init(ExecContext* ctx, const FunctionOptions* options,
- const std::shared_ptr<DataType>&) override {
+ Status Init(ExecContext* ctx, const FunctionOptions* options) override {
options_ = checked_cast<const ScalarAggregateOptions&>(*options);
counts_ = BufferBuilder(ctx->memory_pool());
return Status::OK();
}
- Status Consume(const ExecBatch& batch) override {
- RETURN_NOT_OK(MaybeReserve(num_groups_, batch, [&](int64_t added_groups) {
- num_groups_ += added_groups;
- return counts_.Append(added_groups * sizeof(int64_t), 0);
- }));
+ Status Resize(int64_t new_num_groups) override {
+ auto added_groups = new_num_groups - num_groups_;
+ num_groups_ = new_num_groups;
+ return counts_.Append(added_groups * sizeof(int64_t), 0);
+ }
+
+ Status Merge(GroupedAggregator&& raw_other,
+ const ArrayData& group_id_mapping) override {
+ auto other = checked_cast<GroupedCountImpl*>(&raw_other);
- auto group_ids = batch[1].array()->GetValues<uint32_t>(1);
- auto raw_counts = reinterpret_cast<int64_t*>(counts_.mutable_data());
+ auto counts = reinterpret_cast<int64_t*>(counts_.mutable_data());
+ auto other_counts = reinterpret_cast<const int64_t*>(other->counts_.mutable_data());
+
+ auto g = group_id_mapping.GetValues<uint32_t>(1);
+ for (int64_t other_g = 0; other_g < group_id_mapping.length; ++other_g, ++g) {
+ counts[*g] += other_counts[other_g];
+ }
+ return Status::OK();
+ }
+
+ Status Consume(const ExecBatch& batch) override {
+ auto counts = reinterpret_cast<int64_t*>(counts_.mutable_data());
const auto& input = batch[0].array();
- if (!options_.skip_nulls) {
- if (input->GetNullCount() != 0) {
- for (int64_t i = 0, input_i = input->offset; i < input->length; ++i, ++input_i) {
- auto g = group_ids[i];
- raw_counts[g] += !BitUtil::GetBit(input->buffers[0]->data(), input_i);
- }
+ if (options_.skip_nulls) {
+ auto g_begin =
+ reinterpret_cast<const uint32_t*>(batch[1].array()->buffers[1]->data());
+
+ arrow::internal::VisitSetBitRunsVoid(input->buffers[0], input->offset,
+ input->length,
+ [&](int64_t offset, int64_t length) {
+ auto g = g_begin + offset;
+ for (int64_t i = 0; i < length; ++i, ++g) {
+ counts[*g] += 1;
+ }
+ });
+ } else if (input->MayHaveNulls()) {
Review comment:
Doesn't this fail to count if !skip_nulls && !MayHaveNulls?
##########
File path: cpp/src/arrow/compute/kernels/hash_aggregate.cc
##########
@@ -748,67 +752,128 @@ struct GrouperFastImpl : Grouper {
/// Implementations should be default constructible and perform initialization in
/// Init().
struct GroupedAggregator : KernelState {
- virtual Status Init(ExecContext*, const FunctionOptions*,
- const std::shared_ptr<DataType>&) = 0;
+ virtual Status Init(ExecContext*, const FunctionOptions*) = 0;
+
+ virtual Status Resize(int64_t new_num_groups) = 0;
virtual Status Consume(const ExecBatch& batch) = 0;
- virtual Result<Datum> Finalize() = 0;
+ virtual Status Merge(GroupedAggregator&& other, const ArrayData& group_id_mapping) = 0;
- template <typename Reserve>
- Status MaybeReserve(int64_t old_num_groups, const ExecBatch& batch,
- const Reserve& reserve) {
- int64_t new_num_groups = batch[2].scalar_as<UInt32Scalar>().value;
- if (new_num_groups <= old_num_groups) {
- return Status::OK();
- }
- return reserve(new_num_groups - old_num_groups);
- }
+ virtual Result<Datum> Finalize() = 0;
virtual std::shared_ptr<DataType> out_type() const = 0;
};
+template <typename Impl>
+Result<std::unique_ptr<KernelState>> HashAggregateInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ auto impl = ::arrow::internal::make_unique<Impl>();
+ RETURN_NOT_OK(impl->Init(ctx->exec_context(), args.options));
+ return std::move(impl);
+}
+
+HashAggregateKernel MakeKernel(InputType argument_type, KernelInit init) {
+ HashAggregateKernel kernel;
+
+ kernel.init = std::move(init);
+
+ kernel.signature = KernelSignature::Make(
+ {std::move(argument_type), InputType::Array(Type::UINT32)},
+ OutputType(
+ [](KernelContext* ctx, const std::vector<ValueDescr>&) -> Result<ValueDescr> {
+ return checked_cast<GroupedAggregator*>(ctx->state())->out_type();
+ }));
+
+ kernel.resize = [](KernelContext* ctx, int64_t num_groups) {
+ return checked_cast<GroupedAggregator*>(ctx->state())->Resize(num_groups);
+ };
+
+ kernel.consume = [](KernelContext* ctx, const ExecBatch& batch) {
+ return checked_cast<GroupedAggregator*>(ctx->state())->Consume(batch);
+ };
+
+ kernel.merge = [](KernelContext* ctx, KernelState&& other,
+ const ArrayData& group_id_mapping) {
+ return checked_cast<GroupedAggregator*>(ctx->state())
+ ->Merge(checked_cast<GroupedAggregator&&>(other), group_id_mapping);
+ };
+
+ kernel.finalize = [](KernelContext* ctx, Datum* out) {
+ ARROW_ASSIGN_OR_RAISE(*out,
+ checked_cast<GroupedAggregator*>(ctx->state())->Finalize());
+ return Status::OK();
+ };
+
+ return kernel;
+}
+
+Status AddHashAggKernels(
+ const std::vector<std::shared_ptr<DataType>>& types,
+ Result<HashAggregateKernel> make_kernel(const std::shared_ptr<DataType>&),
+ HashAggregateFunction* function) {
+ for (const auto& ty : types) {
+ ARROW_ASSIGN_OR_RAISE(auto kernel, make_kernel(ty));
+ RETURN_NOT_OK(function->AddKernel(std::move(kernel)));
+ }
+ return Status::OK();
+}
+
// ----------------------------------------------------------------------
// Count implementation
struct GroupedCountImpl : public GroupedAggregator {
- Status Init(ExecContext* ctx, const FunctionOptions* options,
- const std::shared_ptr<DataType>&) override {
+ Status Init(ExecContext* ctx, const FunctionOptions* options) override {
options_ = checked_cast<const ScalarAggregateOptions&>(*options);
counts_ = BufferBuilder(ctx->memory_pool());
return Status::OK();
}
- Status Consume(const ExecBatch& batch) override {
- RETURN_NOT_OK(MaybeReserve(num_groups_, batch, [&](int64_t added_groups) {
- num_groups_ += added_groups;
- return counts_.Append(added_groups * sizeof(int64_t), 0);
- }));
+ Status Resize(int64_t new_num_groups) override {
+ auto added_groups = new_num_groups - num_groups_;
+ num_groups_ = new_num_groups;
+ return counts_.Append(added_groups * sizeof(int64_t), 0);
+ }
+
+ Status Merge(GroupedAggregator&& raw_other,
+ const ArrayData& group_id_mapping) override {
+ auto other = checked_cast<GroupedCountImpl*>(&raw_other);
- auto group_ids = batch[1].array()->GetValues<uint32_t>(1);
- auto raw_counts = reinterpret_cast<int64_t*>(counts_.mutable_data());
+ auto counts = reinterpret_cast<int64_t*>(counts_.mutable_data());
+ auto other_counts = reinterpret_cast<const int64_t*>(other->counts_.mutable_data());
+
+ auto g = group_id_mapping.GetValues<uint32_t>(1);
+ for (int64_t other_g = 0; other_g < group_id_mapping.length; ++other_g, ++g) {
+ counts[*g] += other_counts[other_g];
+ }
+ return Status::OK();
+ }
+
+ Status Consume(const ExecBatch& batch) override {
+ auto counts = reinterpret_cast<int64_t*>(counts_.mutable_data());
const auto& input = batch[0].array();
- if (!options_.skip_nulls) {
- if (input->GetNullCount() != 0) {
- for (int64_t i = 0, input_i = input->offset; i < input->length; ++i, ++input_i) {
- auto g = group_ids[i];
- raw_counts[g] += !BitUtil::GetBit(input->buffers[0]->data(), input_i);
- }
+ if (options_.skip_nulls) {
+ auto g_begin =
+ reinterpret_cast<const uint32_t*>(batch[1].array()->buffers[1]->data());
+
+ arrow::internal::VisitSetBitRunsVoid(input->buffers[0], input->offset,
+ input->length,
+ [&](int64_t offset, int64_t length) {
+ auto g = g_begin + offset;
+ for (int64_t i = 0; i < length; ++i, ++g) {
+ counts[*g] += 1;
+ }
+ });
+ } else if (input->MayHaveNulls()) {
Review comment:
Ah, sorry, I misunderstood what skip_nulls means. Thanks.
--
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.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org