This is an automated email from the ASF dual-hosted git repository.

felipecrv pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new 66580441ad GH-41738: [C++] Fix the issue that temp vector stack may be 
under sized (#41746)
66580441ad is described below

commit 66580441ad674c15050710fc1228f3090c855196
Author: Rossi Sun <[email protected]>
AuthorDate: Tue May 21 08:41:20 2024 +0800

    GH-41738: [C++] Fix the issue that temp vector stack may be under sized 
(#41746)
    
    
    
    ### Rationale for this change
    
    See #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: #41738
    
    Authored-by: Ruoxi Sun <[email protected]>
    Signed-off-by: Felipe Oliveira Carvalho <[email protected]>
---
 cpp/src/arrow/compute/CMakeLists.txt        |  3 +-
 cpp/src/arrow/compute/util_internal.cc      |  4 +--
 cpp/src/arrow/compute/util_internal.h       |  2 ++
 cpp/src/arrow/compute/util_internal_test.cc | 52 +++++++++++++++++++++++++++++
 4 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/cpp/src/arrow/compute/CMakeLists.txt 
b/cpp/src/arrow/compute/CMakeLists.txt
index fb778be113..0a8018cd58 100644
--- a/cpp/src/arrow/compute/CMakeLists.txt
+++ b/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)
 
diff --git a/cpp/src/arrow/compute/util_internal.cc 
b/cpp/src/arrow/compute/util_internal.cc
index 9780d1b2f3..7a7875162c 100644
--- a/cpp/src/arrow/compute/util_internal.cc
+++ b/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();
diff --git a/cpp/src/arrow/compute/util_internal.h 
b/cpp/src/arrow/compute/util_internal.h
index f6c5ac1f83..5e5b15a5ff 100644
--- a/cpp/src/arrow/compute/util_internal.h
+++ b/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>
diff --git a/cpp/src/arrow/compute/util_internal_test.cc 
b/cpp/src/arrow/compute/util_internal_test.cc
new file mode 100644
index 0000000000..fbf34f2228
--- /dev/null
+++ b/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

Reply via email to