[GLUTEN-10933][VL] cuDF: Move lock to ShuffleReader by jinchengchenghh · Pull Request #10934 · apache/gluten

jinchengchenghh · 2025-10-23T15:47:24Z

The lock in WholeStageResultIterator restrict the cpu thread to produce batch, move the lock to here can let threads produce first batch with 1 GB in advance. Maybe the threads should prepare more data and let the GPU consume, this depends on the GPU operator time.

Need to restrict the total stage can offload to GPU, otherwise, after fallback, the lock cannot make effect on GPU execution.

Related issue: #10933

github-actions · 2025-10-23T15:48:02Z

Run Gluten Clickhouse CI on x86

github-actions · 2025-10-23T15:54:35Z

Run Gluten Clickhouse CI on x86

github-actions · 2025-10-24T13:29:05Z

Run Gluten Clickhouse CI on x86

github-actions · 2025-10-28T11:22:00Z

Run Gluten Clickhouse CI on x86

github-actions · 2025-10-28T13:22:25Z

Run Gluten Clickhouse CI on x86

github-actions · 2025-10-28T17:24:41Z

Run Gluten Clickhouse CI on x86

zhztheplayer

Thank you @jinchengchenghh.

A few more questions / comments on this PR's change:

Can we extract common code of GPUColumnarShuffleExchangeExec and ColumnarShuffleExchangeExec out to an abstract class, to avoid duplication?
If putting the lock in shuffle reader, do we still waste CPU if part of the input plan is in vanilla Spark or CPU Velox?
If 2 is true, should we finally implement this concurrency control in Velox? E.g., pass an integer option to control the maximum concurrency of cuDF computation? This way will look more reasonable to me, but I’d appreciate more insights on this.
The implementation doesn't seem to take a corner case into consideration, when there are multiple Velox tasks in the same Spark task. E.g, a Velox plan that has one operator falling back to vanilla Spark between two Velox operators. According to the code, it's possible that two Spark tasks run into race conditions, when the lock is released by task A but the second half of cuDF computation of task A is not finished.

Co-authored-by: Hongze Zhang <hongze.zzz123@gmail.com>

github-actions · 2025-10-30T11:08:34Z

Run Gluten Clickhouse CI on x86

github-actions · 2025-10-30T11:29:04Z

Run Gluten Clickhouse CI on x86

jinchengchenghh · 2025-10-30T11:40:29Z

Refactored.
We will only offload the stage which can totally run in GPU.
Yes, this will cause the lock does not take effect on GPU, now we only check the WholeStageTransformer operator, in my design, the total stage should can offload to GPU, I will add restriction in following PR, do not contain any format conversion in the stage plan.

Thanks for your advice! @zhztheplayer

zhztheplayer · 2025-10-30T12:25:11Z

+arrow::Result<BlockType> readBlockType(arrow::io::InputStream* inputStream) {
+  BlockType type;
+  ARROW_ASSIGN_OR_RAISE(auto bytes, inputStream->Read(sizeof(BlockType), &type));
+  if (bytes == 0) {
+    // Reach EOS.
+    return BlockType::kEndOfStream;
+  }
+  return type;
+}
+
+struct BufferViewReleaser {
+  BufferViewReleaser() : BufferViewReleaser(nullptr) {}
+
+  BufferViewReleaser(std::shared_ptr<arrow::Buffer> arrowBuffer) : bufferReleaser_(std::move(arrowBuffer)) {}
+
+  void addRef() const {}
+
+  void release() const {}
+
+ private:
+  const std::shared_ptr<arrow::Buffer> bufferReleaser_;
+};
+
+BufferPtr wrapInBufferViewAsOwner(const void* buffer, size_t length, std::shared_ptr<arrow::Buffer> bufferReleaser) {
+  return BufferView<BufferViewReleaser>::create(
+      static_cast<const uint8_t*>(buffer), length, {std::move(bufferReleaser)});
+}
+
+BufferPtr convertToVeloxBuffer(std::shared_ptr<arrow::Buffer> buffer) {
+  if (buffer == nullptr) {
+    return nullptr;
+  }
+  return wrapInBufferViewAsOwner(buffer->data(), buffer->size(), buffer);
+}
+
+template <TypeKind Kind, typename T = typename TypeTraits<Kind>::NativeType>
+VectorPtr readFlatVector(
+    std::vector<BufferPtr>& buffers,
+    int32_t& bufferIdx,
+    uint32_t length,
+    std::shared_ptr<const Type> type,
+    const VectorPtr& dictionary,
+    memory::MemoryPool* pool) {
+  auto nulls = buffers[bufferIdx++];
+  auto valuesOrIndices = buffers[bufferIdx++];
+
+  nulls = nulls == nullptr || nulls->size() == 0 ? BufferPtr(nullptr) : nulls;
+
+  if (dictionary != nullptr) {
+    return BaseVector::wrapInDictionary(nulls, valuesOrIndices, length, dictionary);
+  }
+
+  return std::make_shared<FlatVector<T>>(
+      pool, type, nulls, length, std::move(valuesOrIndices), std::vector<BufferPtr>{});
+}
+
+template <>
+VectorPtr readFlatVector<TypeKind::UNKNOWN>(
+    std::vector<BufferPtr>& buffers,
+    int32_t& bufferIdx,
+    uint32_t length,
+    std::shared_ptr<const Type> type,
+    const VectorPtr& dictionary,
+    memory::MemoryPool* pool) {
+  return BaseVector::createNullConstant(type, length, pool);
+}
+
+template <>
+VectorPtr readFlatVector<TypeKind::HUGEINT>(
+    std::vector<BufferPtr>& buffers,
+    int32_t& bufferIdx,
+    uint32_t length,
+    std::shared_ptr<const Type> type,
+    const VectorPtr& dictionary,
+    memory::MemoryPool* pool) {
+  auto nulls = buffers[bufferIdx++];
+  auto valuesOrIndices = buffers[bufferIdx++];
+
+  // Because if buffer does not compress, it will get from netty, the address maynot aligned 16B, which will cause
+  // int128_t = xxx coredump by instruction movdqa
+  const auto* addr = valuesOrIndices->as<facebook::velox::int128_t>();
+  if ((reinterpret_cast<uintptr_t>(addr) & 0xf) != 0) {
+    auto alignedBuffer = AlignedBuffer::allocate<char>(valuesOrIndices->size(), pool);
+    fastCopy(alignedBuffer->asMutable<char>(), valuesOrIndices->as<char>(), valuesOrIndices->size());
+    valuesOrIndices = alignedBuffer;
+  }
+
+  nulls = nulls == nullptr || nulls->size() == 0 ? BufferPtr(nullptr) : nulls;
+
+  if (dictionary != nullptr) {
+    return BaseVector::wrapInDictionary(nulls, valuesOrIndices, length, dictionary);
+  }
+
+  return std::make_shared<FlatVector<int128_t>>(
+      pool, type, nulls, length, std::move(valuesOrIndices), std::vector<BufferPtr>{});
+}
+
+VectorPtr readFlatVectorStringView(
+    std::vector<BufferPtr>& buffers,
+    int32_t& bufferIdx,
+    uint32_t length,
+    std::shared_ptr<const Type> type,
+    const VectorPtr& dictionary,
+    memory::MemoryPool* pool) {
+  auto nulls = buffers[bufferIdx++];
+  auto lengthOrIndices = buffers[bufferIdx++];
+
+  nulls = nulls == nullptr || nulls->size() == 0 ? BufferPtr(nullptr) : nulls;
+
+  if (dictionary != nullptr) {
+    return BaseVector::wrapInDictionary(nulls, lengthOrIndices, length, dictionary);
+  }
+
+  auto valueBuffer = buffers[bufferIdx++];
+
+  const auto* rawLength = lengthOrIndices->as<StringLengthType>();
+  const auto* valueBufferPtr = valueBuffer->as<char>();
+
+  auto values = AlignedBuffer::allocate<char>(sizeof(StringView) * length, pool);
+  auto* rawValues = values->asMutable<StringView>();
+
+  uint64_t offset = 0;
+  for (int32_t i = 0; i < length; ++i) {
+    rawValues[i] = StringView(valueBufferPtr + offset, rawLength[i]);
+    offset += rawLength[i];
+  }
+
+  std::vector<BufferPtr> stringBuffers;
+  stringBuffers.emplace_back(valueBuffer);
+
+  return std::make_shared<FlatVector<StringView>>(
+      pool, type, nulls, length, std::move(values), std::move(stringBuffers));
+}
+
+template <>
+VectorPtr readFlatVector<TypeKind::VARCHAR>(
+    std::vector<BufferPtr>& buffers,
+    int32_t& bufferIdx,
+    uint32_t length,
+    std::shared_ptr<const Type> type,
+    const VectorPtr& dictionary,
+    memory::MemoryPool* pool) {
+  return readFlatVectorStringView(buffers, bufferIdx, length, type, dictionary, pool);
+}
+
+template <>
+VectorPtr readFlatVector<TypeKind::VARBINARY>(
+    std::vector<BufferPtr>& buffers,
+    int32_t& bufferIdx,
+    uint32_t length,
+    std::shared_ptr<const Type> type,
+    const VectorPtr& dictionary,
+    memory::MemoryPool* pool) {
+  return readFlatVectorStringView(buffers, bufferIdx, length, type, dictionary, pool);
+}
+
+std::unique_ptr<ByteInputStream> toByteStream(uint8_t* data, int32_t size) {
+  std::vector<ByteRange> byteRanges;
+  byteRanges.push_back(ByteRange{data, size, 0});
+  auto byteStream = std::make_unique<BufferInputStream>(byteRanges);
+  return byteStream;
+}
+
+RowVectorPtr readComplexType(BufferPtr buffer, RowTypePtr& rowType, memory::MemoryPool* pool) {
+  RowVectorPtr result;
+  auto byteStream = toByteStream(const_cast<uint8_t*>(buffer->as<uint8_t>()), buffer->size());
+  auto serde = std::make_unique<serializer::presto::PrestoVectorSerde>();
+  serializer::presto::PrestoVectorSerde::PrestoOptions options;
+  options.useLosslessTimestamp = true;
+  serde->deserialize(byteStream.get(), pool, rowType, &result, &options);
+  return result;
+}
+
+RowTypePtr getComplexWriteType(const std::vector<TypePtr>& types) {
+  std::vector<std::string> complexTypeColNames;
+  std::vector<TypePtr> complexTypeChildrens;
+  for (int32_t i = 0; i < types.size(); ++i) {
+    auto kind = types[i]->kind();
+    switch (kind) {
+      case TypeKind::ROW:
+      case TypeKind::MAP:
+      case TypeKind::ARRAY: {
+        complexTypeColNames.emplace_back(types[i]->name());
+        complexTypeChildrens.emplace_back(types[i]);
+      } break;
+      default:
+        break;
+    }
+  }
+  return std::make_shared<const RowType>(std::move(complexTypeColNames), std::move(complexTypeChildrens));
+}
+
+RowVectorPtr deserialize(
+    RowTypePtr type,
+    uint32_t numRows,
+    std::vector<BufferPtr>& buffers,
+    const std::vector<int32_t>& dictionaryFields,
+    const std::vector<VectorPtr>& dictionaries,
+    memory::MemoryPool* pool) {
+  std::vector<VectorPtr> children;
+  auto types = type->as<TypeKind::ROW>().children();
+
+  std::vector<VectorPtr> complexChildren;
+  auto complexRowType = getComplexWriteType(types);
+  if (complexRowType->children().size() > 0) {
+    complexChildren = readComplexType(buffers[buffers.size() - 1], complexRowType, pool)->children();
+  }
+
+  int32_t bufferIdx = 0;
+  int32_t complexIdx = 0;
+  int32_t dictionaryIdx = 0;
+  for (size_t i = 0; i < types.size(); ++i) {
+    const auto kind = types[i]->kind();
+    switch (kind) {
+      case TypeKind::ROW:
+      case TypeKind::MAP:
+      case TypeKind::ARRAY: {
+        children.emplace_back(std::move(complexChildren[complexIdx]));
+        complexIdx++;
+      } break;
+      default: {
+        VectorPtr dictionary{nullptr};
+        if (!dictionaryFields.empty() && dictionaryIdx < dictionaryFields.size() &&
+            dictionaryFields[dictionaryIdx] == i) {
+          dictionary = dictionaries[dictionaryIdx++];
+        }
+        auto res = VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH_ALL(
+            readFlatVector, kind, buffers, bufferIdx, numRows, types[i], dictionary, pool);
+        children.emplace_back(std::move(res));
+      } break;
+    }
+  }
+
+  return std::make_shared<RowVector>(pool, type, BufferPtr(nullptr), numRows, children);
+}
+
+std::shared_ptr<VeloxColumnarBatch> makeColumnarBatch(
+    RowTypePtr type,
+    uint32_t numRows,
+    std::vector<std::shared_ptr<arrow::Buffer>> arrowBuffers,
+    const std::vector<int32_t>& dictionaryFields,
+    const std::vector<VectorPtr>& dictionaries,
+    memory::MemoryPool* pool,
+    int64_t& deserializeTime) {
+  ScopedTimer timer(&deserializeTime);
+  std::vector<BufferPtr> veloxBuffers;
+  veloxBuffers.reserve(arrowBuffers.size());
+  for (auto& buffer : arrowBuffers) {
+    veloxBuffers.push_back(convertToVeloxBuffer(std::move(buffer)));
+  }
+  auto rowVector = deserialize(type, numRows, veloxBuffers, dictionaryFields, dictionaries, pool);
+  return std::make_shared<VeloxColumnarBatch>(std::move(rowVector));
+}
+
+} // namespace


Hi @jinchengchenghh, any idea on reducing the code duplication of this part?

Should we just put the GPU reader implementation to VeloxShuffleReader.h / VeloxShuffleReader.cc ? cc @marin-ma

These all will be removed after support Buffer to cudf::table, a draft version is https://github.com/jinchengchenghh/gluten/blob/shuffle/cpp/velox/shuffle/GpuShuffleReader.cc

Thank you for the explanation.

zhztheplayer

Thank you. I only went through the design and code roughly, feel free to call for other committers for review if you think is needed.

marin-ma · 2025-10-30T15:01:13Z

      metrics: Map[String, SQLMetric],
-      shuffleWriterType: ShuffleWriterType): Serializer = {
+      shuffleWriterType: ShuffleWriterType,
+      enableCudf: Boolean): Serializer = {


Just noticed this change. Is it possible to just add one or more new ShuffleWriterType for cuff so we can avoid adding new parameters?

Shuffle writer is same, only the shuffle reader is different. Each ShuffleWriterType matches to its own GpuShuffleReader

The shuffle writer type is also used on the reader side to determine which deserialiser to use. If the deserialisation are different, then would be better to have different shuffle writer type like gpu_hash, gpu_sort, etc.

Make sense, and I notice the shuffle writer type name exists in cpp and scala code, should we use protobuf for them?

Sure. Would be nice to have a better way to align them.

The lock in WholeStageResultIterator restrict the cpu thread to produce batch, move the lock to here can let threads produce first batch with 1 GB in advance. Maybe the threads should prepare more data and let the GPU consume, this depends on the GPU operator time. Need to restrict the total stage can offload to GPU, otherwise, after fallback, the lock cannot make effect on GPU execution.

github-actions Bot added CORE works for Gluten Core VELOX labels Oct 23, 2025

jinchengchenghh force-pushed the validation branch from bc48720 to 3e61e24 Compare October 24, 2025 13:28

jinchengchenghh added 3 commits October 28, 2025 11:20

[GLUTEN-10933][VL] cuDF: Move lock to ShuffleReader

2950b85

clean some

100ff3a

fix test

6b5bad9

jinchengchenghh force-pushed the validation branch from 3e61e24 to 6b5bad9 Compare October 28, 2025 11:21

fix

0261bc7

fix build

95a3ac6

github-actions Bot added the CLICKHOUSE label Oct 28, 2025

jinchengchenghh requested a review from zhztheplayer October 29, 2025 10:15

zhztheplayer reviewed Oct 29, 2025

View reviewed changes

Comment thread cpp/velox/cudf/GpuLock.cc Outdated

Update cpp/velox/cudf/GpuLock.cc

a9ecaed

Co-authored-by: Hongze Zhang <hongze.zzz123@gmail.com>

jinchengchenghh added 2 commits October 30, 2025 11:25

refactor ColumnarShuffleExchangeExec

729089c

address comment to move the static member to static struct

a5ada79

zhztheplayer reviewed Oct 30, 2025

View reviewed changes

zhztheplayer approved these changes Oct 30, 2025

View reviewed changes

jinchengchenghh merged commit aeb2c67 into apache:main Oct 30, 2025
60 checks passed

marin-ma reviewed Oct 30, 2025

View reviewed changes

Uh oh!

Conversation

jinchengchenghh commented Oct 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Oct 23, 2025

Uh oh!

github-actions Bot commented Oct 23, 2025

Uh oh!

github-actions Bot commented Oct 24, 2025

Uh oh!

github-actions Bot commented Oct 28, 2025

Uh oh!

github-actions Bot commented Oct 28, 2025

Uh oh!

github-actions Bot commented Oct 28, 2025

Uh oh!

zhztheplayer left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot commented Oct 30, 2025

Uh oh!

github-actions Bot commented Oct 30, 2025

Uh oh!

jinchengchenghh commented Oct 30, 2025

Uh oh!

zhztheplayer Oct 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jinchengchenghh Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

zhztheplayer Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

zhztheplayer left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

marin-ma Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

jinchengchenghh Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

marin-ma Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

jinchengchenghh Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

marin-ma Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jinchengchenghh commented Oct 23, 2025 •

edited

Loading

zhztheplayer Oct 30, 2025 •

edited

Loading