Move temp stack vector into internal header and reduce including

cpp/src/arrow/CMakeLists.txt

@@ -716,7 +716,8 @@ set(ARROW_COMPUTE_SRCS
  compute/row/compare_internal.cc
  compute/row/grouper.cc
  compute/row/row_internal.cc
- compute/util.cc)
+ compute/util.cc
+ compute/util_internal.cc)
 
 append_runtime_avx2_src(ARROW_COMPUTE_SRCS compute/key_hash_internal_avx2.cc)
 append_runtime_avx2_bmi2_src(ARROW_COMPUTE_SRCS compute/key_map_internal_avx2.cc)

cpp/src/arrow/compute/key_hash_internal.h

@@ -48,8 +48,12 @@ class ARROW_EXPORT Hashing32 {
  static void HashMultiColumn(const std::vector<KeyColumnArray>& cols, LightContext* ctx,
  uint32_t* out_hash);
 
- // Clarify the max temp stack usage for HashBatch so the caller could reserve enough
- // size in advance.
+ // Clarify the max temp stack usage for HashBatch, which might be necessary for the
+ // caller to be aware of at compile time to reserve enough stack size in advance. The
+ // HashBatch implementation uses one uint32 temp vector as a buffer for hash, one uint16
+ // temp vector as a buffer for null indices and one uint32 temp vector as a buffer for
+ // null hash, all are of size kMiniBatchLength. Plus extra kMiniBatchLength to cope with
+ // stack padding and aligning.
  static constexpr auto kHashBatchTempStackUsage =
  (sizeof(uint32_t) + sizeof(uint16_t) + sizeof(uint32_t) + /*extra=*/1) *
  util::MiniBatch::kMiniBatchLength;
@@ -167,8 +171,11 @@ class ARROW_EXPORT Hashing64 {
  static void HashMultiColumn(const std::vector<KeyColumnArray>& cols, LightContext* ctx,
  uint64_t* hashes);
 
- // Clarify the max temp stack usage for HashBatch so the caller could reserve enough
- // size in advance.
+ // Clarify the max temp stack usage for HashBatch, which might be necessary for the
+ // caller to be aware of at compile time to reserve enough stack size in advance. The
+ // HashBatch implementation uses one uint16 temp vector as a buffer for null indices and
+ // one uint64 temp vector as a buffer for null hash, all are of size kMiniBatchLength.
+ // Plus extra kMiniBatchLength to cope with stack padding and aligning.
  static constexpr auto kHashBatchTempStackUsage =
  (sizeof(uint16_t) + sizeof(uint64_t) + /*extra=*/1) *
  util::MiniBatch::kMiniBatchLength;

cpp/src/arrow/compute/key_map_internal.h

@@ -21,6 +21,7 @@
 #include <functional>
 
 #include "arrow/compute/util.h"
+#include "arrow/compute/util_internal.h"
 #include "arrow/result.h"
 #include "arrow/status.h"
 #include "arrow/type_fwd.h"

cpp/src/arrow/compute/light_array_internal.h

@@ -22,6 +22,7 @@
 #include "arrow/array.h"
 #include "arrow/compute/exec.h"
 #include "arrow/compute/util.h"
+#include "arrow/compute/util_internal.h"
 #include "arrow/type.h"
 #include "arrow/util/cpu_info.h"
 #include "arrow/util/logging.h"

cpp/src/arrow/compute/light_array_test.cc

@@ -20,6 +20,7 @@
 #include <gtest/gtest.h>
 #include <numeric>
 
+#include "arrow/memory_pool.h"
 #include "arrow/testing/generator.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/type.h"

cpp/src/arrow/compute/row/compare_internal.h

@@ -32,8 +32,11 @@ namespace compute {
 
 class ARROW_EXPORT KeyCompare {
  public:
- // Clarify the max temp stack usage for CompareColumnsToRows so the caller could reserve
- // enough size in advance.
+ // Clarify the max temp stack usage for CompareColumnsToRows, which might be necessary
+ // for the caller to be aware of (possibly at compile time) to reserve enough stack size
+ // in advance. The CompareColumnsToRows implementation uses three uint8 temp vectors as
+ // buffers for match vectors, all are of size num_rows. Plus extra kMiniBatchLength to
+ // cope with stack padding and aligning.
  constexpr static int64_t CompareColumnsToRowsTempStackUsage(int64_t num_rows) {
  return (sizeof(uint8_t) + sizeof(uint8_t) + sizeof(uint8_t)) * num_rows +
  /*extra=*/util::MiniBatch::kMiniBatchLength;

cpp/src/arrow/compute/util.cc

@@ -17,11 +17,7 @@
 
 #include "arrow/compute/util.h"
 
-#include "arrow/table.h"
-#include "arrow/util/bit_util.h"
-#include "arrow/util/bitmap_ops.h"
 #include "arrow/util/logging.h"
-#include "arrow/util/tracing_internal.h"
 #include "arrow/util/ubsan.h"
 
 namespace arrow {
@@ -31,36 +27,6 @@ using internal::CpuInfo;
 
 namespace util {
 
-void TempVectorStack::alloc(uint32_t num_bytes, uint8_t** data, int* id) {
- int64_t estimated_alloc_size = EstimatedAllocationSize(num_bytes);
- int64_t new_top = top_ + estimated_alloc_size;
- // 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: allocating " << estimated_alloc_size
- << " on top of " << top_ << " in stack of size " << buffer_size_;
- *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
- // result in those known bytes being corrupted).
- reinterpret_cast<uint64_t*>(buffer_->mutable_data() + top_)[0] = kGuard1;
- reinterpret_cast<uint64_t*>(buffer_->mutable_data() + new_top)[-1] = kGuard2;
- *id = num_vectors_++;
- top_ = new_top;
-}
-
-void TempVectorStack::release(int id, uint32_t num_bytes) {
- ARROW_DCHECK(num_vectors_ == id + 1);
- int64_t size = EstimatedAllocationSize(num_bytes);
- ARROW_DCHECK(reinterpret_cast<const uint64_t*>(buffer_->mutable_data() + top_)[-1] ==
- kGuard2);
- ARROW_DCHECK(top_ >= size);
- top_ -= size;
- ARROW_DCHECK(reinterpret_cast<const uint64_t*>(buffer_->mutable_data() + top_)[0] ==
- kGuard1);
- --num_vectors_;
-}
-
 namespace bit_util {
 
 inline uint64_t SafeLoadUpTo8Bytes(const uint8_t* bytes, int num_bytes) {

cpp/src/arrow/compute/util.h

@@ -24,17 +24,10 @@
 #include <unordered_map>
 #include <vector>
 
-#include "arrow/buffer.h"
 #include "arrow/compute/expression.h"
 #include "arrow/compute/type_fwd.h"
-#include "arrow/memory_pool.h"
 #include "arrow/result.h"
-#include "arrow/status.h"
-#include "arrow/util/bit_util.h"
 #include "arrow/util/cpu_info.h"
-#include "arrow/util/mutex.h"
-#include "arrow/util/thread_pool.h"
-#include "arrow/util/type_fwd.h"
 
 #if defined(__clang__) || defined(__GNUC__)
 #define BYTESWAP(x) __builtin_bswap64(x)
@@ -77,72 +70,6 @@ class MiniBatch {
  static constexpr int kMiniBatchLength = 1 << kLogMiniBatchLength;
 };
 
-/// Storage used to allocate temporary vectors of a batch size.
-/// Temporary vectors should resemble allocating temporary variables on the stack
-/// but in the context of vectorized processing where we need to store a vector of
-/// temporaries instead of a single value.
-class ARROW_EXPORT TempVectorStack {
- template <typename>
- friend class TempVectorHolder;
-
- public:
- Status Init(MemoryPool* pool, int64_t size) {
- num_vectors_ = 0;
- top_ = 0;
- buffer_size_ = EstimatedAllocationSize(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);
- buffer_ = std::move(buffer);
- return Status::OK();
- }
-
- private:
- static int64_t EstimatedAllocationSize(int64_t size) {
- return PaddedAllocationSize(size) + 2 * sizeof(uint64_t);
- }
-
- 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.
- //
- // Also add padding at the end to facilitate loads and stores
- // using SIMD when number of vector elements is not divisible
- // by the number of SIMD lanes.
- //
- return ::arrow::bit_util::RoundUp(num_bytes, sizeof(int64_t)) + kPadding;
- }
- void alloc(uint32_t num_bytes, uint8_t** data, int* id);
- void release(int id, uint32_t num_bytes);
- static constexpr uint64_t kGuard1 = 0x3141592653589793ULL;
- static constexpr uint64_t kGuard2 = 0x0577215664901532ULL;
- static constexpr int64_t kPadding = 64;
- int num_vectors_;
- int64_t top_;
- std::unique_ptr<Buffer> buffer_;
- int64_t buffer_size_;
-};
-
-template <typename T>
-class TempVectorHolder {
- friend class TempVectorStack;
-
- public:
- ~TempVectorHolder() { stack_->release(id_, num_elements_ * sizeof(T)); }
- T* mutable_data() { return reinterpret_cast<T*>(data_); }
- TempVectorHolder(TempVectorStack* stack, uint32_t num_elements) {
- stack_ = stack;
- num_elements_ = num_elements;
- stack_->alloc(num_elements * sizeof(T), &data_, &id_);
- }
-
- private:
- TempVectorStack* stack_;
- uint8_t* data_;
- int id_;
- uint32_t num_elements_;
-};
-
 namespace bit_util {
 
 ARROW_EXPORT void bits_to_indexes(int bit_to_search, int64_t hardware_flags,

cpp/src/arrow/compute/util_internal.cc

@@ -0,0 +1,79 @@
+// 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 "arrow/compute/util_internal.h"
+
+#include "arrow/compute/util.h"
+#include "arrow/memory_pool.h"
+
+namespace arrow {
+namespace util {
+
+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));
+ // Ensure later operations don't accidentally read uninitialized memory.
+ std::memset(buffer->mutable_data(), 0xFF, size);
+ buffer_ = std::move(buffer);
+ return Status::OK();
+}
+
+int64_t TempVectorStack::PaddedAllocationSize(int64_t num_bytes) {
+ // Round up allocation size to multiple of 8 bytes
+ // to avoid returning temp vectors with unaligned address.
+ //
+ // Also add padding at the end to facilitate loads and stores
+ // using SIMD when number of vector elements is not divisible
+ // by the number of SIMD lanes.
+ //
+ return ::arrow::bit_util::RoundUp(num_bytes, sizeof(int64_t)) + kPadding;
+}
+
+void TempVectorStack::alloc(uint32_t num_bytes, uint8_t** data, int* id) {
+ int64_t estimated_alloc_size = EstimatedAllocationSize(num_bytes);
+ int64_t new_top = top_ + estimated_alloc_size;
+ // 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: allocating " << estimated_alloc_size
+ << " on top of " << top_ << " in stack of size " << buffer_size_;
+ *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
+ // result in those known bytes being corrupted).
+ reinterpret_cast<uint64_t*>(buffer_->mutable_data() + top_)[0] = kGuard1;
+ reinterpret_cast<uint64_t*>(buffer_->mutable_data() + new_top)[-1] = kGuard2;
+ *id = num_vectors_++;
+ top_ = new_top;
+}
+
+void TempVectorStack::release(int id, uint32_t num_bytes) {
+ ARROW_DCHECK(num_vectors_ == id + 1);
+ int64_t size = EstimatedAllocationSize(num_bytes);
+ ARROW_DCHECK(reinterpret_cast<const uint64_t*>(buffer_->mutable_data() + top_)[-1] ==
+ kGuard2);
+ ARROW_DCHECK(top_ >= size);
+ top_ -= size;
+ ARROW_DCHECK(reinterpret_cast<const uint64_t*>(buffer_->mutable_data() + top_)[0] ==
+ kGuard1);
+ --num_vectors_;
+}
+
+} // namespace util
+} // namespace arrow

cpp/src/arrow/compute/util_internal.h

@@ -17,6 +17,8 @@
 
 #pragma once
 
+#include "arrow/status.h"
+#include "arrow/type_fwd.h"
 #include "arrow/util/logging.h"
 
 namespace arrow {
@@ -27,5 +29,54 @@ void CheckAlignment(const void* ptr) {
  ARROW_DCHECK(reinterpret_cast<uint64_t>(ptr) % sizeof(T) == 0);
 }
 
+/// Storage used to allocate temporary vectors of a batch size.
+/// Temporary vectors should resemble allocating temporary variables on the stack
+/// but in the context of vectorized processing where we need to store a vector of
+/// temporaries instead of a single value.
+class ARROW_EXPORT TempVectorStack {
+ template <typename>
+ friend class TempVectorHolder;
+
+ public:
+ Status Init(MemoryPool* pool, int64_t size);
+
+ private:
+ static int64_t EstimatedAllocationSize(int64_t size) {
+ return PaddedAllocationSize(size) + 2 * sizeof(uint64_t);
+ }
+
+ static int64_t PaddedAllocationSize(int64_t num_bytes);
+
+ void alloc(uint32_t num_bytes, uint8_t** data, int* id);
+ void release(int id, uint32_t num_bytes);
+ static constexpr uint64_t kGuard1 = 0x3141592653589793ULL;
+ static constexpr uint64_t kGuard2 = 0x0577215664901532ULL;
+ static constexpr int64_t kPadding = 64;
+ int num_vectors_;
+ int64_t top_;
+ std::unique_ptr<Buffer> buffer_;
+ int64_t buffer_size_;
+};
+
+template <typename T>
+class TempVectorHolder {
+ friend class TempVectorStack;
+
+ public:
+ ~TempVectorHolder() { stack_->release(id_, num_elements_ * sizeof(T)); }
+ T* mutable_data() { return reinterpret_cast<T*>(data_); }
+ TempVectorHolder(TempVectorStack* stack, uint32_t num_elements) {
+ stack_ = stack;
+ num_elements_ = num_elements;
+ stack_->alloc(num_elements * sizeof(T), &data_, &id_);
+ }
+
+ private:
+ TempVectorStack* stack_;
+ uint8_t* data_;
+ int id_;
+ uint32_t num_elements_;
+};
+
 } // namespace util
 } // namespace arrow