From b8db90cd3fe1e5b288ae4a81ffc41e1d79971d55 Mon Sep 17 00:00:00 2001 From: ZhangHuiGui Date: Wed, 27 Mar 2024 22:02:29 +0800 Subject: [PATCH] fix --- cpp/src/arrow/compute/key_hash.cc | 66 +++++++++++++------------------ cpp/src/arrow/compute/util.cc | 6 +-- cpp/src/arrow/compute/util.h | 6 +-- 3 files changed, 34 insertions(+), 44 deletions(-) diff --git a/cpp/src/arrow/compute/key_hash.cc b/cpp/src/arrow/compute/key_hash.cc index 493df15c8fcd1..87772eed156f7 100644 --- a/cpp/src/arrow/compute/key_hash.cc +++ b/cpp/src/arrow/compute/key_hash.cc @@ -379,38 +379,37 @@ void Hashing32::HashFixed(int64_t hardware_flags, bool combine_hashes, uint32_t } Status Hashing32::HashMultiColumn(const std::vector& cols, - LightContext* ctx, uint32_t* hashes) { - uint32_t num_rows = static_cast(cols[0].length()); - - constexpr uint32_t max_batch_size = util::MiniBatch::kMiniBatchLength; - const auto alloc_batch_size = std::min(num_rows, max_batch_size); + LightContext* ctx, uint32_t* hashes) { + auto num_rows = static_cast(cols[0].length()); + const auto alloc_batch_size = + std::min(num_rows, static_cast(util::MiniBatch::kMiniBatchLength)); // pre calculate alloc size in TempVectorStack for hash_temp_buf, null_hash_temp_buf // and null_indices_buf const auto alloc_hash_temp_buf = - util::TempVectorStack::EstimateAllocSize(alloc_batch_size * sizeof(uint32_t)); + util::TempVectorStack::EstimateAllocationSize(alloc_batch_size * sizeof(uint32_t)); const auto alloc_for_null_indices_buf = - util::TempVectorStack::EstimateAllocSize(alloc_batch_size * sizeof(uint16_t)); + util::TempVectorStack::EstimateAllocationSize(alloc_batch_size * sizeof(uint16_t)); const auto alloc_size = alloc_hash_temp_buf * 2 + alloc_for_null_indices_buf; - std::shared_ptr temp_stack(nullptr); - if (!ctx->stack) { - temp_stack = std::make_shared(); + std::unique_ptr temp_stack(nullptr); + auto stack = ctx->stack; + if (!stack) { + temp_stack = std::make_unique(); RETURN_NOT_OK(temp_stack->Init(default_memory_pool(), alloc_size)); - ctx->stack = temp_stack.get(); + stack = temp_stack.get(); } else { - RETURN_NOT_OK(ctx->stack->CheckAllocOverflow(alloc_size)); + RETURN_NOT_OK(stack->CheckAllocationOverflow(alloc_size)); } - auto hash_temp_buf = util::TempVectorHolder(ctx->stack, alloc_batch_size); + auto hash_temp_buf = util::TempVectorHolder(stack, alloc_batch_size); uint32_t* hash_temp = hash_temp_buf.mutable_data(); - auto null_indices_buf = util::TempVectorHolder(ctx->stack, alloc_batch_size); + auto null_indices_buf = util::TempVectorHolder(stack, alloc_batch_size); uint16_t* null_indices = null_indices_buf.mutable_data(); int num_null_indices; - auto null_hash_temp_buf = - util::TempVectorHolder(ctx->stack, alloc_batch_size); + auto null_hash_temp_buf = util::TempVectorHolder(stack, alloc_batch_size); uint32_t* null_hash_temp = null_hash_temp_buf.mutable_data(); for (uint32_t first_row = 0; first_row < num_rows;) { @@ -479,10 +478,6 @@ Status Hashing32::HashMultiColumn(const std::vector& cols, first_row += batch_size_next; } - - if (temp_stack) { - ctx->stack = nullptr; - } return Status::OK(); } @@ -846,33 +841,32 @@ void Hashing64::HashFixed(bool combine_hashes, uint32_t num_keys, uint64_t key_l Status Hashing64::HashMultiColumn(const std::vector& cols, LightContext* ctx, uint64_t* hashes) { - uint32_t num_rows = static_cast(cols[0].length()); - - constexpr uint32_t max_batch_size = util::MiniBatch::kMiniBatchLength; - const auto alloc_batch_size = std::min(num_rows, max_batch_size); + auto num_rows = static_cast(cols[0].length()); + const auto alloc_batch_size = + std::min(num_rows, static_cast(util::MiniBatch::kMiniBatchLength)); // pre calculate alloc size in TempVectorStack for null_indices_buf, null_hash_temp_buf const auto alloc_for_null_hash_temp_buf = - util::TempVectorStack::EstimateAllocSize(alloc_batch_size * sizeof(uint64_t)); + util::TempVectorStack::EstimateAllocationSize(alloc_batch_size * sizeof(uint64_t)); const auto alloc_for_null_indices_buf = - util::TempVectorStack::EstimateAllocSize(alloc_batch_size * sizeof(uint16_t)); + util::TempVectorStack::EstimateAllocationSize(alloc_batch_size * sizeof(uint16_t)); const auto alloc_size = alloc_for_null_hash_temp_buf + alloc_for_null_indices_buf; - std::shared_ptr temp_stack(nullptr); - if (!ctx->stack) { - temp_stack = std::make_shared(); + std::unique_ptr temp_stack(nullptr); + auto stack = ctx->stack; + if (!stack) { + temp_stack = std::make_unique(); RETURN_NOT_OK(temp_stack->Init(default_memory_pool(), alloc_size)); - ctx->stack = temp_stack.get(); + stack = temp_stack.get(); } else { - RETURN_NOT_OK(ctx->stack->CheckAllocOverflow(alloc_size)); + RETURN_NOT_OK(stack->CheckAllocationOverflow(alloc_size)); } - auto null_indices_buf = util::TempVectorHolder(ctx->stack, alloc_batch_size); + auto null_indices_buf = util::TempVectorHolder(stack, alloc_batch_size); uint16_t* null_indices = null_indices_buf.mutable_data(); int num_null_indices; - auto null_hash_temp_buf = - util::TempVectorHolder(ctx->stack, alloc_batch_size); + auto null_hash_temp_buf = util::TempVectorHolder(stack, alloc_batch_size); uint64_t* null_hash_temp = null_hash_temp_buf.mutable_data(); for (uint32_t first_row = 0; first_row < num_rows;) { @@ -938,10 +932,6 @@ Status Hashing64::HashMultiColumn(const std::vector& cols, first_row += batch_size_next; } - - if (temp_stack) { - ctx->stack = nullptr; - } return Status::OK(); } diff --git a/cpp/src/arrow/compute/util.cc b/cpp/src/arrow/compute/util.cc index 199e0ddace45c..014fbe339dc7f 100644 --- a/cpp/src/arrow/compute/util.cc +++ b/cpp/src/arrow/compute/util.cc @@ -32,9 +32,9 @@ using internal::CpuInfo; namespace util { void TempVectorStack::alloc(uint32_t num_bytes, uint8_t** data, int* id) { - const auto estimate_size = EstimateAllocSize(num_bytes); + const auto estimate_size = EstimateAllocationSize(num_bytes); // XXX cannot return a regular Status because most consumers do not either. - ARROW_DCHECK_OK(CheckAllocOverflow(estimate_size)); + ARROW_CHECK_OK(CheckAllocationOverflow(estimate_size)); int64_t new_top = top_ + estimate_size; *data = buffer_->mutable_data() + top_ + sizeof(uint64_t); // We set 8 bytes before the beginning of the allocated range and @@ -58,7 +58,7 @@ void TempVectorStack::release(int id, uint32_t num_bytes) { --num_vectors_; } -Status TempVectorStack::CheckAllocOverflow(int64_t alloc_size) { +Status TempVectorStack::CheckAllocationOverflow(int64_t alloc_size) { // Stack overflow check (see GH-39582). if ((alloc_size + top_) > buffer_size_) { return Status::Invalid("TempVectorStack alloc overflow. (Actual ", buffer_size_, diff --git a/cpp/src/arrow/compute/util.h b/cpp/src/arrow/compute/util.h index fab86f44dcd34..9817b49ecee50 100644 --- a/cpp/src/arrow/compute/util.h +++ b/cpp/src/arrow/compute/util.h @@ -89,7 +89,7 @@ class ARROW_EXPORT TempVectorStack { Status Init(MemoryPool* pool, int64_t size) { num_vectors_ = 0; top_ = 0; - buffer_size_ = EstimateAllocSize(size); + buffer_size_ = EstimateAllocationSize(size); ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateResizableBuffer(size, pool)); // Ensure later operations don't accidentally read uninitialized memory. std::memset(buffer->mutable_data(), 0xFF, size); @@ -97,11 +97,11 @@ class ARROW_EXPORT TempVectorStack { return Status::OK(); } - static int64_t EstimateAllocSize(int64_t size) { + static int64_t EstimateAllocationSize(int64_t size) { return PaddedAllocationSize(size) + 2 * sizeof(uint64_t); } - Status CheckAllocOverflow(int64_t alloc_size); + Status CheckAllocationOverflow(int64_t alloc_size); private: static int64_t PaddedAllocationSize(int64_t num_bytes) {