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