GH-40069: [C++] Use DCLP to prevent race conditions in reading/writing scalar scratch space

cpp/src/arrow/array/data.cc

@@ -284,24 +284,70 @@ void ArraySpan::SetMembers(const ArrayData& data) {
 namespace {
 
 template <typename offset_type>
-BufferSpan OffsetsForScalar(uint8_t* scratch_space, offset_type value_size) {
- auto* offsets = reinterpret_cast<offset_type*>(scratch_space);
- offsets[0] = 0;
- offsets[1] = static_cast<offset_type>(value_size);
+BufferSpan OffsetsForScalar(const internal::ArraySpanFillFromScalarScratchSpace& scalar,
+ offset_type value_size) {
+ scalar.FillScratchSpace([&](uint8_t* scratch_space) {
+ auto* offsets = reinterpret_cast<offset_type*>(scratch_space);
+ offsets[0] = 0;
+ offsets[1] = static_cast<offset_type>(value_size);
+ });
  static_assert(2 * sizeof(offset_type) <= 16);
- return {scratch_space, sizeof(offset_type) * 2};
+ return {scalar.scratch_space_, sizeof(offset_type) * 2};
 }
 
 template <typename offset_type>
-std::pair<BufferSpan, BufferSpan> OffsetsAndSizesForScalar(uint8_t* scratch_space,
- offset_type value_size) {
- auto* offsets = scratch_space;
- auto* sizes = scratch_space + sizeof(offset_type);
- reinterpret_cast<offset_type*>(offsets)[0] = 0;
- reinterpret_cast<offset_type*>(sizes)[0] = value_size;
+std::pair<BufferSpan, BufferSpan> OffsetsAndSizesForScalar(
+ const internal::ArraySpanFillFromScalarScratchSpace& scalar, offset_type value_size) {
+ scalar.FillScratchSpace([&](uint8_t* scratch_space) {
+ auto* offsets = scratch_space;
+ auto* sizes = scratch_space + sizeof(offset_type);
+ reinterpret_cast<offset_type*>(offsets)[0] = 0;
+ reinterpret_cast<offset_type*>(sizes)[0] = value_size;
+ });
  static_assert(2 * sizeof(offset_type) <= 16);
- return {BufferSpan{offsets, sizeof(offset_type)},
- BufferSpan{sizes, sizeof(offset_type)}};
+ return {BufferSpan{scalar.scratch_space_, sizeof(offset_type)},
+ BufferSpan{scalar.scratch_space_ + sizeof(offset_type), sizeof(offset_type)}};
+}
+
+std::pair<BufferSpan, BufferSpan> TypeCodeAndOffsetsForDenseUnionScalar(
+ const internal::ArraySpanFillFromScalarScratchSpace& scalar, int8_t type_code) {
+ // Dense union needs scratch space to store both offsets and a type code
+ struct UnionScratchSpace {
+ alignas(int64_t) int8_t type_code;
+ alignas(int64_t) uint8_t offsets[sizeof(int32_t) * 2];
+ };
+ auto* union_scratch_space = reinterpret_cast<UnionScratchSpace*>(scalar.scratch_space_);
+ scalar.FillScratchSpace([&](uint8_t*) {
+ union_scratch_space->type_code = type_code;
+ reinterpret_cast<int32_t*>(union_scratch_space->offsets)[0] = 0;
+ reinterpret_cast<int32_t*>(union_scratch_space->offsets)[1] = 1;
+ });
+ static_assert(sizeof(UnionScratchSpace) <=
+ sizeof(internal::ArraySpanFillFromScalarScratchSpace::scratch_space_));
+ return {BufferSpan{reinterpret_cast<uint8_t*>(&union_scratch_space->type_code), 1},
+ BufferSpan{reinterpret_cast<uint8_t*>(&union_scratch_space->offsets),
+ sizeof(int32_t)}};
+}
+
+BufferSpan TypeCodeAndOffsetsForSparseUnionScalar(
+ const internal::ArraySpanFillFromScalarScratchSpace& scalar, int8_t type_code) {
+ scalar.FillScratchSpace([&](uint8_t* scratch_space) {
+ reinterpret_cast<int8_t*>(scratch_space)[0] = type_code;
+ });
+ static_assert(sizeof(int8_t) <=
+ sizeof(internal::ArraySpanFillFromScalarScratchSpace::scratch_space_));
+ return {scalar.scratch_space_, 1};
+}
+
+template <typename run_end_type>
+BufferSpan RunEndForRunEndEncodedScalar(
+ const internal::ArraySpanFillFromScalarScratchSpace& scalar, run_end_type run_end) {
+ scalar.FillScratchSpace([&](uint8_t* scratch_space) {
+ reinterpret_cast<run_end_type*>(scratch_space)[0] = run_end;
+ });
+ static_assert(sizeof(run_end_type) <=
+ sizeof(internal::ArraySpanFillFromScalarScratchSpace::scratch_space_));
+ return {scalar.scratch_space_, sizeof(run_end_type)};
 }
 
 int GetNumBuffers(const DataType& type) {
@@ -415,11 +461,10 @@ void ArraySpan::FillFromScalar(const Scalar& value) {
  data_size = scalar.value->size();
  }
  if (is_binary_like(type_id)) {
- this->buffers[1] =
- OffsetsForScalar(scalar.scratch_space_, static_cast<int32_t>(data_size));
+ this->buffers[1] = OffsetsForScalar(scalar, static_cast<int32_t>(data_size));
  } else {
  // is_large_binary_like
- this->buffers[1] = OffsetsForScalar(scalar.scratch_space_, data_size);
+ this->buffers[1] = OffsetsForScalar(scalar, data_size);
  }
  this->buffers[2].data = const_cast<uint8_t*>(data_buffer);
  this->buffers[2].size = data_size;
@@ -457,16 +502,15 @@ void ArraySpan::FillFromScalar(const Scalar& value) {
  }
 
  if (type_id == Type::LIST || type_id == Type::MAP) {
- this->buffers[1] =
- OffsetsForScalar(scalar.scratch_space_, static_cast<int32_t>(value_length));
+ this->buffers[1] = OffsetsForScalar(scalar, static_cast<int32_t>(value_length));
  } else if (type_id == Type::LARGE_LIST) {
- this->buffers[1] = OffsetsForScalar(scalar.scratch_space_, value_length);
+ this->buffers[1] = OffsetsForScalar(scalar, value_length);
  } else if (type_id == Type::LIST_VIEW) {
- std::tie(this->buffers[1], this->buffers[2]) = OffsetsAndSizesForScalar(
- scalar.scratch_space_, static_cast<int32_t>(value_length));
+ std::tie(this->buffers[1], this->buffers[2]) =
+OffsetsAndSizesForScalar(scalar, static_cast<int32_t>(value_length));
  } else if (type_id == Type::LARGE_LIST_VIEW) {
  std::tie(this->buffers[1], this->buffers[2]) =
- OffsetsAndSizesForScalar(scalar.scratch_space_, value_length);
+ OffsetsAndSizesForScalar(scalar, value_length);
  } else {
  DCHECK_EQ(type_id, Type::FIXED_SIZE_LIST);
  // FIXED_SIZE_LIST: does not have a second buffer
@@ -480,27 +524,14 @@ void ArraySpan::FillFromScalar(const Scalar& value) {
  this->child_data[i].FillFromScalar(*scalar.value[i]);
  }
  } else if (is_union(type_id)) {
- // Dense union needs scratch space to store both offsets and a type code
- struct UnionScratchSpace {
- alignas(int64_t) int8_t type_code;
- alignas(int64_t) uint8_t offsets[sizeof(int32_t) * 2];
- };
- static_assert(sizeof(UnionScratchSpace) <= sizeof(UnionScalar::scratch_space_));
- auto* union_scratch_space = reinterpret_cast<UnionScratchSpace*>(
- &checked_cast<const UnionScalar&>(value).scratch_space_);
-
  // First buffer is kept null since unions have no validity vector
  this->buffers[0] = {};
 
- union_scratch_space->type_code = checked_cast<const UnionScalar&>(value).type_code;
- this->buffers[1].data = reinterpret_cast<uint8_t*>(&union_scratch_space->type_code);
- this->buffers[1].size = 1;
-
  this->child_data.resize(this->type->num_fields());
  if (type_id == Type::DENSE_UNION) {
  const auto& scalar = checked_cast<const DenseUnionScalar&>(value);
- this->buffers[2] =
-OffsetsForScalar(union_scratch_space->offsets, static_cast<int32_t>(1));
+std::tie(this->buffers[1], this->buffers[2]) =
+TypeCodeAndOffsetsForDenseUnionScalar(scalar, scalar.type_code);
  // We can't "see" the other arrays in the union, but we put the "active"
  // union array in the right place and fill zero-length arrays for the
  // others
@@ -517,6 +548,7 @@ void ArraySpan::FillFromScalar(const Scalar& value) {
  }
  } else {
  const auto& scalar = checked_cast<const SparseUnionScalar&>(value);
+ this->buffers[1] = TypeCodeAndOffsetsForSparseUnionScalar(scalar, scalar.type_code);
  // Sparse union scalars have a full complement of child values even
  // though only one of them is relevant, so we just fill them in here
  for (int i = 0; i < static_cast<int>(this->child_data.size()); ++i) {
@@ -539,9 +571,7 @@ void ArraySpan::FillFromScalar(const Scalar& value) {
  e.type = scalar.run_end_type().get();
  e.length = 1;
  e.null_count = 0;
- e.buffers[1].data = scalar.scratch_space_;
- e.buffers[1].size = sizeof(run_end);
- reinterpret_cast<decltype(run_end)*>(scalar.scratch_space_)[0] = run_end;
+ e.buffers[1] = RunEndForRunEndEncodedScalar(scalar, run_end);
  };
 
  switch (scalar.run_end_type()->id()) {

cpp/src/arrow/scalar.cc

@@ -554,6 +554,18 @@ Status Scalar::ValidateFull() const {
  return ScalarValidateImpl(/*full_validation=*/true).Validate(*this);
 }
 
+namespace internal {
+void ArraySpanFillFromScalarScratchSpace::FillScratchSpace(FillScratchSpaceFn fn) const {
+ if (!scratch_space_fill_.filled_.load(std::memory_order_acquire)) {
+ std::lock_guard<std::mutex> lock(scratch_space_fill_.mutex_);
+ if (!scratch_space_fill_.filled_.load(std::memory_order_relaxed)) {
+ fn(scratch_space_);
+ scratch_space_fill_.filled_.store(true, std::memory_order_release);
+ }
+ }
+}
+} // namespace internal
+
 BaseBinaryScalar::BaseBinaryScalar(std::string s, std::shared_ptr<DataType> type)
  : BaseBinaryScalar(Buffer::FromString(std::move(s)), std::move(type)) {}
 

cpp/src/arrow/scalar.h

@@ -22,6 +22,7 @@
 
 #include <iosfwd>
 #include <memory>
+#include <mutex>
 #include <ratio>
 #include <string>
 #include <string_view>
@@ -136,6 +137,27 @@ struct ARROW_EXPORT ArraySpanFillFromScalarScratchSpace {
  // Scalar- including binary scalars where we need to create a buffer
  // that looks like two 32-bit or 64-bit offsets.
  alignas(int64_t) mutable uint8_t scratch_space_[sizeof(int64_t) * 2];
+
+ using FillScratchSpaceFn = std::function<void(uint8_t* scratch_space)>;
+ // Race-free filling of the scratch space.
+ void FillScratchSpace(FillScratchSpaceFn fn) const;
+
+ private:
+ // Structure using DCLP to assists the avoidance of the race conditions in
+ // reading/writing the scratch space.
+ struct ScratchSpaceFill {
+ std::atomic<bool> filled_{false};
+ std::mutex mutex_;
+
+ // Boilerplate code to make this structure default-constructible/copyable/movable, so
+ // that it doesn't affect the default-constructibility/copyability/movability of the
+ // containing class.
+ ScratchSpaceFill() = default;
+ ScratchSpaceFill(const ScratchSpaceFill&) {}
+ ScratchSpaceFill(ScratchSpaceFill&& other) {}
+ ScratchSpaceFill& operator=(const ScratchSpaceFill&) { return *this; }
+ ScratchSpaceFill& operator=(ScratchSpaceFill&&) { return *this; }
+ } mutable scratch_space_fill_;
 };
 
 struct ARROW_EXPORT PrimitiveScalarBase : public Scalar {

cpp/src/arrow/scalar_test.cc

@@ -16,6 +16,7 @@
 // under the License.
 
 #include <chrono>
+#include <future>
 #include <limits>
 #include <memory>
 #include <ostream>
@@ -1984,4 +1985,34 @@ TEST_F(TestExtensionScalar, ValidateErrors) {
  AssertValidationFails(scalar);
 }
 
+TEST(TestScalarFillArraySpan, ParallelFill) {
+ for (auto scalar : {
+ ScalarFromJSON(utf8(), R"("test data")"),
+ // ScalarFromJSON(utf8_view(), R"("test data")"),
+ // ScalarFromJSON(list(int8()), "[1, 2, 3]"),
+ // ScalarFromJSON(list_view(int8()), "[1, 2, 3]"),
+ // TODO: more coming.
+ }) {
+ ARROW_SCOPED_TRACE("Scalar: ", scalar->ToString());
+
+ // Lambda to fill an ArraySpan with the scalar (and consequently fill the scratch
+ // space of the scalar), and use the ArraySpan a bit.
+ auto array_span_from_scalar = [&scalar]() {
+ ArraySpan span;
+ span.FillFromScalar(*scalar);
+ ASSERT_TRUE(span.type->Equals(scalar->type));
+ ASSERT_EQ(span.length, 1);
+ auto values = span.GetValues<int32_t>(1);
+ ASSERT_EQ(values[0], 0);
+ };
+
+ // Two concurrent calls to the lambda are just enough for TSAN to report a race
+ // condition.
+ auto fut1 = std::async(std::launch::async, array_span_from_scalar);
+ auto fut2 = std::async(std::launch::async, array_span_from_scalar);
+ fut1.get();
+ fut2.get();
+ }
+}
+
 } // namespace arrow