You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "pitrou (via GitHub)" <gi...@apache.org> on 2025/02/18 12:50:00 UTC
[PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
pitrou opened a new pull request, #45562:
URL: https://github.com/apache/arrow/pull/45562
### Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
### What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
### Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
### Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please uncomment the line below and explain which changes are breaking.
-->
<!-- **This PR includes breaking changes to public APIs.** -->
<!--
Please uncomment the line below (and provide explanation) if the changes fix either (a) a security vulnerability, (b) a bug that caused incorrect or invalid data to be produced, or (c) a bug that causes a crash (even when the API contract is upheld). We use this to highlight fixes to issues that may affect users without their knowledge. For this reason, fixing bugs that cause errors don't count, since those are usually obvious.
-->
<!-- **This PR contains a "Critical Fix".** -->
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1977557439
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,186 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (batch[0].is_array()) {
+ ARROW_ASSIGN_OR_RAISE(span<const PivotWiderKeyIndex> keys,
+ key_mapper_->MapKeys(batch[0].array));
+ if (batch[1].is_array()) {
+ // Array keys, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Array keys, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ } else {
+ ARROW_ASSIGN_OR_RAISE(PivotWiderKeyIndex key,
+ key_mapper_->MapKey(*batch[0].scalar));
+ if (key != kNullPivotKey) {
+ if (batch[1].is_array()) {
+ // Scalar key, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ if (!values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Scalar key, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ if (batch.length > 1 || values_[key]->is_valid) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeFrom(KernelContext*, KernelState&& src) override {
+ const auto& other_state = checked_cast<const PivotImpl&>(src);
+ for (int64_t key = 0; key < static_cast<int64_t>(values_.size()); ++key) {
+ if (other_state.values_[key]->is_valid) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = other_state.values_[key];
+ }
+ }
+ return Status::OK();
+ }
+
+ Status Finalize(KernelContext* ctx, Datum* out) override {
+ *out = std::make_shared<StructScalar>(std::move(values_), out_type_);
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same pivot key");
+ }
+
+ std::shared_ptr<DataType> out_type() const { return out_type_; }
+
+ std::shared_ptr<DataType> key_type_;
+ std::shared_ptr<DataType> out_type_;
+ const PivotWiderOptions* options_;
+ std::unique_ptr<PivotWiderKeyMapper> key_mapper_;
+ ScalarVector values_;
+};
+
+Result<std::unique_ptr<KernelState>> PivotInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ const auto& options = checked_cast<const PivotWiderOptions&>(*args.options);
+ DCHECK_EQ(args.inputs.size(), 2);
+ DCHECK(is_base_binary_like(args.inputs[0].id()));
+ auto state = std::make_unique<PivotImpl>();
+ RETURN_NOT_OK(state->Init(options, args.inputs));
+ return state;
+}
+
+Result<TypeHolder> ResolveOutputType(KernelContext* ctx, const std::vector<TypeHolder>&) {
+ return checked_cast<PivotImpl*>(ctx->state())->out_type();
+}
+
+const FunctionDoc pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear, null is emitted.\n"
+ "If a pivot key appears twice, KeyError is raised.\n"
+ "Behavior of unexpected pivot keys is controlled by PivotWiderOptions."),
Review Comment:
I think this is covered by "If a pivot key appears twice, KeyError is raised"?
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1979743935
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 1]
+ //
+ // Then each output column is computed by taking the values with the
+ // respective take_indices for the column's keys.
+ //
+
+ DCHECK_EQ(groups.size(), keys.size());
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ std::shared_ptr<DataType> take_index_type;
+ std::vector<std::shared_ptr<Buffer>> take_indices(num_keys_);
+ std::vector<std::shared_ptr<Buffer>> take_bitmaps(num_keys_);
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size =
+ num_keys_ * (take_indices_size + take_bitmap_size);
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ // Slice the scratch space into individual buffers for each output column's
+ // take_indices array.
+ std::vector<TakeIndex*> take_indices_data(num_keys_);
+ std::vector<uint8_t*> take_bitmap_data(num_keys_);
+ int64_t offset = 0;
+ for (int i = 0; i < num_keys_; ++i) {
+ take_indices[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_indices_size);
+ take_indices_data[i] = take_indices[i]->mutable_data_as<TakeIndex>();
+ offset += take_indices_size;
+ take_bitmaps[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_bitmap_size);
+ take_bitmap_data[i] = take_bitmaps[i]->mutable_data();
+ memset(take_bitmap_data[i], 0, take_bitmap_size);
+ offset += take_bitmap_size;
+ }
+ DCHECK_LE(offset, scratch_buffer_.capacity());
+
+ // Populate the take_indices for each output column
+ for (int64_t i = 0; i < values.length; ++i) {
+ const PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values.IsNull(i)) {
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ const uint32_t group = groups[i];
+ if (bit_util::GetBit(take_bitmap_data[key], group)) {
+ return DuplicateValue();
+ }
+ // For row #group in column #key, we are going to take the value at index #i
+ bit_util::SetBit(take_bitmap_data[key], group);
+ take_indices_data[key][group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to compute the output columns for this batch
+ auto values_data = values.ToArrayData();
+ ArrayVector new_columns(num_keys_);
+ for (int i = 0; i < num_keys_; ++i) {
+ auto indices_data =
+ ArrayData::Make(take_index_type, num_groups_,
+ {std::move(take_bitmaps[i]), std::move(take_indices[i])});
+ // If indices_data is all nulls, we can just ignore this column.
+ if (indices_data->GetNullCount() != indices_data->length) {
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column, Take(values_data, indices_data,
+ TakeOptions::Defaults(), ctx_));
+ new_columns[i] = grouped_column.make_array();
+ }
+ }
+ // Merge them with the previous columns
+ return MergeColumns(std::move(new_columns));
+ }
+
+ Status Consume(span<const uint32_t> groups, const PivotWiderKeyIndex key,
+ const ArraySpan& values) {
+ if (key == kNullPivotKey) {
+ // Nothing to update
+ return Status::OK();
+ }
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ // The algorithm is simpler than in the array-taking version of Consume()
+ // below, since only the column #key needs to be updated.
+ std::shared_ptr<DataType> take_index_type;
+ std::shared_ptr<Buffer> take_indices;
+ std::shared_ptr<Buffer> take_bitmap;
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size = take_indices_size + take_bitmap_size;
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ take_indices = std::make_shared<MutableBuffer>(scratch_buffer_.mutable_data(),
+ take_indices_size);
+ take_bitmap = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + take_indices_size, take_bitmap_size);
+ auto take_indices_data = take_indices->mutable_data_as<TakeIndex>();
+ auto take_bitmap_data = take_bitmap->mutable_data();
+ memset(take_bitmap_data, 0, take_bitmap_size);
+
+ for (int64_t i = 0; i < values.length; ++i) {
+ const uint32_t group = groups[i];
+ if (!values.IsNull(i)) {
+ if (bit_util::GetBit(take_bitmap_data, group)) {
+ return DuplicateValue();
+ }
+ bit_util::SetBit(take_bitmap_data, group);
+ take_indices_data[group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to update column #key
+ auto values_data = values.ToArrayData();
+ auto indices_data = ArrayData::Make(
+ take_index_type, num_groups_, {std::move(take_bitmap), std::move(take_indices)});
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column,
+ Take(values_data, indices_data, TakeOptions::Defaults(), ctx_));
+ return MergeColumn(&columns_[key], grouped_column.make_array());
+ }
+
+ Status Resize(int64_t new_num_groups) {
+ if (new_num_groups > std::numeric_limits<int32_t>::max()) {
+ return Status::NotImplemented("Pivot with more 2**31 groups");
+ }
+ return ResizeColumns(new_num_groups);
+ }
+
+ Status Merge(GroupedPivotAccumulator&& other, const ArrayData& group_id_mapping) {
+ // To merge `other` into `*this`, we simply merge their respective columns.
+ // However, we must first transpose `other`'s rows using `group_id_mapping`.
+ // This is a logical "scatter" operation.
+ //
+ // Since `scatter(indices)` is implemented as `take(inverse_permutation(indices))`,
+ // we can save time by computing `inverse_permutation(indices)` once for all
+ // columns.
+
+ // Scatter/InversePermutation only accept signed indices. We checked
+ // in Resize() above that we were inside the limites for int32.
+ auto scatter_indices = group_id_mapping.Copy();
+ scatter_indices->type = int32();
+ std::shared_ptr<DataType> take_indices_type;
+ if (num_groups_ - 1 <= std::numeric_limits<int8_t>::max()) {
+ take_indices_type = int8();
+ } else if (num_groups_ - 1 <= std::numeric_limits<int16_t>::max()) {
+ take_indices_type = int16();
+ } else {
+ DCHECK_GE(num_groups_ - 1, std::numeric_limits<int32_t>::max());
+ take_indices_type = int32();
+ }
+ InversePermutationOptions options(/*max_index=*/num_groups_ - 1, take_indices_type);
+ ARROW_ASSIGN_OR_RAISE(auto take_indices,
+ InversePermutation(scatter_indices, options, ctx_));
+ auto scatter_column =
+ [&](const std::shared_ptr<Array>& column) -> Result<std::shared_ptr<Array>> {
+ ARROW_ASSIGN_OR_RAISE(auto scattered,
+ CallFunction("take", {column, take_indices}, &options, ctx_));
Review Comment:
Oops, missed that actually.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2685349125
@github-actions crossbow submit -g cpp
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "icexelloss (via GitHub)" <gi...@apache.org>.
icexelloss commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1969893817
##########
cpp/src/arrow/compute/kernels/pivot_internal.cc:
##########
@@ -0,0 +1,127 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/kernels/pivot_internal.h"
+
+#include <cstdint>
+
+#include "arrow/compute/exec.h"
+#include "arrow/compute/kernels/codegen_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/type_traits.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/visit_type_inline.h"
+
+namespace arrow::compute::internal {
+
+using ::arrow::util::span;
+
+struct BasePivotKeyMapper : public PivotWiderKeyMapper {
+ Status Init(const PivotWiderOptions* options) override {
+ if (options->key_names.size() > static_cast<size_t>(kMaxPivotKey) + 1) {
+ return Status::NotImplemented("Pivoting to more than ",
+ static_cast<size_t>(kMaxPivotKey) + 1,
+ " columns: got ", options->key_names.size());
+ }
+ key_name_map_.reserve(options->key_names.size());
+ PivotWiderKeyIndex index = 0;
+ for (const auto& key_name : options->key_names) {
+ bool inserted =
+ key_name_map_.try_emplace(std::string_view(key_name), index++).second;
+ if (!inserted) {
+ return Status::KeyError("Duplicate key name '", key_name,
+ "' in PivotWiderOptions");
+ }
+ }
+ unexpected_key_behavior_ = options->unexpected_key_behavior;
+ return Status::OK();
+ }
+
+ protected:
+ Result<PivotWiderKeyIndex> KeyNotFound(std::string_view key_name) {
+ if (unexpected_key_behavior_ == PivotWiderOptions::kIgnore) {
+ return kNullPivotKey;
+ }
+ return Status::KeyError("Unexpected pivot key: ", key_name);
+ }
+
+ Result<PivotWiderKeyIndex> LookupKey(std::string_view key_name) {
+ const auto it = this->key_name_map_.find(key_name);
+ if (ARROW_PREDICT_FALSE(it == this->key_name_map_.end())) {
+ return KeyNotFound(key_name);
+ } else {
+ return it->second;
+ }
+ }
+
+ Status NullKeyName() { return Status::KeyError("pivot key name cannot be null"); }
+
+ static constexpr int kBatchLength = 512;
Review Comment:
What is this used for?
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1976821867
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
+/// the first one points to the pivot key column, the second points to the
+/// pivoted data column.
+/// - The pivot key column must be string-like; its values will be matched
+/// against `key_names` in order to dispatch the pivoted data into the
+/// output.
+///
+/// "pivot_wider" example
+/// ---------------------
+///
+/// Assuming the following two input columns with types utf8 and int16 (respectively):
+/// ```
+/// width | 11
+/// height | 13
+/// ```
+/// and the options `PivotWiderOptions(.key_names = {"height", "width"})`
+///
+/// then the output will be a scalar with the type
+/// `struct{"height": int16, "width": int16}`
+/// and the value `{"height": 13, "width": 11}`.
+///
+/// "hash_pivot_wider" example
+/// --------------------------
+///
+/// Assuming the following input with schema
+/// `{"group": int32, "key": utf8, "value": int16}`:
+/// ```
+/// group | key | value
+/// -----------------------------
+/// 1 | height | 11
+/// 1 | width | 12
+/// 2 | width | 13
+/// 3 | height | 14
+/// 3 | depth | 15
+/// ```
+/// and the following settings:
+/// - a hash grouping key "group"
+/// - Aggregate(
+/// .function = "hash_pivot_wider",
+/// .options = PivotWiderOptions(.key_names = {"height", "width"}),
+/// .target = {"key", "value"},
+/// .name = {"properties"})
+///
+/// then the output will have the schema
+/// `{"group": int32, "properties": struct{"height": int16, "width": int16}}`
+/// and the following value:
+/// ```
+/// group | properties
+/// | height | width
+/// -----------------------------
+/// 1 | 11 | 12
+/// 2 | null | 13
+/// 3 | 14 | null
+/// ```
+class ARROW_EXPORT PivotWiderOptions : public FunctionOptions {
Review Comment:
OK, fair enough.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1979740523
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,186 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (batch[0].is_array()) {
+ ARROW_ASSIGN_OR_RAISE(span<const PivotWiderKeyIndex> keys,
+ key_mapper_->MapKeys(batch[0].array));
+ if (batch[1].is_array()) {
+ // Array keys, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Array keys, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ } else {
+ ARROW_ASSIGN_OR_RAISE(PivotWiderKeyIndex key,
+ key_mapper_->MapKey(*batch[0].scalar));
+ if (key != kNullPivotKey) {
+ if (batch[1].is_array()) {
+ // Scalar key, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ if (!values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Scalar key, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ if (batch.length > 1 || values_[key]->is_valid) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeFrom(KernelContext*, KernelState&& src) override {
+ const auto& other_state = checked_cast<const PivotImpl&>(src);
+ for (int64_t key = 0; key < static_cast<int64_t>(values_.size()); ++key) {
+ if (other_state.values_[key]->is_valid) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = other_state.values_[key];
+ }
+ }
+ return Status::OK();
+ }
+
+ Status Finalize(KernelContext* ctx, Datum* out) override {
+ *out = std::make_shared<StructScalar>(std::move(values_), out_type_);
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same pivot key");
+ }
+
+ std::shared_ptr<DataType> out_type() const { return out_type_; }
+
+ std::shared_ptr<DataType> key_type_;
+ std::shared_ptr<DataType> out_type_;
+ const PivotWiderOptions* options_;
+ std::unique_ptr<PivotWiderKeyMapper> key_mapper_;
+ ScalarVector values_;
+};
+
+Result<std::unique_ptr<KernelState>> PivotInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ const auto& options = checked_cast<const PivotWiderOptions&>(*args.options);
+ DCHECK_EQ(args.inputs.size(), 2);
+ DCHECK(is_base_binary_like(args.inputs[0].id()));
+ auto state = std::make_unique<PivotImpl>();
+ RETURN_NOT_OK(state->Init(options, args.inputs));
+ return state;
+}
+
+Result<TypeHolder> ResolveOutputType(KernelContext* ctx, const std::vector<TypeHolder>&) {
+ return checked_cast<PivotImpl*>(ctx->state())->out_type();
+}
+
+const FunctionDoc pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear, null is emitted.\n"
+ "If a pivot key appears twice, KeyError is raised.\n"
+ "Behavior of unexpected pivot keys is controlled by PivotWiderOptions."),
Review Comment:
Well, not a big deal. The docstring looks good now.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2665638753
Revision: 53cff5570590c6216dda4fa5c8d50b3d6e697a20
Submitted crossbow builds: [ursacomputing/crossbow @ actions-d681b13450](https://github.com/ursacomputing/crossbow/branches/all?query=actions-d681b13450)
|Task|Status|
|----|------|
|example-cpp-minimal-build-static|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033797/job/37398583594)|
|example-cpp-minimal-build-static-system-dependency|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033220/job/37398581800)|
|example-cpp-tutorial|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033197/job/37398581134)|
|test-alpine-linux-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033278/job/37398582229)|
|test-build-cpp-fuzz|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033696/job/37398583781)|
|test-conda-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033530/job/37398583056)|
|test-conda-cpp-valgrind|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033104/job/37398581091)|
|test-cuda-cpp-ubuntu-22.04-cuda-11.7.1|[](https://github.com/ursacomputing/crossbow/actions/runs/13391032920/job/37398580426)|
|test-debian-12-cpp-amd64|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033728/job/37398583465)|
|test-debian-12-cpp-i386|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033619/job/37398583302)|
|test-fedora-39-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033424/job/37398582563)|
|test-ubuntu-22.04-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033282/job/37398582370)|
|test-ubuntu-22.04-cpp-20|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033387/job/37398582493)|
|test-ubuntu-22.04-cpp-bundled|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033662/job/37398583433)|
|test-ubuntu-22.04-cpp-emscripten|[](https://github.com/ursacomputing/crossbow/actions/runs/13391032930/job/37398580454)|
|test-ubuntu-22.04-cpp-no-threading|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033170/job/37398581117)|
|test-ubuntu-24.04-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033262/job/37398582106)|
|test-ubuntu-24.04-cpp-bundled-offline|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033192/job/37398581522)|
|test-ubuntu-24.04-cpp-gcc-13-bundled|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033142/job/37398581192)|
|test-ubuntu-24.04-cpp-gcc-14|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033429/job/37398582857)|
|test-ubuntu-24.04-cpp-minimal-with-formats|[](https://github.com/ursacomputing/crossbow/actions/runs/13391032905/job/37398580440)|
|test-ubuntu-24.04-cpp-thread-sanitizer|[](https://github.com/ursacomputing/crossbow/actions/runs/13391033816/job/37398583995)|
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2665633046
@github-actions crossbow submit -g cpp
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1971272127
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
> And from the code, the case PivotDuplicateValues is using group by thus shouldn't be executing this code path.
Update: I was wrong about this part - the test is indeed doing the scalar agg.
(But unfortunately this doesn't change rest of the my comment.)
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2684520702
@github-actions crossbow submit -g cpp
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1971748121
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,159 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (!batch[1].is_array()) {
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.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2694942577
> I don't think the scalar case for hash kernel is blocking - we can probably do as a follow up if you want to merge this first.
I'll wait for @zanmato1984 's opinion before moving forward on this.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1975726909
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
Review Comment:
Good point. Thank you.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "icexelloss (via GitHub)" <gi...@apache.org>.
icexelloss commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1969909007
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,159 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (!batch[1].is_array()) {
Review Comment:
Can we handle scalar here too? Seems pretty trivial to do (basically put the scalar value to the outputs) or expand the scalar to array.
Although this is probably practical use case, but we have seen lot of cases where user calls functions on scalar values (e.g., when you assign a literal columns in Acero it is scalar) so it is probably better UX/peace of mind to always support scalars
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1975120637
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
Review Comment:
Multiple pivot key columns are not supported at all (the pivot key must be string-like). Multiple pivot value columns can indeed be achieved using a `struct`.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1979148402
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 1]
+ //
+ // Then each output column is computed by taking the values with the
+ // respective take_indices for the column's keys.
+ //
+
+ DCHECK_EQ(groups.size(), keys.size());
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ std::shared_ptr<DataType> take_index_type;
+ std::vector<std::shared_ptr<Buffer>> take_indices(num_keys_);
+ std::vector<std::shared_ptr<Buffer>> take_bitmaps(num_keys_);
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size =
+ num_keys_ * (take_indices_size + take_bitmap_size);
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ // Slice the scratch space into individual buffers for each output column's
+ // take_indices array.
+ std::vector<TakeIndex*> take_indices_data(num_keys_);
+ std::vector<uint8_t*> take_bitmap_data(num_keys_);
+ int64_t offset = 0;
+ for (int i = 0; i < num_keys_; ++i) {
+ take_indices[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_indices_size);
+ take_indices_data[i] = take_indices[i]->mutable_data_as<TakeIndex>();
+ offset += take_indices_size;
+ take_bitmaps[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_bitmap_size);
+ take_bitmap_data[i] = take_bitmaps[i]->mutable_data();
+ memset(take_bitmap_data[i], 0, take_bitmap_size);
+ offset += take_bitmap_size;
+ }
+ DCHECK_LE(offset, scratch_buffer_.capacity());
+
+ // Populate the take_indices for each output column
+ for (int64_t i = 0; i < values.length; ++i) {
+ const PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values.IsNull(i)) {
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ const uint32_t group = groups[i];
+ if (bit_util::GetBit(take_bitmap_data[key], group)) {
+ return DuplicateValue();
+ }
+ // For row #group in column #key, we are going to take the value at index #i
+ bit_util::SetBit(take_bitmap_data[key], group);
+ take_indices_data[key][group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to compute the output columns for this batch
+ auto values_data = values.ToArrayData();
+ ArrayVector new_columns(num_keys_);
+ for (int i = 0; i < num_keys_; ++i) {
+ auto indices_data =
+ ArrayData::Make(take_index_type, num_groups_,
+ {std::move(take_bitmaps[i]), std::move(take_indices[i])});
+ // If indices_data is all nulls, we can just ignore this column.
+ if (indices_data->GetNullCount() != indices_data->length) {
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column, Take(values_data, indices_data,
+ TakeOptions::Defaults(), ctx_));
+ new_columns[i] = grouped_column.make_array();
+ }
+ }
+ // Merge them with the previous columns
+ return MergeColumns(std::move(new_columns));
+ }
+
+ Status Consume(span<const uint32_t> groups, const PivotWiderKeyIndex key,
+ const ArraySpan& values) {
+ if (key == kNullPivotKey) {
+ // Nothing to update
+ return Status::OK();
+ }
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ // The algorithm is simpler than in the array-taking version of Consume()
+ // below, since only the column #key needs to be updated.
+ std::shared_ptr<DataType> take_index_type;
+ std::shared_ptr<Buffer> take_indices;
+ std::shared_ptr<Buffer> take_bitmap;
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size = take_indices_size + take_bitmap_size;
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ take_indices = std::make_shared<MutableBuffer>(scratch_buffer_.mutable_data(),
+ take_indices_size);
+ take_bitmap = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + take_indices_size, take_bitmap_size);
+ auto take_indices_data = take_indices->mutable_data_as<TakeIndex>();
+ auto take_bitmap_data = take_bitmap->mutable_data();
+ memset(take_bitmap_data, 0, take_bitmap_size);
+
+ for (int64_t i = 0; i < values.length; ++i) {
+ const uint32_t group = groups[i];
+ if (!values.IsNull(i)) {
+ if (bit_util::GetBit(take_bitmap_data, group)) {
+ return DuplicateValue();
+ }
+ bit_util::SetBit(take_bitmap_data, group);
+ take_indices_data[group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to update column #key
+ auto values_data = values.ToArrayData();
+ auto indices_data = ArrayData::Make(
+ take_index_type, num_groups_, {std::move(take_bitmap), std::move(take_indices)});
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column,
+ Take(values_data, indices_data, TakeOptions::Defaults(), ctx_));
+ return MergeColumn(&columns_[key], grouped_column.make_array());
+ }
+
+ Status Resize(int64_t new_num_groups) {
+ if (new_num_groups > std::numeric_limits<int32_t>::max()) {
+ return Status::NotImplemented("Pivot with more 2**31 groups");
+ }
+ return ResizeColumns(new_num_groups);
+ }
+
+ Status Merge(GroupedPivotAccumulator&& other, const ArrayData& group_id_mapping) {
+ // To merge `other` into `*this`, we simply merge their respective columns.
+ // However, we must first transpose `other`'s rows using `group_id_mapping`.
+ // This is a logical "scatter" operation.
+ //
+ // Since `scatter(indices)` is implemented as `take(inverse_permutation(indices))`,
+ // we can save time by computing `inverse_permutation(indices)` once for all
+ // columns.
+
+ // Scatter/InversePermutation only accept signed indices. We checked
+ // in Resize() above that we were inside the limites for int32.
+ auto scatter_indices = group_id_mapping.Copy();
+ scatter_indices->type = int32();
+ std::shared_ptr<DataType> take_indices_type;
+ if (num_groups_ - 1 <= std::numeric_limits<int8_t>::max()) {
+ take_indices_type = int8();
+ } else if (num_groups_ - 1 <= std::numeric_limits<int16_t>::max()) {
+ take_indices_type = int16();
+ } else {
+ DCHECK_GE(num_groups_ - 1, std::numeric_limits<int32_t>::max());
+ take_indices_type = int32();
+ }
+ InversePermutationOptions options(/*max_index=*/num_groups_ - 1, take_indices_type);
+ ARROW_ASSIGN_OR_RAISE(auto take_indices,
+ InversePermutation(scatter_indices, options, ctx_));
+ auto scatter_column =
+ [&](const std::shared_ptr<Array>& column) -> Result<std::shared_ptr<Array>> {
+ ARROW_ASSIGN_OR_RAISE(auto scattered,
+ CallFunction("take", {column, take_indices}, &options, ctx_));
Review Comment:
I don't think there was a precise reason, let me try to make this more consistent :)
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,186 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (batch[0].is_array()) {
+ ARROW_ASSIGN_OR_RAISE(span<const PivotWiderKeyIndex> keys,
+ key_mapper_->MapKeys(batch[0].array));
+ if (batch[1].is_array()) {
+ // Array keys, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Array keys, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ } else {
+ ARROW_ASSIGN_OR_RAISE(PivotWiderKeyIndex key,
+ key_mapper_->MapKey(*batch[0].scalar));
+ if (key != kNullPivotKey) {
+ if (batch[1].is_array()) {
+ // Scalar key, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ if (!values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Scalar key, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ if (batch.length > 1 || values_[key]->is_valid) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeFrom(KernelContext*, KernelState&& src) override {
+ const auto& other_state = checked_cast<const PivotImpl&>(src);
+ for (int64_t key = 0; key < static_cast<int64_t>(values_.size()); ++key) {
+ if (other_state.values_[key]->is_valid) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = other_state.values_[key];
+ }
+ }
+ return Status::OK();
+ }
+
+ Status Finalize(KernelContext* ctx, Datum* out) override {
+ *out = std::make_shared<StructScalar>(std::move(values_), out_type_);
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same pivot key");
+ }
+
+ std::shared_ptr<DataType> out_type() const { return out_type_; }
+
+ std::shared_ptr<DataType> key_type_;
+ std::shared_ptr<DataType> out_type_;
+ const PivotWiderOptions* options_;
+ std::unique_ptr<PivotWiderKeyMapper> key_mapper_;
+ ScalarVector values_;
+};
+
+Result<std::unique_ptr<KernelState>> PivotInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ const auto& options = checked_cast<const PivotWiderOptions&>(*args.options);
+ DCHECK_EQ(args.inputs.size(), 2);
+ DCHECK(is_base_binary_like(args.inputs[0].id()));
+ auto state = std::make_unique<PivotImpl>();
+ RETURN_NOT_OK(state->Init(options, args.inputs));
+ return state;
+}
+
+Result<TypeHolder> ResolveOutputType(KernelContext* ctx, const std::vector<TypeHolder>&) {
+ return checked_cast<PivotImpl*>(ctx->state())->out_type();
+}
+
+const FunctionDoc pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear, null is emitted.\n"
+ "If a pivot key appears twice, KeyError is raised.\n"
+ "Behavior of unexpected pivot keys is controlled by PivotWiderOptions."),
Review Comment:
> Well, that's what I originally expected, until I saw `Invalid` not `KeyError` is actually thrown in `DuplicateValue`, so I think this description is not accurate enough.
Good point, I'll correct this.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "conbench-apache-arrow[bot] (via GitHub)" <gi...@apache.org>.
conbench-apache-arrow[bot] commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2699493614
After merging your PR, Conbench analyzed the 4 benchmarking runs that have been run so far on merge-commit b41386d0e520298551783a0313914b9f30e708be.
There were no benchmark performance regressions. 🎉
The [full Conbench report](https://github.com/apache/arrow/runs/38209003986) has more details. It also includes information about 48 possible false positives for unstable benchmarks that are known to sometimes produce them.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1970957244
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
I don't know... why a new scalar agg function makes an existing acero node having to change the internal logic?
Is there any test failing w/o this part of change? Maybe I can take a look and then tell if this change is right.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1971288561
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
Hmm, I can't seem to reproduce locally right now, though I _did_ have to change this to have the tests pass at the time. I'm gonna revert temporarily so that we can check CI.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "icexelloss (via GitHub)" <gi...@apache.org>.
icexelloss commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2682369908
Should we also add python binding for the option and call via pyarrow.compute?
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1975117433
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 2]
+ //
+ // Then each output column is computed by taking the values with the
+ // respective take_indices for the column's keys.
+ //
+
+ DCHECK_EQ(groups.size(), keys.size());
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ std::shared_ptr<DataType> take_index_type;
+ std::vector<std::shared_ptr<Buffer>> take_indices(num_keys_);
+ std::vector<std::shared_ptr<Buffer>> take_bitmaps(num_keys_);
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size =
+ num_keys_ * (take_indices_size + take_bitmap_size);
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ // Slice the scratch space into individual buffers for each output column's
+ // take_indices array.
+ std::vector<TakeIndex*> take_indices_data(num_keys_);
+ std::vector<uint8_t*> take_bitmap_data(num_keys_);
+ int64_t offset = 0;
+ for (int i = 0; i < num_keys_; ++i) {
+ take_indices[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_indices_size);
+ take_indices_data[i] = take_indices[i]->mutable_data_as<TakeIndex>();
+ offset += take_indices_size;
+ take_bitmaps[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_bitmap_size);
+ take_bitmap_data[i] = take_bitmaps[i]->mutable_data();
+ memset(take_bitmap_data[i], 0, take_bitmap_size);
+ offset += take_bitmap_size;
+ }
+ DCHECK_LE(offset, scratch_buffer_.capacity());
+
+ // Populate the take_indices for each output column
+ for (int64_t i = 0; i < values.length; ++i) {
+ const PivotWiderKeyIndex key = keys[i];
+ const uint32_t group = groups[i];
+ if (key != kNullPivotKey && !values.IsNull(i)) {
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ if (bit_util::GetBit(take_bitmap_data[key], group)) {
+ return DuplicateValue();
+ }
+ // For row #group in column #key, we are going to take the value at index #i
+ bit_util::SetBit(take_bitmap_data[key], group);
+ take_indices_data[key][group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to compute the output columns for this batch
+ auto values_data = values.ToArrayData();
+ ArrayVector new_columns(num_keys_);
+ for (int i = 0; i < num_keys_; ++i) {
+ auto indices_data =
+ ArrayData::Make(take_index_type, num_groups_,
+ {std::move(take_bitmaps[i]), std::move(take_indices[i])});
+ // If indices_data is all nulls, we can just ignore this column.
+ if (indices_data->GetNullCount() != indices_data->length) {
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column, Take(values_data, indices_data,
+ TakeOptions::Defaults(), ctx_));
+ new_columns[i] = grouped_column.make_array();
+ }
+ }
+ // Merge them with the previous columns
+ return MergeColumns(std::move(new_columns));
+ }
+
+ Status Consume(span<const uint32_t> groups, const PivotWiderKeyIndex key,
+ const ArraySpan& values) {
+ if (key == kNullPivotKey) {
+ // Nothing to update
+ return Status::OK();
+ }
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ // The algorithm is simpler than in the array-taking version of Consume()
+ // below, since only the column #key needs to be updated.
+ std::shared_ptr<DataType> take_index_type;
+ std::shared_ptr<Buffer> take_indices;
+ std::shared_ptr<Buffer> take_bitmap;
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size = take_indices_size + take_bitmap_size;
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ take_indices = std::make_shared<MutableBuffer>(scratch_buffer_.mutable_data(),
+ take_indices_size);
+ take_bitmap = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + take_indices_size, take_bitmap_size);
+ auto take_indices_data = take_indices->mutable_data_as<TakeIndex>();
+ auto take_bitmap_data = take_bitmap->mutable_data();
+ memset(take_bitmap_data, 0, take_bitmap_size);
+
+ for (int64_t i = 0; i < values.length; ++i) {
+ const uint32_t group = groups[i];
+ if (!values.IsNull(i)) {
+ if (bit_util::GetBit(take_bitmap_data, group)) {
+ return DuplicateValue();
+ }
+ bit_util::SetBit(take_bitmap_data, group);
+ take_indices_data[group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to update column #key
+ auto values_data = values.ToArrayData();
+ auto indices_data = ArrayData::Make(
+ take_index_type, num_groups_, {std::move(take_bitmap), std::move(take_indices)});
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column,
+ Take(values_data, indices_data, TakeOptions::Defaults(), ctx_));
+ return MergeColumn(&columns_[key], grouped_column.make_array());
+ }
+
+ Status Resize(int64_t new_num_groups) {
+ if (new_num_groups > std::numeric_limits<int32_t>::max()) {
+ return Status::NotImplemented("Pivot with more 2**31 groups");
+ }
+ return ResizeColumns(new_num_groups);
+ }
+
+ Status Merge(GroupedPivotAccumulator&& other, const ArrayData& group_id_mapping) {
+ // To merge `other` into `*this`, we simply merge their respective columns.
+ // However, we must first transpose `other`'s rows using `group_id_mapping`.
+ // This is a logical "scatter" operation.
+ //
+ // Since `scatter(indices)` is implemented as `take(inverse_permutation(indices))`,
+ // we can save time by computing `inverse_permutation(indices)` once for all
Review Comment:
Yes! Well, this is really doing a multi-column scatter, but scattering each column separately avoids allocating too much memory, since the scattered column is only an intermediate result before calling the coalesce function (see `MergeColumn` below).
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1970185714
##########
cpp/src/arrow/compute/kernels/pivot_internal.cc:
##########
@@ -0,0 +1,127 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/kernels/pivot_internal.h"
+
+#include <cstdint>
+
+#include "arrow/compute/exec.h"
+#include "arrow/compute/kernels/codegen_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/type_traits.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/visit_type_inline.h"
+
+namespace arrow::compute::internal {
+
+using ::arrow::util::span;
+
+struct BasePivotKeyMapper : public PivotWiderKeyMapper {
+ Status Init(const PivotWiderOptions* options) override {
+ if (options->key_names.size() > static_cast<size_t>(kMaxPivotKey) + 1) {
+ return Status::NotImplemented("Pivoting to more than ",
+ static_cast<size_t>(kMaxPivotKey) + 1,
+ " columns: got ", options->key_names.size());
+ }
+ key_name_map_.reserve(options->key_names.size());
+ PivotWiderKeyIndex index = 0;
+ for (const auto& key_name : options->key_names) {
+ bool inserted =
+ key_name_map_.try_emplace(std::string_view(key_name), index++).second;
+ if (!inserted) {
+ return Status::KeyError("Duplicate key name '", key_name,
+ "' in PivotWiderOptions");
+ }
+ }
+ unexpected_key_behavior_ = options->unexpected_key_behavior;
+ return Status::OK();
+ }
+
+ protected:
+ Result<PivotWiderKeyIndex> KeyNotFound(std::string_view key_name) {
+ if (unexpected_key_behavior_ == PivotWiderOptions::kIgnore) {
+ return kNullPivotKey;
+ }
+ return Status::KeyError("Unexpected pivot key: ", key_name);
+ }
+
+ Result<PivotWiderKeyIndex> LookupKey(std::string_view key_name) {
+ const auto it = this->key_name_map_.find(key_name);
+ if (ARROW_PREDICT_FALSE(it == this->key_name_map_.end())) {
+ return KeyNotFound(key_name);
+ } else {
+ return it->second;
+ }
+ }
+
+ Status NullKeyName() { return Status::KeyError("pivot key name cannot be null"); }
+
+ static constexpr int kBatchLength = 512;
Review Comment:
Probably a leftover, I'll remove it :)
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1970180427
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
I suppose it does not actually stop execution. @zanmato1984 might know more about this.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1971113619
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
That makes sense. Any existing one returning error in merge? Or do I have to do some code modification to reproduce the issue?
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2684527276
Revision: d5841909f20a00612c7eb82620d30122ce0fdf86
Submitted crossbow builds: [ursacomputing/crossbow @ actions-ff0e1159a1](https://github.com/ursacomputing/crossbow/branches/all?query=actions-ff0e1159a1)
|Task|Status|
|----|------|
|example-cpp-minimal-build-static|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540591/job/37843543449)|
|example-cpp-minimal-build-static-system-dependency|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540684/job/37843544054)|
|example-cpp-tutorial|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540668/job/37843544019)|
|test-alpine-linux-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540104/job/37843541646)|
|test-build-cpp-fuzz|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540403/job/37843542624)|
|test-conda-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540697/job/37843544028)|
|test-conda-cpp-meson|[](https://github.com/ursacomputing/crossbow/actions/runs/13541541379/job/37843546759)|
|test-conda-cpp-valgrind|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540829/job/37843544764)|
|test-cuda-cpp-ubuntu-22.04-cuda-11.7.1|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540882/job/37843545442)|
|test-debian-12-cpp-amd64|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540812/job/37843545455)|
|test-debian-12-cpp-i386|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540584/job/37843543448)|
|test-fedora-39-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540086/job/37843541655)|
|test-ubuntu-22.04-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540491/job/37843542711)|
|test-ubuntu-22.04-cpp-20|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540319/job/37843542631)|
|test-ubuntu-22.04-cpp-bundled|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540798/job/37843544842)|
|test-ubuntu-22.04-cpp-emscripten|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540080/job/37843541650)|
|test-ubuntu-22.04-cpp-no-threading|[](https://github.com/ursacomputing/crossbow/actions/runs/13541541273/job/37843546794)|
|test-ubuntu-24.04-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13541541360/job/37843546628)|
|test-ubuntu-24.04-cpp-bundled-offline|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540048/job/37843541870)|
|test-ubuntu-24.04-cpp-gcc-13-bundled|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540814/job/37843545419)|
|test-ubuntu-24.04-cpp-gcc-14|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540426/job/37843542648)|
|test-ubuntu-24.04-cpp-minimal-with-formats|[](https://github.com/ursacomputing/crossbow/actions/runs/13541541192/job/37843546434)|
|test-ubuntu-24.04-cpp-thread-sanitizer|[](https://github.com/ursacomputing/crossbow/actions/runs/13541540721/job/37843544844)|
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1971249009
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
I tried, no luck, sorry. And from the code, the case `PivotDuplicateValues` is using group by thus shouldn't be executing this code path. I can't speak for the validity of this change w/o truly understanding what the problem is.
I do have a suspect though: this might be race related. And if, a big if, this is the case, I'm afraid your change doesn't fix the race. However there are still many things unclear so please don't take this seriously until I have new findings.
Do you still hit the problem if you revert this part of the change? Or shall we try reverting it (maybe temporarily) to see if CI is good?
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "icexelloss (via GitHub)" <gi...@apache.org>.
icexelloss commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1977580887
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 1]
+ //
+ // Then each output column is computed by taking the values with the
+ // respective take_indices for the column's keys.
+ //
+
+ DCHECK_EQ(groups.size(), keys.size());
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ std::shared_ptr<DataType> take_index_type;
+ std::vector<std::shared_ptr<Buffer>> take_indices(num_keys_);
+ std::vector<std::shared_ptr<Buffer>> take_bitmaps(num_keys_);
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size =
+ num_keys_ * (take_indices_size + take_bitmap_size);
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ // Slice the scratch space into individual buffers for each output column's
+ // take_indices array.
+ std::vector<TakeIndex*> take_indices_data(num_keys_);
+ std::vector<uint8_t*> take_bitmap_data(num_keys_);
+ int64_t offset = 0;
+ for (int i = 0; i < num_keys_; ++i) {
+ take_indices[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_indices_size);
+ take_indices_data[i] = take_indices[i]->mutable_data_as<TakeIndex>();
+ offset += take_indices_size;
+ take_bitmaps[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_bitmap_size);
+ take_bitmap_data[i] = take_bitmaps[i]->mutable_data();
+ memset(take_bitmap_data[i], 0, take_bitmap_size);
+ offset += take_bitmap_size;
+ }
+ DCHECK_LE(offset, scratch_buffer_.capacity());
+
+ // Populate the take_indices for each output column
+ for (int64_t i = 0; i < values.length; ++i) {
+ const PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values.IsNull(i)) {
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ const uint32_t group = groups[i];
+ if (bit_util::GetBit(take_bitmap_data[key], group)) {
+ return DuplicateValue();
+ }
+ // For row #group in column #key, we are going to take the value at index #i
+ bit_util::SetBit(take_bitmap_data[key], group);
+ take_indices_data[key][group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to compute the output columns for this batch
+ auto values_data = values.ToArrayData();
+ ArrayVector new_columns(num_keys_);
+ for (int i = 0; i < num_keys_; ++i) {
+ auto indices_data =
+ ArrayData::Make(take_index_type, num_groups_,
+ {std::move(take_bitmaps[i]), std::move(take_indices[i])});
+ // If indices_data is all nulls, we can just ignore this column.
+ if (indices_data->GetNullCount() != indices_data->length) {
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column, Take(values_data, indices_data,
+ TakeOptions::Defaults(), ctx_));
+ new_columns[i] = grouped_column.make_array();
+ }
+ }
+ // Merge them with the previous columns
+ return MergeColumns(std::move(new_columns));
+ }
+
+ Status Consume(span<const uint32_t> groups, const PivotWiderKeyIndex key,
+ const ArraySpan& values) {
+ if (key == kNullPivotKey) {
+ // Nothing to update
+ return Status::OK();
+ }
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ // The algorithm is simpler than in the array-taking version of Consume()
+ // below, since only the column #key needs to be updated.
+ std::shared_ptr<DataType> take_index_type;
+ std::shared_ptr<Buffer> take_indices;
+ std::shared_ptr<Buffer> take_bitmap;
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size = take_indices_size + take_bitmap_size;
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ take_indices = std::make_shared<MutableBuffer>(scratch_buffer_.mutable_data(),
+ take_indices_size);
+ take_bitmap = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + take_indices_size, take_bitmap_size);
+ auto take_indices_data = take_indices->mutable_data_as<TakeIndex>();
+ auto take_bitmap_data = take_bitmap->mutable_data();
+ memset(take_bitmap_data, 0, take_bitmap_size);
+
+ for (int64_t i = 0; i < values.length; ++i) {
+ const uint32_t group = groups[i];
+ if (!values.IsNull(i)) {
+ if (bit_util::GetBit(take_bitmap_data, group)) {
+ return DuplicateValue();
+ }
+ bit_util::SetBit(take_bitmap_data, group);
+ take_indices_data[group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to update column #key
+ auto values_data = values.ToArrayData();
+ auto indices_data = ArrayData::Make(
+ take_index_type, num_groups_, {std::move(take_bitmap), std::move(take_indices)});
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column,
+ Take(values_data, indices_data, TakeOptions::Defaults(), ctx_));
+ return MergeColumn(&columns_[key], grouped_column.make_array());
+ }
+
+ Status Resize(int64_t new_num_groups) {
+ if (new_num_groups > std::numeric_limits<int32_t>::max()) {
+ return Status::NotImplemented("Pivot with more 2**31 groups");
+ }
+ return ResizeColumns(new_num_groups);
+ }
+
+ Status Merge(GroupedPivotAccumulator&& other, const ArrayData& group_id_mapping) {
+ // To merge `other` into `*this`, we simply merge their respective columns.
+ // However, we must first transpose `other`'s rows using `group_id_mapping`.
+ // This is a logical "scatter" operation.
+ //
+ // Since `scatter(indices)` is implemented as `take(inverse_permutation(indices))`,
+ // we can save time by computing `inverse_permutation(indices)` once for all
+ // columns.
+
+ // Scatter/InversePermutation only accept signed indices. We checked
+ // in Resize() above that we were inside the limites for int32.
+ auto scatter_indices = group_id_mapping.Copy();
+ scatter_indices->type = int32();
+ std::shared_ptr<DataType> take_indices_type;
+ if (num_groups_ - 1 <= std::numeric_limits<int8_t>::max()) {
+ take_indices_type = int8();
+ } else if (num_groups_ - 1 <= std::numeric_limits<int16_t>::max()) {
+ take_indices_type = int16();
+ } else {
+ DCHECK_GE(num_groups_ - 1, std::numeric_limits<int32_t>::max());
+ take_indices_type = int32();
+ }
+ InversePermutationOptions options(/*max_index=*/num_groups_ - 1, take_indices_type);
+ ARROW_ASSIGN_OR_RAISE(auto take_indices,
+ InversePermutation(scatter_indices, options, ctx_));
+ auto scatter_column =
+ [&](const std::shared_ptr<Array>& column) -> Result<std::shared_ptr<Array>> {
+ ARROW_ASSIGN_OR_RAISE(auto scattered,
+ CallFunction("take", {column, take_indices}, &options, ctx_));
+ return scattered.make_array();
+ };
+ return MergeColumns(std::move(other.columns_), std::move(scatter_column));
+ }
+
+ Result<ArrayVector> Finalize() {
+ // Ensure that columns are allocated even if num_groups_ == 0
+ RETURN_NOT_OK(ResizeColumns(num_groups_));
+ return std::move(columns_);
+ }
+
+ protected:
+ Status ResizeColumns(int64_t new_num_groups) {
+ if (new_num_groups == num_groups_ && num_groups_ != 0) {
+ return Status::OK();
+ }
+ ARROW_ASSIGN_OR_RAISE(
+ auto array_suffix,
+ MakeArrayOfNull(value_type_, new_num_groups - num_groups_, ctx_->memory_pool()));
+ for (auto& column : columns_) {
+ if (num_groups_ != 0) {
+ DCHECK_NE(column, nullptr);
+ ARROW_ASSIGN_OR_RAISE(
+ column, Concatenate({std::move(column), array_suffix}, ctx_->memory_pool()));
+ } else {
+ column = array_suffix;
+ }
+ DCHECK_EQ(column->length(), new_num_groups);
+ }
+ num_groups_ = new_num_groups;
+ return Status::OK();
+ }
+
+ using ColumnTransform =
+ std::function<Result<std::shared_ptr<Array>>(const std::shared_ptr<Array>&)>;
+
+ Status MergeColumns(ArrayVector&& other_columns,
+ const ColumnTransform& transform = {}) {
+ DCHECK_EQ(columns_.size(), other_columns.size());
+ for (int i = 0; i < num_keys_; ++i) {
+ if (other_columns[i]) {
+ RETURN_NOT_OK(MergeColumn(&columns_[i], std::move(other_columns[i]), transform));
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeColumn(std::shared_ptr<Array>* column, std::shared_ptr<Array> other_column,
+ const ColumnTransform& transform = {}) {
+ if (other_column->null_count() == other_column->length()) {
+ // Avoid paying for the transform step below, since merging will be a no-op anyway.
+ return Status::OK();
+ }
+ if (transform) {
+ ARROW_ASSIGN_OR_RAISE(other_column, transform(other_column));
+ }
+ DCHECK_EQ(num_groups_, other_column->length());
+ if (!*column) {
+ *column = other_column;
+ return Status::OK();
+ }
+ if ((*column)->null_count() == (*column)->length()) {
+ *column = other_column;
+ return Status::OK();
+ }
+ int64_t expected_non_nulls = (num_groups_ - (*column)->null_count()) +
+ (num_groups_ - other_column->null_count());
+ ARROW_ASSIGN_OR_RAISE(auto coalesced,
+ CallFunction("coalesce", {*column, other_column}, ctx_));
+ // Check that all non-null values in other_column and column were kept in the result.
+ if (expected_non_nulls != num_groups_ - coalesced.null_count()) {
+ DCHECK_GT(expected_non_nulls, num_groups_ - coalesced.null_count());
+ return DuplicateValue();
+ }
+ *column = coalesced.make_array();
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same grouped pivot key");
+ }
+
+ ExecContext* ctx_;
+ std::shared_ptr<DataType> value_type_;
+ int num_keys_;
+ int64_t num_groups_;
+ ArrayVector columns_;
+ // A persistent scratch buffer to store the take indices in Consume
+ BufferBuilder scratch_buffer_;
+};
+
+struct GroupedPivotImpl : public GroupedAggregator {
+ Status Init(ExecContext* ctx, const KernelInitArgs& args) override {
+ DCHECK_EQ(args.inputs.size(), 3);
+ key_type_ = args.inputs[0].GetSharedPtr();
+ options_ = checked_cast<const PivotWiderOptions*>(args.options);
+ DCHECK_NE(options_, nullptr);
+ auto value_type = args.inputs[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ out_struct_type_ = checked_cast<const StructType*>(out_type_.get());
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ RETURN_NOT_OK(accumulator_.Init(ctx, value_type, options_));
+ return Status::OK();
+ }
+
+ Status Resize(int64_t new_num_groups) override {
+ num_groups_ = new_num_groups;
+ return accumulator_.Resize(new_num_groups);
+ }
+
+ Status Merge(GroupedAggregator&& raw_other,
+ const ArrayData& group_id_mapping) override {
+ auto other = checked_cast<GroupedPivotImpl*>(&raw_other);
+ return accumulator_.Merge(std::move(other->accumulator_), group_id_mapping);
+ }
+
+ Status Consume(const ExecSpan& batch) override {
+ DCHECK_EQ(batch.values.size(), 3);
+ auto groups = batch[2].array.GetSpan<const uint32_t>(1, batch.length);
+ if (!batch[1].is_array()) {
+ return Status::NotImplemented("Consuming scalar pivot value");
Review Comment:
Does it make sense to support scalar here as well? (similar to the non hashed version)
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "icexelloss (via GitHub)" <gi...@apache.org>.
icexelloss commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2694532303
Thanks @pitrou - left a question on whether we want to handle scalar case for the hash kernel as well - otherwise LGTM.
I don't think the scalar case for hash kernel is blocking - we can probably do as a follow up if you want to merge this first.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1975118467
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
+/// the first one points to the pivot key column, the second points to the
+/// pivoted data column.
+/// - The pivot key column must be string-like; its values will be matched
+/// against `key_names` in order to dispatch the pivoted data into the
+/// output.
+///
+/// "pivot_wider" example
+/// ---------------------
+///
+/// Assuming the following two input columns with types utf8 and int16 (respectively):
+/// ```
+/// width | 11
+/// height | 13
+/// ```
+/// and the options `PivotWiderOptions(.key_names = {"height", "width"})`
+///
+/// then the output will be a scalar with the type
+/// `struct{"height": int16, "width": int16}`
+/// and the value `{"height": 13, "width": 11}`.
+///
+/// "hash_pivot_wider" example
Review Comment:
Yes.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1979322908
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 1]
+ //
+ // Then each output column is computed by taking the values with the
+ // respective take_indices for the column's keys.
+ //
+
+ DCHECK_EQ(groups.size(), keys.size());
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ std::shared_ptr<DataType> take_index_type;
+ std::vector<std::shared_ptr<Buffer>> take_indices(num_keys_);
+ std::vector<std::shared_ptr<Buffer>> take_bitmaps(num_keys_);
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size =
+ num_keys_ * (take_indices_size + take_bitmap_size);
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ // Slice the scratch space into individual buffers for each output column's
+ // take_indices array.
+ std::vector<TakeIndex*> take_indices_data(num_keys_);
+ std::vector<uint8_t*> take_bitmap_data(num_keys_);
+ int64_t offset = 0;
+ for (int i = 0; i < num_keys_; ++i) {
+ take_indices[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_indices_size);
+ take_indices_data[i] = take_indices[i]->mutable_data_as<TakeIndex>();
+ offset += take_indices_size;
+ take_bitmaps[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_bitmap_size);
+ take_bitmap_data[i] = take_bitmaps[i]->mutable_data();
+ memset(take_bitmap_data[i], 0, take_bitmap_size);
+ offset += take_bitmap_size;
+ }
+ DCHECK_LE(offset, scratch_buffer_.capacity());
+
+ // Populate the take_indices for each output column
+ for (int64_t i = 0; i < values.length; ++i) {
+ const PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values.IsNull(i)) {
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ const uint32_t group = groups[i];
+ if (bit_util::GetBit(take_bitmap_data[key], group)) {
+ return DuplicateValue();
+ }
+ // For row #group in column #key, we are going to take the value at index #i
+ bit_util::SetBit(take_bitmap_data[key], group);
+ take_indices_data[key][group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to compute the output columns for this batch
+ auto values_data = values.ToArrayData();
+ ArrayVector new_columns(num_keys_);
+ for (int i = 0; i < num_keys_; ++i) {
+ auto indices_data =
+ ArrayData::Make(take_index_type, num_groups_,
+ {std::move(take_bitmaps[i]), std::move(take_indices[i])});
+ // If indices_data is all nulls, we can just ignore this column.
+ if (indices_data->GetNullCount() != indices_data->length) {
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column, Take(values_data, indices_data,
+ TakeOptions::Defaults(), ctx_));
+ new_columns[i] = grouped_column.make_array();
+ }
+ }
+ // Merge them with the previous columns
+ return MergeColumns(std::move(new_columns));
+ }
+
+ Status Consume(span<const uint32_t> groups, const PivotWiderKeyIndex key,
+ const ArraySpan& values) {
+ if (key == kNullPivotKey) {
+ // Nothing to update
+ return Status::OK();
+ }
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ // The algorithm is simpler than in the array-taking version of Consume()
+ // below, since only the column #key needs to be updated.
+ std::shared_ptr<DataType> take_index_type;
+ std::shared_ptr<Buffer> take_indices;
+ std::shared_ptr<Buffer> take_bitmap;
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size = take_indices_size + take_bitmap_size;
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ take_indices = std::make_shared<MutableBuffer>(scratch_buffer_.mutable_data(),
+ take_indices_size);
+ take_bitmap = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + take_indices_size, take_bitmap_size);
+ auto take_indices_data = take_indices->mutable_data_as<TakeIndex>();
+ auto take_bitmap_data = take_bitmap->mutable_data();
+ memset(take_bitmap_data, 0, take_bitmap_size);
+
+ for (int64_t i = 0; i < values.length; ++i) {
+ const uint32_t group = groups[i];
+ if (!values.IsNull(i)) {
+ if (bit_util::GetBit(take_bitmap_data, group)) {
+ return DuplicateValue();
+ }
+ bit_util::SetBit(take_bitmap_data, group);
+ take_indices_data[group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to update column #key
+ auto values_data = values.ToArrayData();
+ auto indices_data = ArrayData::Make(
+ take_index_type, num_groups_, {std::move(take_bitmap), std::move(take_indices)});
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column,
+ Take(values_data, indices_data, TakeOptions::Defaults(), ctx_));
+ return MergeColumn(&columns_[key], grouped_column.make_array());
+ }
+
+ Status Resize(int64_t new_num_groups) {
+ if (new_num_groups > std::numeric_limits<int32_t>::max()) {
+ return Status::NotImplemented("Pivot with more 2**31 groups");
+ }
+ return ResizeColumns(new_num_groups);
+ }
+
+ Status Merge(GroupedPivotAccumulator&& other, const ArrayData& group_id_mapping) {
+ // To merge `other` into `*this`, we simply merge their respective columns.
+ // However, we must first transpose `other`'s rows using `group_id_mapping`.
+ // This is a logical "scatter" operation.
+ //
+ // Since `scatter(indices)` is implemented as `take(inverse_permutation(indices))`,
+ // we can save time by computing `inverse_permutation(indices)` once for all
+ // columns.
+
+ // Scatter/InversePermutation only accept signed indices. We checked
+ // in Resize() above that we were inside the limites for int32.
+ auto scatter_indices = group_id_mapping.Copy();
+ scatter_indices->type = int32();
+ std::shared_ptr<DataType> take_indices_type;
+ if (num_groups_ - 1 <= std::numeric_limits<int8_t>::max()) {
+ take_indices_type = int8();
+ } else if (num_groups_ - 1 <= std::numeric_limits<int16_t>::max()) {
+ take_indices_type = int16();
+ } else {
+ DCHECK_GE(num_groups_ - 1, std::numeric_limits<int32_t>::max());
+ take_indices_type = int32();
+ }
+ InversePermutationOptions options(/*max_index=*/num_groups_ - 1, take_indices_type);
+ ARROW_ASSIGN_OR_RAISE(auto take_indices,
+ InversePermutation(scatter_indices, options, ctx_));
+ auto scatter_column =
+ [&](const std::shared_ptr<Array>& column) -> Result<std::shared_ptr<Array>> {
+ ARROW_ASSIGN_OR_RAISE(auto scattered,
+ CallFunction("take", {column, take_indices}, &options, ctx_));
Review Comment:
Ahah, and this was actually hiding a bug. It was passing a `InversePermutationOptions` instead of `TakeOptions`.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1979313138
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,186 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (batch[0].is_array()) {
+ ARROW_ASSIGN_OR_RAISE(span<const PivotWiderKeyIndex> keys,
+ key_mapper_->MapKeys(batch[0].array));
+ if (batch[1].is_array()) {
+ // Array keys, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Array keys, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ } else {
+ ARROW_ASSIGN_OR_RAISE(PivotWiderKeyIndex key,
+ key_mapper_->MapKey(*batch[0].scalar));
+ if (key != kNullPivotKey) {
+ if (batch[1].is_array()) {
+ // Scalar key, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ if (!values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Scalar key, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ if (batch.length > 1 || values_[key]->is_valid) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeFrom(KernelContext*, KernelState&& src) override {
+ const auto& other_state = checked_cast<const PivotImpl&>(src);
+ for (int64_t key = 0; key < static_cast<int64_t>(values_.size()); ++key) {
+ if (other_state.values_[key]->is_valid) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = other_state.values_[key];
+ }
+ }
+ return Status::OK();
+ }
+
+ Status Finalize(KernelContext* ctx, Datum* out) override {
+ *out = std::make_shared<StructScalar>(std::move(values_), out_type_);
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same pivot key");
+ }
+
+ std::shared_ptr<DataType> out_type() const { return out_type_; }
+
+ std::shared_ptr<DataType> key_type_;
+ std::shared_ptr<DataType> out_type_;
+ const PivotWiderOptions* options_;
+ std::unique_ptr<PivotWiderKeyMapper> key_mapper_;
+ ScalarVector values_;
+};
+
+Result<std::unique_ptr<KernelState>> PivotInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ const auto& options = checked_cast<const PivotWiderOptions&>(*args.options);
+ DCHECK_EQ(args.inputs.size(), 2);
+ DCHECK(is_base_binary_like(args.inputs[0].id()));
+ auto state = std::make_unique<PivotImpl>();
+ RETURN_NOT_OK(state->Init(options, args.inputs));
+ return state;
+}
+
+Result<TypeHolder> ResolveOutputType(KernelContext* ctx, const std::vector<TypeHolder>&) {
+ return checked_cast<PivotImpl*>(ctx->state())->out_type();
+}
+
+const FunctionDoc pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear, null is emitted.\n"
+ "If a pivot key appears twice, KeyError is raised.\n"
+ "Behavior of unexpected pivot keys is controlled by PivotWiderOptions."),
Review Comment:
Can you take a look at the updated wording @zanmato1984 ?
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1979729936
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,186 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (batch[0].is_array()) {
+ ARROW_ASSIGN_OR_RAISE(span<const PivotWiderKeyIndex> keys,
+ key_mapper_->MapKeys(batch[0].array));
+ if (batch[1].is_array()) {
+ // Array keys, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Array keys, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ } else {
+ ARROW_ASSIGN_OR_RAISE(PivotWiderKeyIndex key,
+ key_mapper_->MapKey(*batch[0].scalar));
+ if (key != kNullPivotKey) {
+ if (batch[1].is_array()) {
+ // Scalar key, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ if (!values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Scalar key, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ if (batch.length > 1 || values_[key]->is_valid) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeFrom(KernelContext*, KernelState&& src) override {
+ const auto& other_state = checked_cast<const PivotImpl&>(src);
+ for (int64_t key = 0; key < static_cast<int64_t>(values_.size()); ++key) {
+ if (other_state.values_[key]->is_valid) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = other_state.values_[key];
+ }
+ }
+ return Status::OK();
+ }
+
+ Status Finalize(KernelContext* ctx, Datum* out) override {
+ *out = std::make_shared<StructScalar>(std::move(values_), out_type_);
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same pivot key");
+ }
+
+ std::shared_ptr<DataType> out_type() const { return out_type_; }
+
+ std::shared_ptr<DataType> key_type_;
+ std::shared_ptr<DataType> out_type_;
+ const PivotWiderOptions* options_;
+ std::unique_ptr<PivotWiderKeyMapper> key_mapper_;
+ ScalarVector values_;
+};
+
+Result<std::unique_ptr<KernelState>> PivotInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ const auto& options = checked_cast<const PivotWiderOptions&>(*args.options);
+ DCHECK_EQ(args.inputs.size(), 2);
+ DCHECK(is_base_binary_like(args.inputs[0].id()));
+ auto state = std::make_unique<PivotImpl>();
+ RETURN_NOT_OK(state->Init(options, args.inputs));
+ return state;
+}
+
+Result<TypeHolder> ResolveOutputType(KernelContext* ctx, const std::vector<TypeHolder>&) {
+ return checked_cast<PivotImpl*>(ctx->state())->out_type();
+}
+
+const FunctionDoc pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear, null is emitted.\n"
+ "If a pivot key appears twice, KeyError is raised.\n"
+ "Behavior of unexpected pivot keys is controlled by PivotWiderOptions."),
Review Comment:
Hmm, seems this time `KeyError` for specifying duplicate key names is not mentioned?
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2682681291
> Should we also add python binding for the option and call via pyarrow.compute?
We can indeed, for the aggregate function that is.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1975043110
##########
cpp/src/arrow/compute/kernels/pivot_internal.cc:
##########
@@ -0,0 +1,126 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/kernels/pivot_internal.h"
+
+#include <cstdint>
+
+#include "arrow/compute/exec.h"
+#include "arrow/compute/kernels/codegen_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/type_traits.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/visit_type_inline.h"
+
+namespace arrow::compute::internal {
+
+using ::arrow::util::span;
+
+struct BasePivotKeyMapper : public PivotWiderKeyMapper {
+ Status Init(const PivotWiderOptions* options) override {
+ if (options->key_names.size() > static_cast<size_t>(kMaxPivotKey) + 1) {
+ return Status::NotImplemented("Pivoting to more than ",
+ static_cast<size_t>(kMaxPivotKey) + 1,
+ " columns: got ", options->key_names.size());
+ }
+ key_name_map_.reserve(options->key_names.size());
+ PivotWiderKeyIndex index = 0;
+ for (const auto& key_name : options->key_names) {
+ bool inserted =
+ key_name_map_.try_emplace(std::string_view(key_name), index++).second;
+ if (!inserted) {
+ return Status::KeyError("Duplicate key name '", key_name,
+ "' in PivotWiderOptions");
+ }
+ }
+ unexpected_key_behavior_ = options->unexpected_key_behavior;
+ return Status::OK();
+ }
+
+ protected:
+ Result<PivotWiderKeyIndex> KeyNotFound(std::string_view key_name) {
+ if (unexpected_key_behavior_ == PivotWiderOptions::kIgnore) {
+ return kNullPivotKey;
+ }
Review Comment:
```suggestion
}
DCHECK_EQ(unexpected_key_behavior_, PivotWiderOptions::kRaise);
```
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
+/// the first one points to the pivot key column, the second points to the
+/// pivoted data column.
+/// - The pivot key column must be string-like; its values will be matched
+/// against `key_names` in order to dispatch the pivoted data into the
+/// output.
+///
+/// "pivot_wider" example
+/// ---------------------
+///
+/// Assuming the following two input columns with types utf8 and int16 (respectively):
+/// ```
+/// width | 11
+/// height | 13
+/// ```
+/// and the options `PivotWiderOptions(.key_names = {"height", "width"})`
+///
+/// then the output will be a scalar with the type
+/// `struct{"height": int16, "width": int16}`
+/// and the value `{"height": 13, "width": 11}`.
+///
+/// "hash_pivot_wider" example
+/// --------------------------
+///
+/// Assuming the following input with schema
+/// `{"group": int32, "key": utf8, "value": int16}`:
+/// ```
+/// group | key | value
+/// -----------------------------
+/// 1 | height | 11
+/// 1 | width | 12
+/// 2 | width | 13
+/// 3 | height | 14
+/// 3 | depth | 15
+/// ```
+/// and the following settings:
+/// - a hash grouping key "group"
+/// - Aggregate(
+/// .function = "hash_pivot_wider",
+/// .options = PivotWiderOptions(.key_names = {"height", "width"}),
+/// .target = {"key", "value"},
+/// .name = {"properties"})
+///
+/// then the output will have the schema
+/// `{"group": int32, "properties": struct{"height": int16, "width": int16}}`
+/// and the following value:
+/// ```
+/// group | properties
+/// | height | width
+/// -----------------------------
+/// 1 | 11 | 12
+/// 2 | null | 13
+/// 3 | 14 | null
+/// ```
+class ARROW_EXPORT PivotWiderOptions : public FunctionOptions {
Review Comment:
Any reason why this is named `PivotWider`? Will `PivotWiden` or just `Pivot` be more appropriate?
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 2]
Review Comment:
```suggestion
// - take_indices[key = 1] = [3, 2, null, 1]
```
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3456,6 +3853,18 @@ const FunctionDoc hash_one_doc{"Get one value from each group",
const FunctionDoc hash_list_doc{"List all values in each group",
("Null values are also returned."),
{"array", "group_id_array"}};
+
+const FunctionDoc hash_pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct array with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear in a given group, null is emitted.\n"
+ "If a pivot key appears twice in a given group, KeyError is raised.\n"
Review Comment:
Ditto, add description about "duplicate value".
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 2]
+ //
+ // Then each output column is computed by taking the values with the
+ // respective take_indices for the column's keys.
+ //
+
+ DCHECK_EQ(groups.size(), keys.size());
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ std::shared_ptr<DataType> take_index_type;
+ std::vector<std::shared_ptr<Buffer>> take_indices(num_keys_);
+ std::vector<std::shared_ptr<Buffer>> take_bitmaps(num_keys_);
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size =
+ num_keys_ * (take_indices_size + take_bitmap_size);
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ // Slice the scratch space into individual buffers for each output column's
+ // take_indices array.
+ std::vector<TakeIndex*> take_indices_data(num_keys_);
+ std::vector<uint8_t*> take_bitmap_data(num_keys_);
+ int64_t offset = 0;
+ for (int i = 0; i < num_keys_; ++i) {
+ take_indices[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_indices_size);
+ take_indices_data[i] = take_indices[i]->mutable_data_as<TakeIndex>();
+ offset += take_indices_size;
+ take_bitmaps[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_bitmap_size);
+ take_bitmap_data[i] = take_bitmaps[i]->mutable_data();
+ memset(take_bitmap_data[i], 0, take_bitmap_size);
+ offset += take_bitmap_size;
+ }
+ DCHECK_LE(offset, scratch_buffer_.capacity());
+
+ // Populate the take_indices for each output column
+ for (int64_t i = 0; i < values.length; ++i) {
+ const PivotWiderKeyIndex key = keys[i];
+ const uint32_t group = groups[i];
Review Comment:
Can be moved into the following `if`?
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,186 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (batch[0].is_array()) {
+ ARROW_ASSIGN_OR_RAISE(span<const PivotWiderKeyIndex> keys,
+ key_mapper_->MapKeys(batch[0].array));
+ if (batch[1].is_array()) {
+ // Array keys, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Array keys, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ } else {
+ ARROW_ASSIGN_OR_RAISE(PivotWiderKeyIndex key,
+ key_mapper_->MapKey(*batch[0].scalar));
+ if (key != kNullPivotKey) {
+ if (batch[1].is_array()) {
+ // Scalar key, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ if (!values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Scalar key, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ if (batch.length > 1 || values_[key]->is_valid) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeFrom(KernelContext*, KernelState&& src) override {
+ const auto& other_state = checked_cast<const PivotImpl&>(src);
+ for (int64_t key = 0; key < static_cast<int64_t>(values_.size()); ++key) {
+ if (other_state.values_[key]->is_valid) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = other_state.values_[key];
+ }
+ }
+ return Status::OK();
+ }
+
+ Status Finalize(KernelContext* ctx, Datum* out) override {
+ *out = std::make_shared<StructScalar>(std::move(values_), out_type_);
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same pivot key");
+ }
+
+ std::shared_ptr<DataType> out_type() const { return out_type_; }
+
+ std::shared_ptr<DataType> key_type_;
+ std::shared_ptr<DataType> out_type_;
+ const PivotWiderOptions* options_;
+ std::unique_ptr<PivotWiderKeyMapper> key_mapper_;
+ ScalarVector values_;
+};
+
+Result<std::unique_ptr<KernelState>> PivotInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ const auto& options = checked_cast<const PivotWiderOptions&>(*args.options);
+ DCHECK_EQ(args.inputs.size(), 2);
+ DCHECK(is_base_binary_like(args.inputs[0].id()));
+ auto state = std::make_unique<PivotImpl>();
+ RETURN_NOT_OK(state->Init(options, args.inputs));
+ return state;
+}
+
+Result<TypeHolder> ResolveOutputType(KernelContext* ctx, const std::vector<TypeHolder>&) {
+ return checked_cast<PivotImpl*>(ctx->state())->out_type();
+}
+
+const FunctionDoc pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear, null is emitted.\n"
+ "If a pivot key appears twice, KeyError is raised.\n"
+ "Behavior of unexpected pivot keys is controlled by PivotWiderOptions."),
Review Comment:
We may need to add some description about "duplicate values".
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2685576670
I've addressed your review comments and also added Python bindings. @icexelloss @zanmato1984
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1977943197
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 1]
+ //
+ // Then each output column is computed by taking the values with the
+ // respective take_indices for the column's keys.
+ //
+
+ DCHECK_EQ(groups.size(), keys.size());
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ std::shared_ptr<DataType> take_index_type;
+ std::vector<std::shared_ptr<Buffer>> take_indices(num_keys_);
+ std::vector<std::shared_ptr<Buffer>> take_bitmaps(num_keys_);
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size =
+ num_keys_ * (take_indices_size + take_bitmap_size);
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ // Slice the scratch space into individual buffers for each output column's
+ // take_indices array.
+ std::vector<TakeIndex*> take_indices_data(num_keys_);
+ std::vector<uint8_t*> take_bitmap_data(num_keys_);
+ int64_t offset = 0;
+ for (int i = 0; i < num_keys_; ++i) {
+ take_indices[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_indices_size);
+ take_indices_data[i] = take_indices[i]->mutable_data_as<TakeIndex>();
+ offset += take_indices_size;
+ take_bitmaps[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_bitmap_size);
+ take_bitmap_data[i] = take_bitmaps[i]->mutable_data();
+ memset(take_bitmap_data[i], 0, take_bitmap_size);
+ offset += take_bitmap_size;
+ }
+ DCHECK_LE(offset, scratch_buffer_.capacity());
+
+ // Populate the take_indices for each output column
+ for (int64_t i = 0; i < values.length; ++i) {
+ const PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values.IsNull(i)) {
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ const uint32_t group = groups[i];
+ if (bit_util::GetBit(take_bitmap_data[key], group)) {
+ return DuplicateValue();
+ }
+ // For row #group in column #key, we are going to take the value at index #i
+ bit_util::SetBit(take_bitmap_data[key], group);
+ take_indices_data[key][group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to compute the output columns for this batch
+ auto values_data = values.ToArrayData();
+ ArrayVector new_columns(num_keys_);
+ for (int i = 0; i < num_keys_; ++i) {
+ auto indices_data =
+ ArrayData::Make(take_index_type, num_groups_,
+ {std::move(take_bitmaps[i]), std::move(take_indices[i])});
+ // If indices_data is all nulls, we can just ignore this column.
+ if (indices_data->GetNullCount() != indices_data->length) {
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column, Take(values_data, indices_data,
+ TakeOptions::Defaults(), ctx_));
+ new_columns[i] = grouped_column.make_array();
+ }
+ }
+ // Merge them with the previous columns
+ return MergeColumns(std::move(new_columns));
+ }
+
+ Status Consume(span<const uint32_t> groups, const PivotWiderKeyIndex key,
+ const ArraySpan& values) {
+ if (key == kNullPivotKey) {
+ // Nothing to update
+ return Status::OK();
+ }
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ // The algorithm is simpler than in the array-taking version of Consume()
+ // below, since only the column #key needs to be updated.
+ std::shared_ptr<DataType> take_index_type;
+ std::shared_ptr<Buffer> take_indices;
+ std::shared_ptr<Buffer> take_bitmap;
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size = take_indices_size + take_bitmap_size;
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ take_indices = std::make_shared<MutableBuffer>(scratch_buffer_.mutable_data(),
+ take_indices_size);
+ take_bitmap = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + take_indices_size, take_bitmap_size);
+ auto take_indices_data = take_indices->mutable_data_as<TakeIndex>();
+ auto take_bitmap_data = take_bitmap->mutable_data();
+ memset(take_bitmap_data, 0, take_bitmap_size);
+
+ for (int64_t i = 0; i < values.length; ++i) {
+ const uint32_t group = groups[i];
+ if (!values.IsNull(i)) {
+ if (bit_util::GetBit(take_bitmap_data, group)) {
+ return DuplicateValue();
+ }
+ bit_util::SetBit(take_bitmap_data, group);
+ take_indices_data[group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to update column #key
+ auto values_data = values.ToArrayData();
+ auto indices_data = ArrayData::Make(
+ take_index_type, num_groups_, {std::move(take_bitmap), std::move(take_indices)});
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column,
+ Take(values_data, indices_data, TakeOptions::Defaults(), ctx_));
+ return MergeColumn(&columns_[key], grouped_column.make_array());
+ }
+
+ Status Resize(int64_t new_num_groups) {
+ if (new_num_groups > std::numeric_limits<int32_t>::max()) {
+ return Status::NotImplemented("Pivot with more 2**31 groups");
+ }
+ return ResizeColumns(new_num_groups);
+ }
+
+ Status Merge(GroupedPivotAccumulator&& other, const ArrayData& group_id_mapping) {
+ // To merge `other` into `*this`, we simply merge their respective columns.
+ // However, we must first transpose `other`'s rows using `group_id_mapping`.
+ // This is a logical "scatter" operation.
+ //
+ // Since `scatter(indices)` is implemented as `take(inverse_permutation(indices))`,
+ // we can save time by computing `inverse_permutation(indices)` once for all
+ // columns.
+
+ // Scatter/InversePermutation only accept signed indices. We checked
+ // in Resize() above that we were inside the limites for int32.
+ auto scatter_indices = group_id_mapping.Copy();
+ scatter_indices->type = int32();
+ std::shared_ptr<DataType> take_indices_type;
+ if (num_groups_ - 1 <= std::numeric_limits<int8_t>::max()) {
+ take_indices_type = int8();
+ } else if (num_groups_ - 1 <= std::numeric_limits<int16_t>::max()) {
+ take_indices_type = int16();
+ } else {
+ DCHECK_GE(num_groups_ - 1, std::numeric_limits<int32_t>::max());
+ take_indices_type = int32();
+ }
+ InversePermutationOptions options(/*max_index=*/num_groups_ - 1, take_indices_type);
+ ARROW_ASSIGN_OR_RAISE(auto take_indices,
+ InversePermutation(scatter_indices, options, ctx_));
+ auto scatter_column =
+ [&](const std::shared_ptr<Array>& column) -> Result<std::shared_ptr<Array>> {
+ ARROW_ASSIGN_OR_RAISE(auto scattered,
+ CallFunction("take", {column, take_indices}, &options, ctx_));
Review Comment:
Why is this not using `Take(...)` like the other two invocations?
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,186 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (batch[0].is_array()) {
+ ARROW_ASSIGN_OR_RAISE(span<const PivotWiderKeyIndex> keys,
+ key_mapper_->MapKeys(batch[0].array));
+ if (batch[1].is_array()) {
+ // Array keys, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Array keys, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ } else {
+ ARROW_ASSIGN_OR_RAISE(PivotWiderKeyIndex key,
+ key_mapper_->MapKey(*batch[0].scalar));
+ if (key != kNullPivotKey) {
+ if (batch[1].is_array()) {
+ // Scalar key, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ if (!values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Scalar key, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ if (batch.length > 1 || values_[key]->is_valid) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeFrom(KernelContext*, KernelState&& src) override {
+ const auto& other_state = checked_cast<const PivotImpl&>(src);
+ for (int64_t key = 0; key < static_cast<int64_t>(values_.size()); ++key) {
+ if (other_state.values_[key]->is_valid) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = other_state.values_[key];
+ }
+ }
+ return Status::OK();
+ }
+
+ Status Finalize(KernelContext* ctx, Datum* out) override {
+ *out = std::make_shared<StructScalar>(std::move(values_), out_type_);
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same pivot key");
+ }
+
+ std::shared_ptr<DataType> out_type() const { return out_type_; }
+
+ std::shared_ptr<DataType> key_type_;
+ std::shared_ptr<DataType> out_type_;
+ const PivotWiderOptions* options_;
+ std::unique_ptr<PivotWiderKeyMapper> key_mapper_;
+ ScalarVector values_;
+};
+
+Result<std::unique_ptr<KernelState>> PivotInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ const auto& options = checked_cast<const PivotWiderOptions&>(*args.options);
+ DCHECK_EQ(args.inputs.size(), 2);
+ DCHECK(is_base_binary_like(args.inputs[0].id()));
+ auto state = std::make_unique<PivotImpl>();
+ RETURN_NOT_OK(state->Init(options, args.inputs));
+ return state;
+}
+
+Result<TypeHolder> ResolveOutputType(KernelContext* ctx, const std::vector<TypeHolder>&) {
+ return checked_cast<PivotImpl*>(ctx->state())->out_type();
+}
+
+const FunctionDoc pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear, null is emitted.\n"
+ "If a pivot key appears twice, KeyError is raised.\n"
+ "Behavior of unexpected pivot keys is controlled by PivotWiderOptions."),
Review Comment:
Well, that's what I originally expected, until I saw `Invalid` not `KeyError` is actually thrown in `DuplicateValue`, so I think this description is not accurate enough.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1975121356
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
Review Comment:
But, of course, you can pass several aggregates to a group_by, so you can have several hash_pivot's at once if you want.
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
Review Comment:
But, of course, you can pass several aggregates to a group_by, so you can have several hash_pivot_wider aggregates at once if you want.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1971364285
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
Okay, even with Valgrind I can't reproduce anymore.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1976370007
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
+/// the first one points to the pivot key column, the second points to the
+/// pivoted data column.
+/// - The pivot key column must be string-like; its values will be matched
+/// against `key_names` in order to dispatch the pivoted data into the
+/// output.
+///
+/// "pivot_wider" example
+/// ---------------------
+///
+/// Assuming the following two input columns with types utf8 and int16 (respectively):
+/// ```
+/// width | 11
+/// height | 13
+/// ```
+/// and the options `PivotWiderOptions(.key_names = {"height", "width"})`
+///
+/// then the output will be a scalar with the type
+/// `struct{"height": int16, "width": int16}`
+/// and the value `{"height": 13, "width": 11}`.
+///
+/// "hash_pivot_wider" example
+/// --------------------------
+///
+/// Assuming the following input with schema
+/// `{"group": int32, "key": utf8, "value": int16}`:
+/// ```
+/// group | key | value
+/// -----------------------------
+/// 1 | height | 11
+/// 1 | width | 12
+/// 2 | width | 13
+/// 3 | height | 14
+/// 3 | depth | 15
+/// ```
+/// and the following settings:
+/// - a hash grouping key "group"
+/// - Aggregate(
+/// .function = "hash_pivot_wider",
+/// .options = PivotWiderOptions(.key_names = {"height", "width"}),
+/// .target = {"key", "value"},
+/// .name = {"properties"})
+///
+/// then the output will have the schema
+/// `{"group": int32, "properties": struct{"height": int16, "width": int16}}`
+/// and the following value:
+/// ```
+/// group | properties
+/// | height | width
+/// -----------------------------
+/// 1 | 11 | 12
+/// 2 | null | 13
+/// 3 | 14 | null
+/// ```
+class ARROW_EXPORT PivotWiderOptions : public FunctionOptions {
Review Comment:
Well, because there is already a node named "pivot_longer".
https://github.com/apache/arrow/blob/0fbf9823542233c5f32c26534c34cc97ce3f0be2/cpp/src/arrow/acero/pivot_longer_node.cc#L4
Besides, I find "pivot_wider" / "pivot_longer" more explanatory than "pivot" / "melt" or other alternatives.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1977602793
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 1]
+ //
+ // Then each output column is computed by taking the values with the
+ // respective take_indices for the column's keys.
+ //
+
+ DCHECK_EQ(groups.size(), keys.size());
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ std::shared_ptr<DataType> take_index_type;
+ std::vector<std::shared_ptr<Buffer>> take_indices(num_keys_);
+ std::vector<std::shared_ptr<Buffer>> take_bitmaps(num_keys_);
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size =
+ num_keys_ * (take_indices_size + take_bitmap_size);
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ // Slice the scratch space into individual buffers for each output column's
+ // take_indices array.
+ std::vector<TakeIndex*> take_indices_data(num_keys_);
+ std::vector<uint8_t*> take_bitmap_data(num_keys_);
+ int64_t offset = 0;
+ for (int i = 0; i < num_keys_; ++i) {
+ take_indices[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_indices_size);
+ take_indices_data[i] = take_indices[i]->mutable_data_as<TakeIndex>();
+ offset += take_indices_size;
+ take_bitmaps[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_bitmap_size);
+ take_bitmap_data[i] = take_bitmaps[i]->mutable_data();
+ memset(take_bitmap_data[i], 0, take_bitmap_size);
+ offset += take_bitmap_size;
+ }
+ DCHECK_LE(offset, scratch_buffer_.capacity());
+
+ // Populate the take_indices for each output column
+ for (int64_t i = 0; i < values.length; ++i) {
+ const PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values.IsNull(i)) {
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ const uint32_t group = groups[i];
+ if (bit_util::GetBit(take_bitmap_data[key], group)) {
+ return DuplicateValue();
+ }
+ // For row #group in column #key, we are going to take the value at index #i
+ bit_util::SetBit(take_bitmap_data[key], group);
+ take_indices_data[key][group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to compute the output columns for this batch
+ auto values_data = values.ToArrayData();
+ ArrayVector new_columns(num_keys_);
+ for (int i = 0; i < num_keys_; ++i) {
+ auto indices_data =
+ ArrayData::Make(take_index_type, num_groups_,
+ {std::move(take_bitmaps[i]), std::move(take_indices[i])});
+ // If indices_data is all nulls, we can just ignore this column.
+ if (indices_data->GetNullCount() != indices_data->length) {
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column, Take(values_data, indices_data,
+ TakeOptions::Defaults(), ctx_));
+ new_columns[i] = grouped_column.make_array();
+ }
+ }
+ // Merge them with the previous columns
+ return MergeColumns(std::move(new_columns));
+ }
+
+ Status Consume(span<const uint32_t> groups, const PivotWiderKeyIndex key,
+ const ArraySpan& values) {
+ if (key == kNullPivotKey) {
+ // Nothing to update
+ return Status::OK();
+ }
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ // The algorithm is simpler than in the array-taking version of Consume()
+ // below, since only the column #key needs to be updated.
+ std::shared_ptr<DataType> take_index_type;
+ std::shared_ptr<Buffer> take_indices;
+ std::shared_ptr<Buffer> take_bitmap;
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size = take_indices_size + take_bitmap_size;
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ take_indices = std::make_shared<MutableBuffer>(scratch_buffer_.mutable_data(),
+ take_indices_size);
+ take_bitmap = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + take_indices_size, take_bitmap_size);
+ auto take_indices_data = take_indices->mutable_data_as<TakeIndex>();
+ auto take_bitmap_data = take_bitmap->mutable_data();
+ memset(take_bitmap_data, 0, take_bitmap_size);
+
+ for (int64_t i = 0; i < values.length; ++i) {
+ const uint32_t group = groups[i];
+ if (!values.IsNull(i)) {
+ if (bit_util::GetBit(take_bitmap_data, group)) {
+ return DuplicateValue();
+ }
+ bit_util::SetBit(take_bitmap_data, group);
+ take_indices_data[group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to update column #key
+ auto values_data = values.ToArrayData();
+ auto indices_data = ArrayData::Make(
+ take_index_type, num_groups_, {std::move(take_bitmap), std::move(take_indices)});
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column,
+ Take(values_data, indices_data, TakeOptions::Defaults(), ctx_));
+ return MergeColumn(&columns_[key], grouped_column.make_array());
+ }
+
+ Status Resize(int64_t new_num_groups) {
+ if (new_num_groups > std::numeric_limits<int32_t>::max()) {
+ return Status::NotImplemented("Pivot with more 2**31 groups");
+ }
+ return ResizeColumns(new_num_groups);
+ }
+
+ Status Merge(GroupedPivotAccumulator&& other, const ArrayData& group_id_mapping) {
+ // To merge `other` into `*this`, we simply merge their respective columns.
+ // However, we must first transpose `other`'s rows using `group_id_mapping`.
+ // This is a logical "scatter" operation.
+ //
+ // Since `scatter(indices)` is implemented as `take(inverse_permutation(indices))`,
+ // we can save time by computing `inverse_permutation(indices)` once for all
+ // columns.
+
+ // Scatter/InversePermutation only accept signed indices. We checked
+ // in Resize() above that we were inside the limites for int32.
+ auto scatter_indices = group_id_mapping.Copy();
+ scatter_indices->type = int32();
+ std::shared_ptr<DataType> take_indices_type;
+ if (num_groups_ - 1 <= std::numeric_limits<int8_t>::max()) {
+ take_indices_type = int8();
+ } else if (num_groups_ - 1 <= std::numeric_limits<int16_t>::max()) {
+ take_indices_type = int16();
+ } else {
+ DCHECK_GE(num_groups_ - 1, std::numeric_limits<int32_t>::max());
+ take_indices_type = int32();
+ }
+ InversePermutationOptions options(/*max_index=*/num_groups_ - 1, take_indices_type);
+ ARROW_ASSIGN_OR_RAISE(auto take_indices,
+ InversePermutation(scatter_indices, options, ctx_));
+ auto scatter_column =
+ [&](const std::shared_ptr<Array>& column) -> Result<std::shared_ptr<Array>> {
+ ARROW_ASSIGN_OR_RAISE(auto scattered,
+ CallFunction("take", {column, take_indices}, &options, ctx_));
+ return scattered.make_array();
+ };
+ return MergeColumns(std::move(other.columns_), std::move(scatter_column));
+ }
+
+ Result<ArrayVector> Finalize() {
+ // Ensure that columns are allocated even if num_groups_ == 0
+ RETURN_NOT_OK(ResizeColumns(num_groups_));
+ return std::move(columns_);
+ }
+
+ protected:
+ Status ResizeColumns(int64_t new_num_groups) {
+ if (new_num_groups == num_groups_ && num_groups_ != 0) {
+ return Status::OK();
+ }
+ ARROW_ASSIGN_OR_RAISE(
+ auto array_suffix,
+ MakeArrayOfNull(value_type_, new_num_groups - num_groups_, ctx_->memory_pool()));
+ for (auto& column : columns_) {
+ if (num_groups_ != 0) {
+ DCHECK_NE(column, nullptr);
+ ARROW_ASSIGN_OR_RAISE(
+ column, Concatenate({std::move(column), array_suffix}, ctx_->memory_pool()));
+ } else {
+ column = array_suffix;
+ }
+ DCHECK_EQ(column->length(), new_num_groups);
+ }
+ num_groups_ = new_num_groups;
+ return Status::OK();
+ }
+
+ using ColumnTransform =
+ std::function<Result<std::shared_ptr<Array>>(const std::shared_ptr<Array>&)>;
+
+ Status MergeColumns(ArrayVector&& other_columns,
+ const ColumnTransform& transform = {}) {
+ DCHECK_EQ(columns_.size(), other_columns.size());
+ for (int i = 0; i < num_keys_; ++i) {
+ if (other_columns[i]) {
+ RETURN_NOT_OK(MergeColumn(&columns_[i], std::move(other_columns[i]), transform));
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeColumn(std::shared_ptr<Array>* column, std::shared_ptr<Array> other_column,
+ const ColumnTransform& transform = {}) {
+ if (other_column->null_count() == other_column->length()) {
+ // Avoid paying for the transform step below, since merging will be a no-op anyway.
+ return Status::OK();
+ }
+ if (transform) {
+ ARROW_ASSIGN_OR_RAISE(other_column, transform(other_column));
+ }
+ DCHECK_EQ(num_groups_, other_column->length());
+ if (!*column) {
+ *column = other_column;
+ return Status::OK();
+ }
+ if ((*column)->null_count() == (*column)->length()) {
+ *column = other_column;
+ return Status::OK();
+ }
+ int64_t expected_non_nulls = (num_groups_ - (*column)->null_count()) +
+ (num_groups_ - other_column->null_count());
+ ARROW_ASSIGN_OR_RAISE(auto coalesced,
+ CallFunction("coalesce", {*column, other_column}, ctx_));
+ // Check that all non-null values in other_column and column were kept in the result.
+ if (expected_non_nulls != num_groups_ - coalesced.null_count()) {
+ DCHECK_GT(expected_non_nulls, num_groups_ - coalesced.null_count());
+ return DuplicateValue();
+ }
+ *column = coalesced.make_array();
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same grouped pivot key");
+ }
+
+ ExecContext* ctx_;
+ std::shared_ptr<DataType> value_type_;
+ int num_keys_;
+ int64_t num_groups_;
+ ArrayVector columns_;
+ // A persistent scratch buffer to store the take indices in Consume
+ BufferBuilder scratch_buffer_;
+};
+
+struct GroupedPivotImpl : public GroupedAggregator {
+ Status Init(ExecContext* ctx, const KernelInitArgs& args) override {
+ DCHECK_EQ(args.inputs.size(), 3);
+ key_type_ = args.inputs[0].GetSharedPtr();
+ options_ = checked_cast<const PivotWiderOptions*>(args.options);
+ DCHECK_NE(options_, nullptr);
+ auto value_type = args.inputs[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ out_struct_type_ = checked_cast<const StructType*>(out_type_.get());
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ RETURN_NOT_OK(accumulator_.Init(ctx, value_type, options_));
+ return Status::OK();
+ }
+
+ Status Resize(int64_t new_num_groups) override {
+ num_groups_ = new_num_groups;
+ return accumulator_.Resize(new_num_groups);
+ }
+
+ Status Merge(GroupedAggregator&& raw_other,
+ const ArrayData& group_id_mapping) override {
+ auto other = checked_cast<GroupedPivotImpl*>(&raw_other);
+ return accumulator_.Merge(std::move(other->accumulator_), group_id_mapping);
+ }
+
+ Status Consume(const ExecSpan& batch) override {
+ DCHECK_EQ(batch.values.size(), 3);
+ auto groups = batch[2].array.GetSpan<const uint32_t>(1, batch.length);
+ if (!batch[1].is_array()) {
+ return Status::NotImplemented("Consuming scalar pivot value");
Review Comment:
Good question. It might, depending on whether that is easily testable.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1977539472
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 2]
Review Comment:
Nice catch!
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "icexelloss (via GitHub)" <gi...@apache.org>.
icexelloss commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1969909007
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,159 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (!batch[1].is_array()) {
Review Comment:
Can we handle scalar here too? Seems pretty trivial to do (basically put the scalar value to the outputs) or expand the scalar to array.
Although this is probably practical use case, but we have seen lot of cases where user calls functions on scalar values (e.g., when you assign a literal columns in Acero it is scalar) so it is probably better UX to always support scalars
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1971109477
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
Any test where the merge step returns an error could be affected. I guess other hash-aggregate functions never fail in their merge step, so this was not exercised before.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1971137508
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
You can look at `PivotDuplicateValues`. Especially the second case in that test ("Duplicate values in different chunks"). That should hopefully do the trick.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1959711502
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
Without this change, any error in `MergeAll` below could lead to a crash.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1971307687
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
> I do have a suspect though: this might be race related. And if, a big if, this is the case, I'm afraid your change doesn't fix the race. However there are still many things unclear so please don't take this seriously until I have new findings.
Update: Turns out my suspect of race won't happen actually.
So I'm at my wits end by now.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2685356889
Revision: d97a9a7d5b0abdc9b42faffaccfaed12bbd3362e
Submitted crossbow builds: [ursacomputing/crossbow @ actions-e51f9aecb5](https://github.com/ursacomputing/crossbow/branches/all?query=actions-e51f9aecb5)
|Task|Status|
|----|------|
|example-cpp-minimal-build-static|[](https://github.com/ursacomputing/crossbow/actions/runs/13547029466/job/37861103098)|
|example-cpp-minimal-build-static-system-dependency|[](https://github.com/ursacomputing/crossbow/actions/runs/13547027161/job/37861095853)|
|example-cpp-tutorial|[](https://github.com/ursacomputing/crossbow/actions/runs/13547028190/job/37861098980)|
|test-alpine-linux-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13547028516/job/37861100689)|
|test-build-cpp-fuzz|[](https://github.com/ursacomputing/crossbow/actions/runs/13547028162/job/37861098687)|
|test-conda-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13547026932/job/37861095256)|
|test-conda-cpp-meson|[](https://github.com/ursacomputing/crossbow/actions/runs/13547024484/job/37861087005)|
|test-conda-cpp-valgrind|[](https://github.com/ursacomputing/crossbow/actions/runs/13547029721/job/37861104623)|
|test-cuda-cpp-ubuntu-22.04-cuda-11.7.1|[](https://github.com/ursacomputing/crossbow/actions/runs/13547024867/job/37861088511)|
|test-debian-12-cpp-amd64|[](https://github.com/ursacomputing/crossbow/actions/runs/13547026516/job/37861094842)|
|test-debian-12-cpp-i386|[](https://github.com/ursacomputing/crossbow/actions/runs/13547026929/job/37861095199)|
|test-fedora-39-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13547028151/job/37861098624)|
|test-ubuntu-22.04-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13547030148/job/37861106128)|
|test-ubuntu-22.04-cpp-20|[](https://github.com/ursacomputing/crossbow/actions/runs/13547024389/job/37861086969)|
|test-ubuntu-22.04-cpp-bundled|[](https://github.com/ursacomputing/crossbow/actions/runs/13547028332/job/37861100666)|
|test-ubuntu-22.04-cpp-emscripten|[](https://github.com/ursacomputing/crossbow/actions/runs/13547028291/job/37861100018)|
|test-ubuntu-22.04-cpp-no-threading|[](https://github.com/ursacomputing/crossbow/actions/runs/13547027929/job/37861098029)|
|test-ubuntu-24.04-cpp|[](https://github.com/ursacomputing/crossbow/actions/runs/13547029943/job/37861105173)|
|test-ubuntu-24.04-cpp-bundled-offline|[](https://github.com/ursacomputing/crossbow/actions/runs/13547028207/job/37861099627)|
|test-ubuntu-24.04-cpp-gcc-13-bundled|[](https://github.com/ursacomputing/crossbow/actions/runs/13547024920/job/37861089217)|
|test-ubuntu-24.04-cpp-gcc-14|[](https://github.com/ursacomputing/crossbow/actions/runs/13547027533/job/37861097085)|
|test-ubuntu-24.04-cpp-minimal-with-formats|[](https://github.com/ursacomputing/crossbow/actions/runs/13547028450/job/37861100631)|
|test-ubuntu-24.04-cpp-thread-sanitizer|[](https://github.com/ursacomputing/crossbow/actions/runs/13547026157/job/37861093108)|
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1974826981
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
Review Comment:
Do we not support multiple "key" or "data" columns? Or is the user expected to use a `struct` to wrap the desired key/data columns to achieve that?
##########
cpp/src/arrow/compute/api_aggregate.h:
##########
@@ -175,6 +175,88 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
uint32_t min_count;
};
+/// \brief Control Pivot kernel behavior
+///
+/// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
+///
+/// Constraints:
+/// - The corresponding `Aggregate::target` must have two FieldRef elements;
+/// the first one points to the pivot key column, the second points to the
+/// pivoted data column.
+/// - The pivot key column must be string-like; its values will be matched
+/// against `key_names` in order to dispatch the pivoted data into the
+/// output.
+///
+/// "pivot_wider" example
+/// ---------------------
+///
+/// Assuming the following two input columns with types utf8 and int16 (respectively):
+/// ```
+/// width | 11
+/// height | 13
+/// ```
+/// and the options `PivotWiderOptions(.key_names = {"height", "width"})`
+///
+/// then the output will be a scalar with the type
+/// `struct{"height": int16, "width": int16}`
+/// and the value `{"height": 13, "width": 11}`.
+///
+/// "hash_pivot_wider" example
Review Comment:
So the pandas's `pivot` is more like our `hash` version of `pivot_wider`?
##########
cpp/src/arrow/compute/kernels/hash_aggregate.cc:
##########
@@ -3319,9 +3324,401 @@ struct GroupedListFactory {
HashAggregateKernel kernel;
InputType argument_type;
};
-} // namespace
-namespace {
+// ----------------------------------------------------------------------
+// Pivot implementation
+
+struct GroupedPivotAccumulator {
+ Status Init(ExecContext* ctx, std::shared_ptr<DataType> value_type,
+ const PivotWiderOptions* options) {
+ ctx_ = ctx;
+ value_type_ = std::move(value_type);
+ num_keys_ = static_cast<int>(options->key_names.size());
+ num_groups_ = 0;
+ columns_.resize(num_keys_);
+ scratch_buffer_ = BufferBuilder(ctx_->memory_pool());
+ return Status::OK();
+ }
+
+ Status Consume(span<const uint32_t> groups, span<const PivotWiderKeyIndex> keys,
+ const ArraySpan& values) {
+ // To dispatch the values into the right (group, key) coordinates,
+ // we first compute a vector of take indices for each output column.
+ //
+ // For each index #i, we set take_indices[keys[#i]][groups[#i]] = #i.
+ // Unpopulated take_indices entries are null.
+ //
+ // For example, assuming we get:
+ // groups | keys
+ // ===================
+ // 1 | 0
+ // 3 | 1
+ // 1 | 1
+ // 0 | 1
+ //
+ // We are going to compute:
+ // - take_indices[key = 0] = [null, 0, null, null]
+ // - take_indices[key = 1] = [3, 2, null, 2]
+ //
+ // Then each output column is computed by taking the values with the
+ // respective take_indices for the column's keys.
+ //
+
+ DCHECK_EQ(groups.size(), keys.size());
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ std::shared_ptr<DataType> take_index_type;
+ std::vector<std::shared_ptr<Buffer>> take_indices(num_keys_);
+ std::vector<std::shared_ptr<Buffer>> take_bitmaps(num_keys_);
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size =
+ num_keys_ * (take_indices_size + take_bitmap_size);
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ // Slice the scratch space into individual buffers for each output column's
+ // take_indices array.
+ std::vector<TakeIndex*> take_indices_data(num_keys_);
+ std::vector<uint8_t*> take_bitmap_data(num_keys_);
+ int64_t offset = 0;
+ for (int i = 0; i < num_keys_; ++i) {
+ take_indices[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_indices_size);
+ take_indices_data[i] = take_indices[i]->mutable_data_as<TakeIndex>();
+ offset += take_indices_size;
+ take_bitmaps[i] = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + offset, take_bitmap_size);
+ take_bitmap_data[i] = take_bitmaps[i]->mutable_data();
+ memset(take_bitmap_data[i], 0, take_bitmap_size);
+ offset += take_bitmap_size;
+ }
+ DCHECK_LE(offset, scratch_buffer_.capacity());
+
+ // Populate the take_indices for each output column
+ for (int64_t i = 0; i < values.length; ++i) {
+ const PivotWiderKeyIndex key = keys[i];
+ const uint32_t group = groups[i];
+ if (key != kNullPivotKey && !values.IsNull(i)) {
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ if (bit_util::GetBit(take_bitmap_data[key], group)) {
+ return DuplicateValue();
+ }
+ // For row #group in column #key, we are going to take the value at index #i
+ bit_util::SetBit(take_bitmap_data[key], group);
+ take_indices_data[key][group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to compute the output columns for this batch
+ auto values_data = values.ToArrayData();
+ ArrayVector new_columns(num_keys_);
+ for (int i = 0; i < num_keys_; ++i) {
+ auto indices_data =
+ ArrayData::Make(take_index_type, num_groups_,
+ {std::move(take_bitmaps[i]), std::move(take_indices[i])});
+ // If indices_data is all nulls, we can just ignore this column.
+ if (indices_data->GetNullCount() != indices_data->length) {
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column, Take(values_data, indices_data,
+ TakeOptions::Defaults(), ctx_));
+ new_columns[i] = grouped_column.make_array();
+ }
+ }
+ // Merge them with the previous columns
+ return MergeColumns(std::move(new_columns));
+ }
+
+ Status Consume(span<const uint32_t> groups, const PivotWiderKeyIndex key,
+ const ArraySpan& values) {
+ if (key == kNullPivotKey) {
+ // Nothing to update
+ return Status::OK();
+ }
+ DCHECK_LT(static_cast<int>(key), num_keys_);
+ DCHECK_EQ(groups.size(), static_cast<size_t>(values.length));
+
+ // The algorithm is simpler than in the array-taking version of Consume()
+ // below, since only the column #key needs to be updated.
+ std::shared_ptr<DataType> take_index_type;
+ std::shared_ptr<Buffer> take_indices;
+ std::shared_ptr<Buffer> take_bitmap;
+
+ // A generic lambda that computes the take indices with the desired integer width
+ auto compute_take_indices = [&](auto typed_index) {
+ ARROW_UNUSED(typed_index);
+ using TakeIndex = std::decay_t<decltype(typed_index)>;
+ take_index_type = CTypeTraits<TakeIndex>::type_singleton();
+
+ const int64_t take_indices_size =
+ bit_util::RoundUpToMultipleOf64(num_groups_ * sizeof(TakeIndex));
+ const int64_t take_bitmap_size =
+ bit_util::RoundUpToMultipleOf64(bit_util::BytesForBits(num_groups_));
+ const int64_t total_scratch_size = take_indices_size + take_bitmap_size;
+ RETURN_NOT_OK(scratch_buffer_.Resize(total_scratch_size, /*shrink_to_fit=*/false));
+
+ take_indices = std::make_shared<MutableBuffer>(scratch_buffer_.mutable_data(),
+ take_indices_size);
+ take_bitmap = std::make_shared<MutableBuffer>(
+ scratch_buffer_.mutable_data() + take_indices_size, take_bitmap_size);
+ auto take_indices_data = take_indices->mutable_data_as<TakeIndex>();
+ auto take_bitmap_data = take_bitmap->mutable_data();
+ memset(take_bitmap_data, 0, take_bitmap_size);
+
+ for (int64_t i = 0; i < values.length; ++i) {
+ const uint32_t group = groups[i];
+ if (!values.IsNull(i)) {
+ if (bit_util::GetBit(take_bitmap_data, group)) {
+ return DuplicateValue();
+ }
+ bit_util::SetBit(take_bitmap_data, group);
+ take_indices_data[group] = static_cast<TakeIndex>(i);
+ }
+ }
+ return Status::OK();
+ };
+
+ // Call compute_take_indices with the optimal integer width
+ if (values.length <= static_cast<int64_t>(std::numeric_limits<uint8_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint8_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint16_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint16_t{}));
+ } else if (values.length <=
+ static_cast<int64_t>(std::numeric_limits<uint32_t>::max())) {
+ RETURN_NOT_OK(compute_take_indices(uint32_t{}));
+ } else {
+ RETURN_NOT_OK(compute_take_indices(uint64_t{}));
+ }
+
+ // Use take_indices to update column #key
+ auto values_data = values.ToArrayData();
+ auto indices_data = ArrayData::Make(
+ take_index_type, num_groups_, {std::move(take_bitmap), std::move(take_indices)});
+ ARROW_ASSIGN_OR_RAISE(Datum grouped_column,
+ Take(values_data, indices_data, TakeOptions::Defaults(), ctx_));
+ return MergeColumn(&columns_[key], grouped_column.make_array());
+ }
+
+ Status Resize(int64_t new_num_groups) {
+ if (new_num_groups > std::numeric_limits<int32_t>::max()) {
+ return Status::NotImplemented("Pivot with more 2**31 groups");
+ }
+ return ResizeColumns(new_num_groups);
+ }
+
+ Status Merge(GroupedPivotAccumulator&& other, const ArrayData& group_id_mapping) {
+ // To merge `other` into `*this`, we simply merge their respective columns.
+ // However, we must first transpose `other`'s rows using `group_id_mapping`.
+ // This is a logical "scatter" operation.
+ //
+ // Since `scatter(indices)` is implemented as `take(inverse_permutation(indices))`,
+ // we can save time by computing `inverse_permutation(indices)` once for all
Review Comment:
Wow, so `inverse_permutation` is actually more useful than directly invoking `scatter` in this case?
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "icexelloss (via GitHub)" <gi...@apache.org>.
icexelloss commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1969885407
##########
cpp/src/arrow/acero/scalar_aggregate_node.cc:
##########
@@ -294,6 +294,14 @@ Status ScalarAggregateNode::OutputResult(bool is_last) {
// First, insert segment keys
PlaceFields(batch, /*base=*/0, segmenter_values_);
+ // Move away the states and recreate them eagerly, to make sure that any error
+ // below does not leave us with empty states.
+ auto states = std::move(states_);
+ states_.resize(kernels_.size());
+ if (!is_last) {
+ RETURN_NOT_OK(ResetKernelStates());
+ }
Review Comment:
Not sure I understand this - shouldn't error in `MergeAll` be bubbled up and stop execution? (Why does it lead to crash?)
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1970182595
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,159 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (!batch[1].is_array()) {
Review Comment:
Good point, I'll look into it.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "zanmato1984 (via GitHub)" <gi...@apache.org>.
zanmato1984 commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1978618339
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,186 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (batch[0].is_array()) {
+ ARROW_ASSIGN_OR_RAISE(span<const PivotWiderKeyIndex> keys,
+ key_mapper_->MapKeys(batch[0].array));
+ if (batch[1].is_array()) {
+ // Array keys, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Array keys, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ } else {
+ ARROW_ASSIGN_OR_RAISE(PivotWiderKeyIndex key,
+ key_mapper_->MapKey(*batch[0].scalar));
+ if (key != kNullPivotKey) {
+ if (batch[1].is_array()) {
+ // Scalar key, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ if (!values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Scalar key, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ if (batch.length > 1 || values_[key]->is_valid) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeFrom(KernelContext*, KernelState&& src) override {
+ const auto& other_state = checked_cast<const PivotImpl&>(src);
+ for (int64_t key = 0; key < static_cast<int64_t>(values_.size()); ++key) {
+ if (other_state.values_[key]->is_valid) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = other_state.values_[key];
+ }
+ }
+ return Status::OK();
+ }
+
+ Status Finalize(KernelContext* ctx, Datum* out) override {
+ *out = std::make_shared<StructScalar>(std::move(values_), out_type_);
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same pivot key");
+ }
+
+ std::shared_ptr<DataType> out_type() const { return out_type_; }
+
+ std::shared_ptr<DataType> key_type_;
+ std::shared_ptr<DataType> out_type_;
+ const PivotWiderOptions* options_;
+ std::unique_ptr<PivotWiderKeyMapper> key_mapper_;
+ ScalarVector values_;
+};
+
+Result<std::unique_ptr<KernelState>> PivotInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ const auto& options = checked_cast<const PivotWiderOptions&>(*args.options);
+ DCHECK_EQ(args.inputs.size(), 2);
+ DCHECK(is_base_binary_like(args.inputs[0].id()));
+ auto state = std::make_unique<PivotImpl>();
+ RETURN_NOT_OK(state->Init(options, args.inputs));
+ return state;
+}
+
+Result<TypeHolder> ResolveOutputType(KernelContext* ctx, const std::vector<TypeHolder>&) {
+ return checked_cast<PivotImpl*>(ctx->state())->out_type();
+}
+
+const FunctionDoc pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear, null is emitted.\n"
+ "If a pivot key appears twice, KeyError is raised.\n"
+ "Behavior of unexpected pivot keys is controlled by PivotWiderOptions."),
Review Comment:
At least we should either change the error being thrown or refine the wording here.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #45562:
URL: https://github.com/apache/arrow/pull/45562#issuecomment-2698222850
Thanks a lot for the reviews @zanmato1984 @icexelloss ! I will merge this PR now.
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou merged PR #45562:
URL: https://github.com/apache/arrow/pull/45562
--
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
Re: [PR] GH-45269: [C++][Compute] Add "pivot_wider" and "hash_pivot_wider" functions [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #45562:
URL: https://github.com/apache/arrow/pull/45562#discussion_r1979732388
##########
cpp/src/arrow/compute/kernels/aggregate_pivot.cc:
##########
@@ -0,0 +1,186 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/common_internal.h"
+#include "arrow/compute/kernels/pivot_internal.h"
+#include "arrow/scalar.h"
+#include "arrow/util/bit_run_reader.h"
+#include "arrow/util/logging.h"
+
+namespace arrow::compute::internal {
+namespace {
+
+using arrow::internal::VisitSetBitRunsVoid;
+using arrow::util::span;
+
+struct PivotImpl : public ScalarAggregator {
+ Status Init(const PivotWiderOptions& options, const std::vector<TypeHolder>& in_types) {
+ options_ = &options;
+ key_type_ = in_types[0].GetSharedPtr();
+ auto value_type = in_types[1].GetSharedPtr();
+ FieldVector fields;
+ fields.reserve(options_->key_names.size());
+ values_.reserve(options_->key_names.size());
+ for (const auto& key_name : options_->key_names) {
+ fields.push_back(field(key_name, value_type));
+ values_.push_back(MakeNullScalar(value_type));
+ }
+ out_type_ = struct_(std::move(fields));
+ ARROW_ASSIGN_OR_RAISE(key_mapper_, PivotWiderKeyMapper::Make(*key_type_, options_));
+ return Status::OK();
+ }
+
+ Status Consume(KernelContext*, const ExecSpan& batch) override {
+ DCHECK_EQ(batch.num_values(), 2);
+ if (batch[0].is_array()) {
+ ARROW_ASSIGN_OR_RAISE(span<const PivotWiderKeyIndex> keys,
+ key_mapper_->MapKeys(batch[0].array));
+ if (batch[1].is_array()) {
+ // Array keys, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey && !values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Array keys, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ for (int64_t i = 0; i < batch.length; ++i) {
+ PivotWiderKeyIndex key = keys[i];
+ if (key != kNullPivotKey) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ } else {
+ ARROW_ASSIGN_OR_RAISE(PivotWiderKeyIndex key,
+ key_mapper_->MapKey(*batch[0].scalar));
+ if (key != kNullPivotKey) {
+ if (batch[1].is_array()) {
+ // Scalar key, array values
+ auto values = batch[1].array.ToArray();
+ for (int64_t i = 0; i < batch.length; ++i) {
+ if (!values->IsNull(i)) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ ARROW_ASSIGN_OR_RAISE(values_[key], values->GetScalar(i));
+ DCHECK(values_[key]->is_valid);
+ }
+ }
+ } else {
+ // Scalar key, scalar value
+ const Scalar* value = batch[1].scalar;
+ if (value->is_valid) {
+ if (batch.length > 1 || values_[key]->is_valid) {
+ return DuplicateValue();
+ }
+ values_[key] = value->GetSharedPtr();
+ }
+ }
+ }
+ }
+ return Status::OK();
+ }
+
+ Status MergeFrom(KernelContext*, KernelState&& src) override {
+ const auto& other_state = checked_cast<const PivotImpl&>(src);
+ for (int64_t key = 0; key < static_cast<int64_t>(values_.size()); ++key) {
+ if (other_state.values_[key]->is_valid) {
+ if (ARROW_PREDICT_FALSE(values_[key]->is_valid)) {
+ return DuplicateValue();
+ }
+ values_[key] = other_state.values_[key];
+ }
+ }
+ return Status::OK();
+ }
+
+ Status Finalize(KernelContext* ctx, Datum* out) override {
+ *out = std::make_shared<StructScalar>(std::move(values_), out_type_);
+ return Status::OK();
+ }
+
+ Status DuplicateValue() {
+ return Status::Invalid(
+ "Encountered more than one non-null value for the same pivot key");
+ }
+
+ std::shared_ptr<DataType> out_type() const { return out_type_; }
+
+ std::shared_ptr<DataType> key_type_;
+ std::shared_ptr<DataType> out_type_;
+ const PivotWiderOptions* options_;
+ std::unique_ptr<PivotWiderKeyMapper> key_mapper_;
+ ScalarVector values_;
+};
+
+Result<std::unique_ptr<KernelState>> PivotInit(KernelContext* ctx,
+ const KernelInitArgs& args) {
+ const auto& options = checked_cast<const PivotWiderOptions&>(*args.options);
+ DCHECK_EQ(args.inputs.size(), 2);
+ DCHECK(is_base_binary_like(args.inputs[0].id()));
+ auto state = std::make_unique<PivotImpl>();
+ RETURN_NOT_OK(state->Init(options, args.inputs));
+ return state;
+}
+
+Result<TypeHolder> ResolveOutputType(KernelContext* ctx, const std::vector<TypeHolder>&) {
+ return checked_cast<PivotImpl*>(ctx->state())->out_type();
+}
+
+const FunctionDoc pivot_doc{
+ "Pivot values according to a pivot key column",
+ ("Output is a struct with as many fields as `PivotWiderOptions.key_names`.\n"
+ "All output struct fields have the same type as `pivot_values`.\n"
+ "Each pivot key decides in which output field the corresponding pivot value\n"
+ "is emitted. If a pivot key doesn't appear, null is emitted.\n"
+ "If a pivot key appears twice, KeyError is raised.\n"
+ "Behavior of unexpected pivot keys is controlled by PivotWiderOptions."),
Review Comment:
Yes, should we mention it? It's not really likely to occur. But we can still be explicit about it.
--
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