fix

cpp/src/arrow/compute/key_hash.cc

@@ -384,6 +384,15 @@ void Hashing32::HashMultiColumn(const std::vector<KeyColumnArray>& cols,
 
   constexpr uint32_t max_batch_size = util::MiniBatch::kMiniBatchLength;
   const uint32_t alloc_batch_size = std::min(num_rows, max_batch_size);
+  const int64_t estimate_alloc_size = EstimateBatchStackSize<uint32_t>(alloc_batch_size);
+
+  util::TempVectorStack temp_stack;
+  if (!ctx->stack) {
+    ARROW_CHECK_OK(temp_stack.Init(default_memory_pool(), estimate_alloc_size));
+    ctx->stack = &temp_stack;
+  } else {
+    ctx->stack->CheckAllocSizeValid(estimate_alloc_size);
+  }
 
   auto hash_temp_buf = util::TempVectorHolder<uint32_t>(ctx->stack, alloc_batch_size);
   uint32_t* hash_temp = hash_temp_buf.mutable_data();
@@ -473,30 +482,7 @@ Status Hashing32::HashBatch(const ExecBatch& key_batch, uint32_t* hashes,
 
   LightContext ctx;
   ctx.hardware_flags = hardware_flags;
-
-  const int64_t alloc_entry_length = column_arrays[0].length();
-  auto estimate_size = [&] {
-    // An estimate TempVectorStack usage size for Hashing32::HashMultiColumm.
-    const int64_t alloc_size1 =
-        2 * (alloc_entry_length * sizeof(uint32_t) + util::TempVectorStack::meta_size());
-    const int64_t alloc_size2 =
-        alloc_entry_length * sizeof(uint16_t) + util::TempVectorStack::meta_size();
-    return alloc_size1 + alloc_size2;
-  };
-
-  if (!temp_stack) {
-    util::TempVectorStack stack;
-    RETURN_NOT_OK(stack.Init(default_memory_pool(), estimate_size()));
-    ctx.stack = std::move(&stack);
-  } else {
-    auto estimate_alloc_size = estimate_size();
-    ARROW_CHECK_GE(temp_stack->buffer_size(), estimate_alloc_size)
-        << "TempVectorStack's init"
-           " size is not enough. (actual "
-        << temp_stack->buffer_size() << "Bytes, expect " << estimate_alloc_size
-        << "Bytes)";
-    ctx.stack = temp_stack;
-  }
+  ctx.stack = temp_stack;
 
   HashMultiColumn(column_arrays, &ctx, hashes);
   return Status::OK();
@@ -853,6 +839,15 @@ void Hashing64::HashMultiColumn(const std::vector<KeyColumnArray>& cols,
 
   constexpr uint32_t max_batch_size = util::MiniBatch::kMiniBatchLength;
   const uint32_t alloc_batch_size = std::min(num_rows, max_batch_size);
+  const uint64_t estimate_alloc_size = EstimateBatchStackSize<uint64_t>(alloc_batch_size);
+
+  util::TempVectorStack temp_stack;
+  if (!ctx->stack) {
+    ARROW_CHECK_OK(temp_stack.Init(default_memory_pool(), estimate_alloc_size));
+    ctx->stack = &temp_stack;
+  } else {
+    ctx->stack->CheckAllocSizeValid(estimate_alloc_size);
+  }
 
   auto null_indices_buf = util::TempVectorHolder<uint16_t>(ctx->stack, alloc_batch_size);
   uint16_t* null_indices = null_indices_buf.mutable_data();
@@ -936,30 +931,7 @@ Status Hashing64::HashBatch(const ExecBatch& key_batch, uint64_t* hashes,
 
   LightContext ctx;
   ctx.hardware_flags = hardware_flags;
-
-  const int64_t alloc_entry_length = column_arrays[0].length();
-  auto estimate_size = [&] {
-    // An estimate TempVectorStack usage size for Hashing64::HashMultiColumm.
-    const int64_t alloc_size1 =
-        alloc_entry_length * sizeof(uint64_t) + util::TempVectorStack::meta_size();
-    const int64_t alloc_size2 =
-        alloc_entry_length * sizeof(uint16_t) + util::TempVectorStack::meta_size();
-    return alloc_size1 + alloc_size2;
-  };
-
-  if (!temp_stack) {
-    util::TempVectorStack stack;
-    RETURN_NOT_OK(stack.Init(default_memory_pool(), estimate_size()));
-    ctx.stack = std::move(&stack);
-  } else {
-    auto estimate_alloc_size = estimate_size();
-    ARROW_CHECK_GE(temp_stack->buffer_size(), estimate_alloc_size)
-        << "TempVectorStack's init"
-           " size is not enough. (actual "
-        << temp_stack->buffer_size() << "Bytes, expect " << estimate_alloc_size
-        << "Bytes)";
-    ctx.stack = temp_stack;
-  }
+  ctx.stack = temp_stack;
 
   HashMultiColumn(column_arrays, &ctx, hashes);
   return Status::OK();

cpp/src/arrow/compute/key_hash.h

@@ -219,5 +219,24 @@ class ARROW_EXPORT Hashing64 {
                       const uint8_t* keys, uint64_t* hashes);
 };
 
+template <typename T = uint32_t>
+static int64_t EstimateBatchStackSize(int32_t batch_size) {
+  if (sizeof(T) == sizeof(uint32_t)) {
+    const int64_t alloc_for_hash_temp_buf =
+        util::TempVectorStack::EstimateAllocSize(batch_size * sizeof(uint32_t));
+    const int64_t alloc_for_null_hash_temp_buf = alloc_for_hash_temp_buf;
+    const int64_t alloc_for_null_indices_buf =
+        util::TempVectorStack::EstimateAllocSize(batch_size * sizeof(uint16_t));
+    return alloc_for_hash_temp_buf + alloc_for_null_hash_temp_buf +
+           alloc_for_null_indices_buf;
+  } else {
+    const int64_t alloc_for_null_hash_temp_buf =
+        util::TempVectorStack::EstimateAllocSize(batch_size * sizeof(uint64_t));
+    const int64_t alloc_for_null_indices_buf =
+        util::TempVectorStack::EstimateAllocSize(batch_size * sizeof(uint16_t));
+    return alloc_for_null_hash_temp_buf + alloc_for_null_indices_buf;
+  }
+}
+
 }  // namespace compute
 }  // namespace arrow

cpp/src/arrow/compute/key_hash_test.cc

@@ -311,5 +311,32 @@ TEST(VectorHash, FixedLengthTailByteSafety) {
   HashFixedLengthFrom(/*key_length=*/19, /*num_rows=*/64, /*start_row=*/63);
 }
 
+TEST(HashBatch, AllocTempStackAsNeeded) {
+  auto arr = arrow::ArrayFromJSON(arrow::int32(), "[9,2,6]");
+  const int batch_len = arr->length();
+  arrow::compute::ExecBatch exec_batch({arr}, batch_len);
+  auto ctx = arrow::compute::default_exec_context();
+  std::vector<arrow::compute::KeyColumnArray> temp_column_arrays;
+
+  // alloc stack by HashBatch internal
+  std::vector<uint32_t> h1(batch_len);
+  ASSERT_OK(arrow::compute::Hashing32::HashBatch(
+      exec_batch, h1.data(), temp_column_arrays, ctx->cpu_info()->hardware_flags(),
+      nullptr, 0, batch_len));
+
+  // alloc stack as HashBatch needed.
+  util::TempVectorStack stack;
+  ASSERT_OK(
+      stack.Init(default_memory_pool(), EstimateBatchStackSize<int32_t>(batch_len)));
+  std::vector<uint32_t> h2(batch_len);
+  ASSERT_OK(arrow::compute::Hashing32::HashBatch(
+      exec_batch, h2.data(), temp_column_arrays, ctx->cpu_info()->hardware_flags(),
+      &stack, 0, batch_len));
+
+  for (int i = 0; i < batch_len; i++) {
+    EXPECT_EQ(h1[i], h2[i]);
+  }
+}
+
 }  // namespace compute
 }  // namespace arrow

cpp/src/arrow/compute/util.cc

@@ -35,7 +35,7 @@ void TempVectorStack::alloc(uint32_t num_bytes, uint8_t** data, int* id) {
   int64_t new_top = top_ + PaddedAllocationSize(num_bytes) + 2 * sizeof(uint64_t);
   // Stack overflow check (see GH-39582).
   // XXX cannot return a regular Status because most consumers do not either.
-  ARROW_CHECK_LE(new_top, buffer_size_) << "TempVectorStack::alloc overflow";
+  CheckAllocSizeValid(new_top);
   *data = buffer_->mutable_data() + top_ + sizeof(uint64_t);
   // We set 8 bytes before the beginning of the allocated range and
   // 8 bytes after the end to check for stack overflow (which would
@@ -58,6 +58,13 @@ void TempVectorStack::release(int id, uint32_t num_bytes) {
   --num_vectors_;
 }
 
+void TempVectorStack::CheckAllocSizeValid(int64_t estimate_alloc_size) {
+  ARROW_DCHECK_LE(estimate_alloc_size, buffer_size_)
+      << "TempVectorStack alloc overflow."
+         "(Actual "
+      << buffer_size_ << "Bytes, expect " << estimate_alloc_size << "Bytes)";
+}
+
 namespace bit_util {
 
 inline uint64_t SafeLoadUpTo8Bytes(const uint8_t* bytes, int num_bytes) {

cpp/src/arrow/compute/util.h

@@ -89,19 +89,23 @@ class ARROW_EXPORT TempVectorStack {
   Status Init(MemoryPool* pool, int64_t size) {
     num_vectors_ = 0;
     top_ = 0;
-    buffer_size_ = PaddedAllocationSize(size) + kPadding + 2 * sizeof(uint64_t);
+    buffer_size_ = PaddedAllocationSize(size) + 2 * sizeof(uint64_t);
     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);
     buffer_ = std::move(buffer);
     return Status::OK();
   }
 
-  const int64_t buffer_size() const { return buffer_size_; }
-  static int64_t meta_size() { return kPadding + 2 * sizeof(uint64_t); }
+  static int64_t EstimateAllocSize(int64_t size) {
+    return PaddedAllocationSize(size) + 2 * sizeof(uint64_t);
+  }
+
+  int64_t StackBufferSize() const { return buffer_size_; }
+  void CheckAllocSizeValid(int64_t estimate_alloc_size);
 
  private:
-  int64_t PaddedAllocationSize(int64_t num_bytes) {
+  static int64_t PaddedAllocationSize(int64_t num_bytes) {
     // Round up allocation size to multiple of 8 bytes
     // to avoid returning temp vectors with unaligned address.
     //