apacheGH-41738: [C++] Fix the issue that temp vector stack may be und…

…er sized (apache#41746)

### Rationale for this change

See apache#41738.

### What changes are included in this PR?

Allocate the underlying buffer of temp stack vector using padded size.

### Are these changes tested?

UT included.

### Are there any user-facing changes?

None.

* GitHub Issue: apache#41738

Authored-by: Ruoxi Sun <[email protected]>
Signed-off-by: Felipe Oliveira Carvalho <[email protected]>

cpp/src/arrow/compute/CMakeLists.txt

@@ -91,7 +91,8 @@ add_arrow_test(internals_test
                registry_test.cc
                key_hash_test.cc
                row/compare_test.cc
-               row/grouper_test.cc)
+               row/grouper_test.cc
+               util_internal_test.cc)
 
 add_arrow_compute_test(expression_test SOURCES expression_test.cc)
 

cpp/src/arrow/compute/util_internal.cc

@@ -39,9 +39,9 @@ Status TempVectorStack::Init(MemoryPool* pool, int64_t size) {
   num_vectors_ = 0;
   top_ = 0;
   buffer_size_ = EstimatedAllocationSize(size);
-  ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateResizableBuffer(size, pool));
+  ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateResizableBuffer(buffer_size_, pool));
 #ifdef ADDRESS_SANITIZER
-  ASAN_POISON_MEMORY_REGION(buffer->mutable_data(), size);
+  ASAN_POISON_MEMORY_REGION(buffer->mutable_data(), buffer_size_);
 #endif
   buffer_ = std::move(buffer);
   return Status::OK();

cpp/src/arrow/compute/util_internal.h

@@ -66,6 +66,8 @@ class ARROW_EXPORT TempVectorStack {
   int64_t top_;
   std::unique_ptr<Buffer> buffer_;
   int64_t buffer_size_;
+
+  friend class TempVectorStackTest;
 };
 
 template <typename T>

cpp/src/arrow/compute/util_internal_test.cc

@@ -0,0 +1,52 @@
+// 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 <gtest/gtest.h>
+
+#include "arrow/buffer.h"
+#include "arrow/compute/util_internal.h"
+#include "arrow/testing/gtest_util.h"
+
+namespace arrow {
+namespace util {
+
+class TempVectorStackTest : public ::testing::Test {
+ protected:
+  static const uint8_t* BufferData(const TempVectorStack& stack) {
+    return stack.buffer_->data();
+  }
+
+  static int64_t BufferCapacity(const TempVectorStack& stack) {
+    return stack.buffer_->capacity();
+  }
+};
+
+// GH-41738: Test the underlying buffer capacity is sufficient to hold the requested
+// vector.
+TEST_F(TempVectorStackTest, BufferCapacitySufficiency) {
+  for (uint32_t stack_size : {1, 7, 8, 63, 64, 65535, 65536}) {
+    ARROW_SCOPED_TRACE("stack_size = ", stack_size);
+    TempVectorStack stack;
+    ASSERT_OK(stack.Init(default_memory_pool(), stack_size));
+
+    TempVectorHolder<uint8_t> v(&stack, stack_size);
+    ASSERT_LE(v.mutable_data() + stack_size, BufferData(stack) + BufferCapacity(stack));
+  }
+}
+
+}  // namespace util
+}  // namespace arrow