feat(shuffle): Add codec in bolt to support checksum during shuffle (#180)

Author: zhangxffff

bolt/shuffle/sparksql/BoltShuffleReader.cpp

@@ -41,6 +41,7 @@
 #include "bolt/shuffle/sparksql/BoltArrowMemoryPool.h"
 #include "bolt/shuffle/sparksql/Payload.h"
 #include "bolt/shuffle/sparksql/Utils.h"
+#include "bolt/shuffle/sparksql/compression/Codec.h"
 #include "bolt/shuffle/sparksql/compression/Compression.h"
 #include "bolt/vector/ComplexVector.h"
 #include "bolt/vector/FlatVector.h"
@@ -373,7 +374,7 @@ RowVectorPtr makeColumnarBatch(
 BoltColumnarBatchDeserializer::BoltColumnarBatchDeserializer(
     std::shared_ptr<arrow::io::InputStream> in,
     const std::shared_ptr<arrow::Schema>& schema,
-    const std::shared_ptr<arrow::util::Codec>& codec,
+    const std::shared_ptr<Codec>& codec,
     const bytedance::bolt::RowTypePtr& rowType,
     int32_t batchSize,
     int32_t shuffleBatchByteSize,
@@ -384,7 +385,7 @@ BoltColumnarBatchDeserializer::BoltColumnarBatchDeserializer(
     uint64_t& deserializeTime,
     uint64_t& decompressTime,
     bool isRowFormat,
-    ZstdStreamCodec* zstdCodec,
+    AdaptiveParallelZstdCodec* zstdCodec,
     RowBufferPool* rowBufferPool,
     ShuffleRowToColumnarConverter* row2ColConverter)
     : schema_(schema),
@@ -701,19 +702,21 @@ bool BoltColumnarBatchDeserializer::isCompositeRowVectorLayout(int64_t& bytes) {
 
 BoltColumnarBatchDeserializerFactory::BoltColumnarBatchDeserializerFactory(
     const std::shared_ptr<arrow::Schema>& schema,
-    const std::shared_ptr<arrow::util::Codec>& codec,
+    const std::shared_ptr<Codec>& codec,
     const RowTypePtr& rowType,
     int32_t batchSize,
     int32_t shuffleBatchByteSize,
     arrow::MemoryPool* memoryPool,
-    bytedance::bolt::memory::MemoryPool* boltPool)
+    bytedance::bolt::memory::MemoryPool* boltPool,
+    bool checksumEnabled)
     : schema_(schema),
       codec_(codec),
       rowType_(rowType),
       batchSize_(batchSize),
       shuffleBatchByteSize_(shuffleBatchByteSize),
       memoryPool_(memoryPool),
-      boltPool_(boltPool) {
+      boltPool_(boltPool),
+      checksumEnabled_(checksumEnabled) {
   initFromSchema();
 }
 
@@ -729,8 +732,8 @@ BoltColumnarBatchDeserializerFactory::createDeserializer(
         schema_->num_fields() >= rowBaseColumnNumThreshold) ||
        (shuffleWriterType_ == ShuffleWriterType::RowBased));
   if (!zstdCodec_) {
-    zstdCodec_ =
-        std::make_shared<ZstdStreamCodec>(1 /*not used*/, false, memoryPool_);
+    zstdCodec_ = std::make_shared<AdaptiveParallelZstdCodec>(
+        1 /*not used*/, false, memoryPool_, checksumEnabled_);
     rowBufferPool_ = std::make_shared<RowBufferPool>(memoryPool_);
     row2ColConverter_ =
         std::make_shared<ShuffleRowToColumnarConverter>(rowType_, boltPool_);
@@ -821,14 +824,18 @@ BoltShuffleReader::BoltShuffleReader(
     bytedance::bolt::memory::MemoryPool* boltPool)
     : factory_(std::make_unique<BoltColumnarBatchDeserializerFactory>(
           schema,
-          createArrowIpcCodec(
+          createCodec(
               options.compressionType,
-              getCodecBackend(options.codecBackend)),
+              CodecOptions{
+                  getCodecBackend(options.codecBackend),
+                  kDefaultCompressionLevel,
+                  options.checksumEnabled}),
           bytedance::bolt::asRowType(fromBoltTypeToArrowSchema(schema)),
           options.batchSize,
           options.shuffleBatchByteSize,
           pool,
-          boltPool)) {
+          boltPool,
+          options.checksumEnabled)) {
   factory_->setNumPartitions(options.numPartitions);
   factory_->setShuffleWriterType(options.forceShuffleWriterType);
   factory_->setpartitioningShortName(options.partitionShortName);

bolt/shuffle/sparksql/BoltShuffleReader.h

@@ -100,7 +100,7 @@ class BoltColumnarBatchDeserializer {
   BoltColumnarBatchDeserializer(
       std::shared_ptr<arrow::io::InputStream> in,
       const std::shared_ptr<arrow::Schema>& schema,
-      const std::shared_ptr<arrow::util::Codec>& codec,
+      const std::shared_ptr<Codec>& codec,
       const bytedance::bolt::RowTypePtr& rowType,
       int32_t batchSize,
       int32_t shuffleBatchByteSize,
@@ -111,7 +111,7 @@ class BoltColumnarBatchDeserializer {
       uint64_t& deserializeTime,
       uint64_t& decompressTime,
       bool isRowFormat = false,
-      ZstdStreamCodec* zstdCodec = nullptr,
+      AdaptiveParallelZstdCodec* zstdCodec = nullptr,
       RowBufferPool* rowBufferPool = nullptr,
       ShuffleRowToColumnarConverter* row2ColConverter = nullptr);
 
@@ -123,7 +123,7 @@ class BoltColumnarBatchDeserializer {
 
   std::shared_ptr<arrow::io::BufferedInputStream> in_;
   std::shared_ptr<arrow::Schema> schema_;
-  std::shared_ptr<arrow::util::Codec> codec_;
+  std::shared_ptr<Codec> codec_;
   bytedance::bolt::RowTypePtr rowType_;
   int32_t batchSize_;
   int32_t shuffleBatchByteSize_;
@@ -140,7 +140,7 @@ class BoltColumnarBatchDeserializer {
 
   // for row format shuffle read
   const bool isRowFormat_{false};
-  ZstdStreamCodec* zstdCodec_{nullptr};
+  AdaptiveParallelZstdCodec* zstdCodec_{nullptr};
   RowBufferPool* rowBufferPool_{nullptr};
   uint8_t* partialRow_{nullptr};
   int32_t partialRowSize_{0};
@@ -157,12 +157,13 @@ class BoltColumnarBatchDeserializerFactory {
  public:
   BoltColumnarBatchDeserializerFactory(
       const std::shared_ptr<arrow::Schema>& schema,
-      const std::shared_ptr<arrow::util::Codec>& codec,
+      const std::shared_ptr<Codec>& codec,
       const bytedance::bolt::RowTypePtr& rowType,
       int32_t batchSize,
       int32_t shuffleBatchByteSize,
       arrow::MemoryPool* memoryPool,
-      bytedance::bolt::memory::MemoryPool* boltPool);
+      bytedance::bolt::memory::MemoryPool* boltPool,
+      bool checksumEnabled = true);
 
   std::unique_ptr<BoltColumnarBatchDeserializer> createDeserializer(
       std::shared_ptr<arrow::io::InputStream> in);
@@ -205,7 +206,7 @@ class BoltColumnarBatchDeserializerFactory {
 
  private:
   std::shared_ptr<arrow::Schema> schema_;
-  std::shared_ptr<arrow::util::Codec> codec_;
+  std::shared_ptr<Codec> codec_;
   bytedance::bolt::RowTypePtr rowType_;
   int32_t batchSize_;
   int32_t shuffleBatchByteSize_;
@@ -223,9 +224,10 @@ class BoltColumnarBatchDeserializerFactory {
 
   void initFromSchema();
   // for rowbased shuffle
-  std::shared_ptr<ZstdStreamCodec> zstdCodec_{nullptr};
+  std::shared_ptr<AdaptiveParallelZstdCodec> zstdCodec_{nullptr};
   std::shared_ptr<RowBufferPool> rowBufferPool_{nullptr};
   std::shared_ptr<ShuffleRowToColumnarConverter> row2ColConverter_{nullptr};
+  bool checksumEnabled_{true};
 };
 
 class BoltShuffleReader {

bolt/shuffle/sparksql/CMakeLists.txt

@@ -34,6 +34,15 @@ bolt_add_library(
   BoltShuffleWriterV2.cpp
   CelebornReaderStreamIterator.cpp
   compression/Compression.cpp
+  compression/Codec.cpp
+  compression/GzipCodec.cpp
+  compression/Lz4Codec.cpp
+  compression/SnappyCodec.cpp
+  compression/ZstdCodec.cpp
+  compression/StreamCodec.cpp
+  compression/ZstdStreamCodec.cpp
+  compression/GzipStreamCodec.cpp
+  compression/Lz4FrameStreamCodec.cpp
   partition_writer/LocalPartitionWriter.cpp
   partition_writer/PartitionWriter.cpp
   partition_writer/rss/CelebornPartitionWriter.cpp
@@ -84,3 +93,7 @@ endif()
 if(${BOLT_BUILD_TESTING})
   add_subdirectory(tests)
 endif()
+
+if(${BOLT_ENABLE_BENCHMARKS})
+  add_subdirectory(benchmarks)
+endif()