This is an automated email from the ASF dual-hosted git repository.
marong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 1544307d0c [GLUTEN-6896] Add buffered read for hash/sort shuffle
reader (#7897)
1544307d0c is described below
commit 1544307d0c91133fbc059ec3f149869c9701c0f9
Author: Rong Ma <[email protected]>
AuthorDate: Fri Nov 15 10:42:07 2024 +0800
[GLUTEN-6896] Add buffered read for hash/sort shuffle reader (#7897)
---
.../vectorized/ColumnarBatchSerializer.scala | 2 ++
cpp/core/jni/JniWrapper.cc | 3 +-
cpp/core/shuffle/Options.h | 2 ++
cpp/core/shuffle/Payload.cc | 18 +++++++----
cpp/core/shuffle/Payload.h | 1 +
cpp/velox/compute/VeloxRuntime.cc | 1 +
cpp/velox/shuffle/VeloxShuffleReader.cc | 37 ++++++++++++++--------
cpp/velox/shuffle/VeloxShuffleReader.h | 7 +++-
cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h | 1 +
docs/Configuration.md | 1 +
.../gluten/vectorized/ShuffleReaderJniWrapper.java | 1 +
.../VeloxCelebornColumnarBatchSerializer.scala | 2 ++
.../scala/org/apache/gluten/GlutenConfig.scala | 13 ++++++++
13 files changed, 68 insertions(+), 21 deletions(-)
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala
b/backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala
index fa4d92652c..88215e36bb 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala
@@ -98,6 +98,7 @@ private class ColumnarBatchSerializerInstance(
val compressionCodecBackend =
GlutenConfig.getConf.columnarShuffleCodecBackend.orNull
val batchSize = GlutenConfig.getConf.maxBatchSize
+ val bufferSize = GlutenConfig.getConf.columnarShuffleReaderBufferSize
val runtime = Runtimes.contextInstance("ShuffleReader")
val jniWrapper = ShuffleReaderJniWrapper.create(runtime)
val shuffleReaderHandle = jniWrapper.make(
@@ -105,6 +106,7 @@ private class ColumnarBatchSerializerInstance(
compressionCodec,
compressionCodecBackend,
batchSize,
+ bufferSize,
shuffleWriterType)
// Close shuffle reader instance as lately as the end of task processing,
// since the native reader could hold a reference to memory pool that
diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc
index 6a0a5b0057..9da5589486 100644
--- a/cpp/core/jni/JniWrapper.cc
+++ b/cpp/core/jni/JniWrapper.cc
@@ -996,6 +996,7 @@ JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_ShuffleReaderJniWrappe
jstring compressionType,
jstring compressionBackend,
jint batchSize,
+ jlong bufferSize,
jstring shuffleWriterType) {
JNI_METHOD_START
auto ctx = getRuntime(env, wrapper);
@@ -1007,7 +1008,7 @@ JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_ShuffleReaderJniWrappe
options.codecBackend = getCodecBackend(env, compressionBackend);
}
options.batchSize = batchSize;
- // TODO: Add coalesce option and maximum coalesced size.
+ options.bufferSize = bufferSize;
options.shuffleWriterType =
ShuffleWriter::stringToType(jStringToCString(env, shuffleWriterType));
std::shared_ptr<arrow::Schema> schema =
diff --git a/cpp/core/shuffle/Options.h b/cpp/core/shuffle/Options.h
index a3dc9f6260..6a9e0ec4b3 100644
--- a/cpp/core/shuffle/Options.h
+++ b/cpp/core/shuffle/Options.h
@@ -38,6 +38,7 @@ static constexpr double kDefaultMergeBufferThreshold = 0.25;
static constexpr bool kEnableBufferedWrite = true;
static constexpr bool kDefaultUseRadixSort = true;
static constexpr int32_t kDefaultSortBufferSize = 4096;
+static constexpr int64_t kDefaultReadBufferSize = 1 << 20;
enum ShuffleWriterType { kHashShuffle, kSortShuffle, kRssSortShuffle };
enum PartitionWriterType { kLocal, kRss };
@@ -49,6 +50,7 @@ struct ShuffleReaderOptions {
ShuffleWriterType shuffleWriterType = kHashShuffle;
CodecBackend codecBackend = CodecBackend::NONE;
int32_t batchSize = kDefaultBatchSize;
+ int64_t bufferSize = kDefaultReadBufferSize;
};
struct ShuffleWriterOptions {
diff --git a/cpp/core/shuffle/Payload.cc b/cpp/core/shuffle/Payload.cc
index d5cdb6d08e..55f3a43396 100644
--- a/cpp/core/shuffle/Payload.cc
+++ b/cpp/core/shuffle/Payload.cc
@@ -118,9 +118,9 @@ arrow::Status compressAndFlush(
return arrow::Status::OK();
}
-arrow::Result<std::shared_ptr<arrow::Buffer>> readUncompressedBuffer(
- arrow::io::InputStream* inputStream,
- arrow::MemoryPool* pool) {
+arrow::Result<std::shared_ptr<arrow::Buffer>>
+readUncompressedBuffer(arrow::io::InputStream* inputStream, arrow::MemoryPool*
pool, int64_t& deserializedTime) {
+ ScopedTimer timer(&deserializedTime);
int64_t bufferLength;
RETURN_NOT_OK(inputStream->Read(sizeof(int64_t), &bufferLength));
if (bufferLength == kNullBuffer) {
@@ -135,7 +135,9 @@ arrow::Result<std::shared_ptr<arrow::Buffer>>
readCompressedBuffer(
arrow::io::InputStream* inputStream,
const std::shared_ptr<arrow::util::Codec>& codec,
arrow::MemoryPool* pool,
+ int64_t& deserializeTime,
int64_t& decompressTime) {
+ ScopedTimer timer(&deserializeTime);
int64_t compressedLength;
RETURN_NOT_OK(inputStream->Read(sizeof(int64_t), &compressedLength));
if (compressedLength == kNullBuffer) {
@@ -155,7 +157,7 @@ arrow::Result<std::shared_ptr<arrow::Buffer>>
readCompressedBuffer(
ARROW_ASSIGN_OR_RAISE(auto compressed,
arrow::AllocateResizableBuffer(compressedLength, pool));
RETURN_NOT_OK(inputStream->Read(compressedLength,
compressed->mutable_data()));
- ScopedTimer timer(&decompressTime);
+ timer.switchTo(&decompressTime);
ARROW_ASSIGN_OR_RAISE(auto output,
arrow::AllocateResizableBuffer(uncompressedLength, pool));
RETURN_NOT_OK(codec->Decompress(compressedLength, compressed->data(),
uncompressedLength, output->mutable_data()));
return output;
@@ -296,7 +298,9 @@ arrow::Result<std::vector<std::shared_ptr<arrow::Buffer>>>
BlockPayload::deseria
const std::shared_ptr<arrow::util::Codec>& codec,
arrow::MemoryPool* pool,
uint32_t& numRows,
+ int64_t& deserializeTime,
int64_t& decompressTime) {
+ auto timer = std::make_unique<ScopedTimer>(&deserializeTime);
static const std::vector<std::shared_ptr<arrow::Buffer>> kEmptyBuffers{};
ARROW_ASSIGN_OR_RAISE(auto type, readType(inputStream));
if (type == 0) {
@@ -306,6 +310,7 @@ arrow::Result<std::vector<std::shared_ptr<arrow::Buffer>>>
BlockPayload::deseria
RETURN_NOT_OK(inputStream->Read(sizeof(uint32_t), &numRows));
uint32_t numBuffers;
RETURN_NOT_OK(inputStream->Read(sizeof(uint32_t), &numBuffers));
+ timer.reset();
bool isCompressionEnabled = type == Type::kCompressed;
std::vector<std::shared_ptr<arrow::Buffer>> buffers;
@@ -313,9 +318,10 @@ arrow::Result<std::vector<std::shared_ptr<arrow::Buffer>>>
BlockPayload::deseria
for (auto i = 0; i < numBuffers; ++i) {
buffers.emplace_back();
if (isCompressionEnabled) {
- ARROW_ASSIGN_OR_RAISE(buffers.back(), readCompressedBuffer(inputStream,
codec, pool, decompressTime));
+ ARROW_ASSIGN_OR_RAISE(
+ buffers.back(), readCompressedBuffer(inputStream, codec, pool,
deserializeTime, decompressTime));
} else {
- ARROW_ASSIGN_OR_RAISE(buffers.back(),
readUncompressedBuffer(inputStream, pool));
+ ARROW_ASSIGN_OR_RAISE(buffers.back(),
readUncompressedBuffer(inputStream, pool, deserializeTime));
}
}
return buffers;
diff --git a/cpp/core/shuffle/Payload.h b/cpp/core/shuffle/Payload.h
index ea8c897e96..611b2310d5 100644
--- a/cpp/core/shuffle/Payload.h
+++ b/cpp/core/shuffle/Payload.h
@@ -92,6 +92,7 @@ class BlockPayload final : public Payload {
const std::shared_ptr<arrow::util::Codec>& codec,
arrow::MemoryPool* pool,
uint32_t& numRows,
+ int64_t& deserializeTime,
int64_t& decompressTime);
static int64_t maxCompressedLength(
diff --git a/cpp/velox/compute/VeloxRuntime.cc
b/cpp/velox/compute/VeloxRuntime.cc
index 332c75dbd7..4c6b52e6fe 100644
--- a/cpp/velox/compute/VeloxRuntime.cc
+++ b/cpp/velox/compute/VeloxRuntime.cc
@@ -259,6 +259,7 @@ std::shared_ptr<ShuffleReader>
VeloxRuntime::createShuffleReader(
veloxCompressionType,
rowType,
options.batchSize,
+ options.bufferSize,
memoryManager()->getArrowMemoryPool(),
ctxVeloxPool,
options.shuffleWriterType);
diff --git a/cpp/velox/shuffle/VeloxShuffleReader.cc
b/cpp/velox/shuffle/VeloxShuffleReader.cc
index 92751d454a..0407be736a 100644
--- a/cpp/velox/shuffle/VeloxShuffleReader.cc
+++ b/cpp/velox/shuffle/VeloxShuffleReader.cc
@@ -38,7 +38,6 @@
#include <algorithm>
#include <iostream>
-// using namespace facebook;
using namespace facebook::velox;
namespace gluten {
@@ -291,14 +290,14 @@
VeloxHashShuffleReaderDeserializer::VeloxHashShuffleReaderDeserializer(
const std::shared_ptr<arrow::util::Codec>& codec,
const facebook::velox::RowTypePtr& rowType,
int32_t batchSize,
+ int64_t bufferSize,
arrow::MemoryPool* memoryPool,
facebook::velox::memory::MemoryPool* veloxPool,
std::vector<bool>* isValidityBuffer,
bool hasComplexType,
int64_t& deserializeTime,
int64_t& decompressTime)
- : in_(std::move(in)),
- schema_(schema),
+ : schema_(schema),
codec_(codec),
rowType_(rowType),
batchSize_(batchSize),
@@ -307,13 +306,16 @@
VeloxHashShuffleReaderDeserializer::VeloxHashShuffleReaderDeserializer(
isValidityBuffer_(isValidityBuffer),
hasComplexType_(hasComplexType),
deserializeTime_(deserializeTime),
- decompressTime_(decompressTime) {}
+ decompressTime_(decompressTime) {
+ GLUTEN_ASSIGN_OR_THROW(in_,
arrow::io::BufferedInputStream::Create(bufferSize, memoryPool, std::move(in)));
+}
std::shared_ptr<ColumnarBatch> VeloxHashShuffleReaderDeserializer::next() {
if (hasComplexType_) {
uint32_t numRows = 0;
GLUTEN_ASSIGN_OR_THROW(
- auto arrowBuffers, BlockPayload::deserialize(in_.get(), codec_,
memoryPool_, numRows, decompressTime_));
+ auto arrowBuffers,
+ BlockPayload::deserialize(in_.get(), codec_, memoryPool_, numRows,
deserializeTime_, decompressTime_));
if (arrowBuffers.empty()) {
// Reach EOS.
return nullptr;
@@ -332,7 +334,8 @@ std::shared_ptr<ColumnarBatch>
VeloxHashShuffleReaderDeserializer::next() {
uint32_t numRows = 0;
while (!merged_ || merged_->numRows() < batchSize_) {
GLUTEN_ASSIGN_OR_THROW(
- arrowBuffers, BlockPayload::deserialize(in_.get(), codec_,
memoryPool_, numRows, decompressTime_));
+ arrowBuffers,
+ BlockPayload::deserialize(in_.get(), codec_, memoryPool_, numRows,
deserializeTime_, decompressTime_));
if (arrowBuffers.empty()) {
reachEos_ = true;
break;
@@ -372,22 +375,24 @@
VeloxSortShuffleReaderDeserializer::VeloxSortShuffleReaderDeserializer(
const std::shared_ptr<arrow::util::Codec>& codec,
const RowTypePtr& rowType,
int32_t batchSize,
+ int64_t bufferSize,
arrow::MemoryPool* memoryPool,
facebook::velox::memory::MemoryPool* veloxPool,
int64_t& deserializeTime,
int64_t& decompressTime)
- : in_(std::move(in)),
- schema_(schema),
+ : schema_(schema),
codec_(codec),
rowType_(rowType),
batchSize_(batchSize),
arrowPool_(memoryPool),
veloxPool_(veloxPool),
deserializeTime_(deserializeTime),
- decompressTime_(decompressTime) {}
+ decompressTime_(decompressTime) {
+ GLUTEN_ASSIGN_OR_THROW(in_,
arrow::io::BufferedInputStream::Create(bufferSize, memoryPool, std::move(in)));
+}
std::shared_ptr<ColumnarBatch> VeloxSortShuffleReaderDeserializer::next() {
- if (reachEos_) {
+ if (reachedEos_) {
if (cachedRows_ > 0) {
return deserializeToBatch();
}
@@ -401,10 +406,11 @@ std::shared_ptr<ColumnarBatch>
VeloxSortShuffleReaderDeserializer::next() {
while (cachedRows_ < batchSize_) {
uint32_t numRows = 0;
GLUTEN_ASSIGN_OR_THROW(
- auto arrowBuffers, BlockPayload::deserialize(in_.get(), codec_,
arrowPool_, numRows, decompressTime_));
+ auto arrowBuffers,
+ BlockPayload::deserialize(in_.get(), codec_, arrowPool_, numRows,
deserializeTime_, decompressTime_));
if (arrowBuffers.empty()) {
- reachEos_ = true;
+ reachedEos_ = true;
if (cachedRows_ > 0) {
return deserializeToBatch();
}
@@ -467,7 +473,8 @@ void
VeloxSortShuffleReaderDeserializer::readLargeRow(std::vector<std::shared_pt
uint32_t numRows;
while (bufferSize < rowSize) {
GLUTEN_ASSIGN_OR_THROW(
- arrowBuffers, BlockPayload::deserialize(in_.get(), codec_, arrowPool_,
numRows, decompressTime_));
+ arrowBuffers,
+ BlockPayload::deserialize(in_.get(), codec_, arrowPool_, numRows,
deserializeTime_, decompressTime_));
VELOX_DCHECK_EQ(numRows, 0);
bufferSize += arrowBuffers[0]->size();
buffers.emplace_back(std::move(arrowBuffers[0]));
@@ -575,6 +582,7 @@
VeloxColumnarBatchDeserializerFactory::VeloxColumnarBatchDeserializerFactory(
const facebook::velox::common::CompressionKind veloxCompressionType,
const RowTypePtr& rowType,
int32_t batchSize,
+ int64_t bufferSize,
arrow::MemoryPool* memoryPool,
std::shared_ptr<facebook::velox::memory::MemoryPool> veloxPool,
ShuffleWriterType shuffleWriterType)
@@ -583,6 +591,7 @@
VeloxColumnarBatchDeserializerFactory::VeloxColumnarBatchDeserializerFactory(
veloxCompressionType_(veloxCompressionType),
rowType_(rowType),
batchSize_(batchSize),
+ bufferSize_(bufferSize),
memoryPool_(memoryPool),
veloxPool_(veloxPool),
shuffleWriterType_(shuffleWriterType) {
@@ -599,6 +608,7 @@ std::unique_ptr<ColumnarBatchIterator>
VeloxColumnarBatchDeserializerFactory::cr
codec_,
rowType_,
batchSize_,
+ bufferSize_,
memoryPool_,
veloxPool_.get(),
&isValidityBuffer_,
@@ -612,6 +622,7 @@ std::unique_ptr<ColumnarBatchIterator>
VeloxColumnarBatchDeserializerFactory::cr
codec_,
rowType_,
batchSize_,
+ bufferSize_,
memoryPool_,
veloxPool_.get(),
deserializeTime_,
diff --git a/cpp/velox/shuffle/VeloxShuffleReader.h
b/cpp/velox/shuffle/VeloxShuffleReader.h
index d39e389363..af35f97712 100644
--- a/cpp/velox/shuffle/VeloxShuffleReader.h
+++ b/cpp/velox/shuffle/VeloxShuffleReader.h
@@ -21,6 +21,7 @@
#include "shuffle/Payload.h"
#include "shuffle/ShuffleReader.h"
#include "shuffle/VeloxSortShuffleWriter.h"
+#include "utils/Timer.h"
#include "velox/type/Type.h"
#include "velox/vector/ComplexVector.h"
@@ -36,6 +37,7 @@ class VeloxHashShuffleReaderDeserializer final : public
ColumnarBatchIterator {
const std::shared_ptr<arrow::util::Codec>& codec,
const facebook::velox::RowTypePtr& rowType,
int32_t batchSize,
+ int64_t bufferSize,
arrow::MemoryPool* memoryPool,
facebook::velox::memory::MemoryPool* veloxPool,
std::vector<bool>* isValidityBuffer,
@@ -73,6 +75,7 @@ class VeloxSortShuffleReaderDeserializer final : public
ColumnarBatchIterator {
const std::shared_ptr<arrow::util::Codec>& codec,
const facebook::velox::RowTypePtr& rowType,
int32_t batchSize,
+ int64_t bufferSize,
arrow::MemoryPool* memoryPool,
facebook::velox::memory::MemoryPool* veloxPool,
int64_t& deserializeTime,
@@ -97,7 +100,7 @@ class VeloxSortShuffleReaderDeserializer final : public
ColumnarBatchIterator {
std::list<std::pair<uint32_t, facebook::velox::BufferPtr>> cachedInputs_;
uint32_t cachedRows_{0};
- bool reachEos_{false};
+ bool reachedEos_{false};
uint32_t rowOffset_{0};
size_t byteOffset_{0};
@@ -139,6 +142,7 @@ class VeloxColumnarBatchDeserializerFactory : public
DeserializerFactory {
const facebook::velox::common::CompressionKind veloxCompressionType,
const facebook::velox::RowTypePtr& rowType,
int32_t batchSize,
+ int64_t bufferSize,
arrow::MemoryPool* memoryPool,
std::shared_ptr<facebook::velox::memory::MemoryPool> veloxPool,
ShuffleWriterType shuffleWriterType);
@@ -161,6 +165,7 @@ class VeloxColumnarBatchDeserializerFactory : public
DeserializerFactory {
facebook::velox::common::CompressionKind veloxCompressionType_;
facebook::velox::RowTypePtr rowType_;
int32_t batchSize_;
+ int64_t bufferSize_;
arrow::MemoryPool* memoryPool_;
std::shared_ptr<facebook::velox::memory::MemoryPool> veloxPool_;
diff --git a/cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h
b/cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h
index 4b4abfb525..d2995e251c 100644
--- a/cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h
+++ b/cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h
@@ -366,6 +366,7 @@ class VeloxShuffleWriterTest : public
::testing::TestWithParam<ShuffleTestParams
veloxCompressionType,
rowType,
std::numeric_limits<int32_t>::max(),
+ kDefaultReadBufferSize,
defaultArrowMemoryPool().get(),
pool_,
GetParam().shuffleWriterType);
diff --git a/docs/Configuration.md b/docs/Configuration.md
index 76549dd4fe..d4ca9c10fb 100644
--- a/docs/Configuration.md
+++ b/docs/Configuration.md
@@ -51,6 +51,7 @@ You can add these configurations into spark-defaults.conf to
enable or disable t
| spark.gluten.sql.columnar.shuffle.compression.threshold | If number of
rows in a batch falls below this threshold, will copy all buffers into one
buffer to compress.
[...]
| spark.gluten.sql.columnar.shuffle.realloc.threshold | Set the
threshold to dynamically adjust the size of shuffle split buffers. The size of
each split buffer is recalculated for each incoming batch of data. If the new
size deviates from the current partition buffer size by a factor outside the
range of [1 - threshold, 1 + threshold], the split buffer will be re-allocated
using the newly calculated size
[...]
| spark.gluten.sql.columnar.shuffle.merge.threshold | Set the
threshold control the minimum merged size. When a partition buffer is full, and
the number of rows is below (`threshold *
spark.gluten.sql.columnar.maxBatchSize`), it will be saved for merging.
[...]
+| spark.gluten.sql.columnar.shuffle.readerBufferSize | Buffer size
in bytes for shuffle reader reading input stream from local or remote.
[...]
| spark.gluten.sql.columnar.numaBinding | Set up
NUMABinding, default is false
[...]
| spark.gluten.sql.columnar.coreRange | Set up the
core range for NUMABinding, only works when numaBinding set to true. <br /> The
setting is based on the number of cores in your system. Use 72 cores as an
example.
[...]
| spark.gluten.sql.columnar.wholeStage.fallback.threshold | Configure the
threshold for whether whole stage will fall back in AQE supported case by
counting the number of ColumnarToRow & vanilla leaf node
[...]
diff --git
a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleReaderJniWrapper.java
b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleReaderJniWrapper.java
index 09f97cbd06..8d031cf079 100644
---
a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleReaderJniWrapper.java
+++
b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleReaderJniWrapper.java
@@ -40,6 +40,7 @@ public class ShuffleReaderJniWrapper implements RuntimeAware {
String compressionType,
String compressionCodecBackend,
int batchSize,
+ long bufferSize,
String shuffleWriterType);
public native long readStream(long shuffleReaderHandle, JniByteInputStream
jniIn);
diff --git
a/gluten-celeborn/velox/src/main/scala/org/apache/spark/shuffle/VeloxCelebornColumnarBatchSerializer.scala
b/gluten-celeborn/velox/src/main/scala/org/apache/spark/shuffle/VeloxCelebornColumnarBatchSerializer.scala
index a1fb190e45..d5f20c8dea 100644
---
a/gluten-celeborn/velox/src/main/scala/org/apache/spark/shuffle/VeloxCelebornColumnarBatchSerializer.scala
+++
b/gluten-celeborn/velox/src/main/scala/org/apache/spark/shuffle/VeloxCelebornColumnarBatchSerializer.scala
@@ -88,12 +88,14 @@ private class CelebornColumnarBatchSerializerInstance(
.replace(GLUTEN_SORT_SHUFFLE_WRITER, GLUTEN_RSS_SORT_SHUFFLE_WRITER)
val jniWrapper = ShuffleReaderJniWrapper.create(runtime)
val batchSize = GlutenConfig.getConf.maxBatchSize
+ val bufferSize = GlutenConfig.getConf.columnarShuffleReaderBufferSize
val handle = jniWrapper
.make(
cSchema.memoryAddress(),
compressionCodec,
compressionCodecBackend,
batchSize,
+ bufferSize,
shuffleWriterType
)
// Close shuffle reader instance as lately as the end of task processing,
diff --git a/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala
b/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala
index 0e41fc6537..f049e045c7 100644
--- a/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala
+++ b/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala
@@ -193,6 +193,9 @@ class GlutenConfig(conf: SQLConf) extends Logging {
def columnarShuffleCompressionThreshold: Int =
conf.getConf(COLUMNAR_SHUFFLE_COMPRESSION_THRESHOLD)
+ def columnarShuffleReaderBufferSize: Long =
+ conf.getConf(COLUMNAR_SHUFFLE_READER_BUFFER_SIZE)
+
def maxBatchSize: Int = conf.getConf(COLUMNAR_MAX_BATCH_SIZE)
def columnarToRowMemThreshold: Long =
@@ -607,6 +610,9 @@ object GlutenConfig {
val GLUTEN_SHUFFLE_WRITER_MERGE_THRESHOLD =
"spark.gluten.sql.columnar.shuffle.merge.threshold"
val GLUTEN_SHUFFLE_DEFUALT_COMPRESSION_BUFFER_SIZE = 32 * 1024
+ // Shuffle reader buffer size.
+ val GLUTEN_SHUFFLE_READER_BUFFER_SIZE =
"spark.gluten.sql.columnar.shuffle.readerBufferSize"
+
// Controls whether to load DLL from jars. User can get dependent native
libs packed into a jar
// by executing dev/package.sh. Then, with that jar configured, Gluten can
load the native libs
// at runtime. This config is just for velox backend. And it is NOT
applicable to the situation
@@ -1155,6 +1161,13 @@ object GlutenConfig {
.checkValue(v => v >= 0 && v <= 1, "Shuffle writer merge threshold must
between [0, 1]")
.createWithDefault(0.25)
+ val COLUMNAR_SHUFFLE_READER_BUFFER_SIZE =
+ buildConf(GLUTEN_SHUFFLE_READER_BUFFER_SIZE)
+ .internal()
+ .doc("Buffer size in bytes for shuffle reader reading input stream from
local or remote.")
+ .bytesConf(ByteUnit.BYTE)
+ .createWithDefaultString("1MB")
+
val COLUMNAR_MAX_BATCH_SIZE =
buildConf(GLUTEN_MAX_BATCH_SIZE_KEY)
.internal()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
