[GLUTEN-8481][VL] Clean up shuffle reader cpp code

cpp/core/CMakeLists.txt

@@ -120,7 +120,6 @@ set(SPARK_COLUMNAR_PLUGIN_SRCS
     shuffle/RandomPartitioner.cc
     shuffle/RoundRobinPartitioner.cc
     shuffle/ShuffleMemoryPool.cc
-    shuffle/ShuffleReader.cc
     shuffle/ShuffleWriter.cc
     shuffle/SinglePartitioner.cc
     shuffle/Spill.cc

cpp/core/jni/JniWrapper.cc

@@ -1132,7 +1132,6 @@ JNIEXPORT void JNICALL Java_org_apache_gluten_vectorized_ShuffleReaderJniWrapper
     jlong shuffleReaderHandle) {
   JNI_METHOD_START
   auto reader = ObjectStore::retrieve<ShuffleReader>(shuffleReaderHandle);
-  GLUTEN_THROW_NOT_OK(reader->close());
   ObjectStore::release(shuffleReaderHandle);
   JNI_METHOD_END()
 }

cpp/core/shuffle/ShuffleReader.h

@@ -17,63 +17,22 @@
 
 #pragma once
 
-#include "memory/ColumnarBatch.h"
-
-#include <arrow/ipc/message.h>
-#include <arrow/ipc/options.h>
-
-#include "Options.h"
 #include "compute/ResultIterator.h"
-#include "utils/Compression.h"
 
 namespace gluten {
 
-class DeserializerFactory {
- public:
-  virtual ~DeserializerFactory() = default;
-
-  virtual std::unique_ptr<ColumnarBatchIterator> createDeserializer(std::shared_ptr<arrow::io::InputStream> in) = 0;
-
-  virtual arrow::MemoryPool* getPool() = 0;
-
-  virtual int64_t getDecompressTime() = 0;
-
-  virtual int64_t getDeserializeTime() = 0;
-
-  virtual ShuffleWriterType getShuffleWriterType() = 0;
-};
-
 class ShuffleReader {
  public:
-  explicit ShuffleReader(std::unique_ptr<DeserializerFactory> factory);
-
   virtual ~ShuffleReader() = default;
 
   // FIXME iterator should be unique_ptr or un-copyable singleton
-  virtual std::shared_ptr<ResultIterator> readStream(std::shared_ptr<arrow::io::InputStream> in);
-
-  arrow::Status close();
-
-  int64_t getDecompressTime() const;
-
-  int64_t getIpcTime() const;
-
-  int64_t getDeserializeTime() const;
-
-  arrow::MemoryPool* getPool() const;
-
-  ShuffleWriterType getShuffleWriterType() const;
+  virtual std::shared_ptr<ResultIterator> readStream(std::shared_ptr<arrow::io::InputStream> in) = 0;
 
- protected:
-  arrow::MemoryPool* pool_;
-  int64_t decompressTime_ = 0;
-  int64_t deserializeTime_ = 0;
+  virtual int64_t getDecompressTime() const = 0;
 
-  ShuffleWriterType shuffleWriterType_;
+  virtual int64_t getDeserializeTime() const = 0;
 
- private:
-  std::shared_ptr<arrow::Schema> schema_;
-  std::unique_ptr<DeserializerFactory> factory_;
+  virtual arrow::MemoryPool* getPool() const = 0;
 };
 
 } // namespace gluten

cpp/velox/compute/VeloxRuntime.cc

@@ -251,7 +251,7 @@ std::shared_ptr<ShuffleReader> VeloxRuntime::createShuffleReader(
   auto codec = gluten::createArrowIpcCodec(options.compressionType, options.codecBackend);
   auto ctxVeloxPool = memoryManager()->getLeafMemoryPool();
   auto veloxCompressionType = facebook::velox::common::stringToCompressionKind(options.compressionTypeStr);
-  auto deserializerFactory = std::make_unique<gluten::VeloxColumnarBatchDeserializerFactory>(
+  auto deserializerFactory = std::make_unique<gluten::VeloxShuffleReaderDeserializerFactory>(
       schema,
       std::move(codec),
       veloxCompressionType,

cpp/velox/shuffle/VeloxShuffleReader.cc

@@ -15,13 +15,13 @@
  * limitations under the License.
  */
 
-#include "VeloxShuffleReader.h"
-#include "GlutenByteStream.h"
+#include "shuffle/VeloxShuffleReader.h"
 
 #include <arrow/array/array_binary.h>
 #include <arrow/io/buffered.h>
 
 #include "memory/VeloxColumnarBatch.h"
+#include "shuffle/GlutenByteStream.h"
 #include "shuffle/Payload.h"
 #include "shuffle/Utils.h"
 #include "utils/Common.h"
@@ -576,7 +576,7 @@ std::shared_ptr<ColumnarBatch> VeloxRssSortShuffleReaderDeserializer::next() {
   return std::make_shared<VeloxColumnarBatch>(std::move(rowVector));
 }
 
-VeloxColumnarBatchDeserializerFactory::VeloxColumnarBatchDeserializerFactory(
+VeloxShuffleReaderDeserializerFactory::VeloxShuffleReaderDeserializerFactory(
     const std::shared_ptr<arrow::Schema>& schema,
     const std::shared_ptr<arrow::util::Codec>& codec,
     const facebook::velox::common::CompressionKind veloxCompressionType,
@@ -598,7 +598,7 @@ VeloxColumnarBatchDeserializerFactory::VeloxColumnarBatchDeserializerFactory(
   initFromSchema();
 }
 
-std::unique_ptr<ColumnarBatchIterator> VeloxColumnarBatchDeserializerFactory::createDeserializer(
+std::unique_ptr<ColumnarBatchIterator> VeloxShuffleReaderDeserializerFactory::createDeserializer(
     std::shared_ptr<arrow::io::InputStream> in) {
   switch (shuffleWriterType_) {
     case ShuffleWriterType::kHashShuffle:
@@ -635,23 +635,19 @@ std::unique_ptr<ColumnarBatchIterator> VeloxColumnarBatchDeserializerFactory::cr
   }
 }
 
-arrow::MemoryPool* VeloxColumnarBatchDeserializerFactory::getPool() {
+arrow::MemoryPool* VeloxShuffleReaderDeserializerFactory::getPool() {
   return memoryPool_;
 }
 
-ShuffleWriterType VeloxColumnarBatchDeserializerFactory::getShuffleWriterType() {
-  return shuffleWriterType_;
-}
-
-int64_t VeloxColumnarBatchDeserializerFactory::getDecompressTime() {
+int64_t VeloxShuffleReaderDeserializerFactory::getDecompressTime() {
   return decompressTime_;
 }
 
-int64_t VeloxColumnarBatchDeserializerFactory::getDeserializeTime() {
+int64_t VeloxShuffleReaderDeserializerFactory::getDeserializeTime() {
   return deserializeTime_;
 }
 
-void VeloxColumnarBatchDeserializerFactory::initFromSchema() {
+void VeloxShuffleReaderDeserializerFactory::initFromSchema() {
   GLUTEN_ASSIGN_OR_THROW(auto arrowColumnTypes, toShuffleTypeId(schema_->fields()));
   isValidityBuffer_.reserve(arrowColumnTypes.size());
   for (size_t i = 0; i < arrowColumnTypes.size(); ++i) {
@@ -681,7 +677,23 @@ void VeloxColumnarBatchDeserializerFactory::initFromSchema() {
   }
 }
 
-VeloxShuffleReader::VeloxShuffleReader(std::unique_ptr<DeserializerFactory> factory)
-    : ShuffleReader(std::move(factory)) {}
+VeloxShuffleReader::VeloxShuffleReader(std::unique_ptr<VeloxShuffleReaderDeserializerFactory> factory)
+    : factory_(std::move(factory)) {}
+
+std::shared_ptr<ResultIterator> VeloxShuffleReader::readStream(std::shared_ptr<arrow::io::InputStream> in) {
+  return std::make_shared<ResultIterator>(factory_->createDeserializer(in));
+}
+
+arrow::MemoryPool* VeloxShuffleReader::getPool() const {
+  return factory_->getPool();
+}
+
+int64_t VeloxShuffleReader::getDecompressTime() const {
+  return factory_->getDecompressTime();
+}
+
+int64_t VeloxShuffleReader::getDeserializeTime() const {
+  return factory_->getDeserializeTime();
+}
 
 } // namespace gluten

cpp/velox/shuffle/VeloxShuffleReader.h

@@ -17,16 +17,14 @@
 
 #pragma once
 
-#include "operators/serializer/VeloxColumnarBatchSerializer.h"
 #include "shuffle/Payload.h"
 #include "shuffle/ShuffleReader.h"
 #include "shuffle/VeloxSortShuffleWriter.h"
-#include "utils/Timer.h"
+
+#include "velox/serializers/PrestoSerializer.h"
 #include "velox/type/Type.h"
 #include "velox/vector/ComplexVector.h"
 
-#include <velox/serializers/PrestoSerializer.h>
-
 namespace gluten {
 
 class VeloxHashShuffleReaderDeserializer final : public ColumnarBatchIterator {
@@ -134,9 +132,9 @@ class VeloxRssSortShuffleReaderDeserializer : public ColumnarBatchIterator {
   std::shared_ptr<VeloxInputStream> in_;
 };
 
-class VeloxColumnarBatchDeserializerFactory : public DeserializerFactory {
+class VeloxShuffleReaderDeserializerFactory {
  public:
-  VeloxColumnarBatchDeserializerFactory(
+  VeloxShuffleReaderDeserializerFactory(
       const std::shared_ptr<arrow::Schema>& schema,
       const std::shared_ptr<arrow::util::Codec>& codec,
       const facebook::velox::common::CompressionKind veloxCompressionType,
@@ -147,15 +145,13 @@ class VeloxColumnarBatchDeserializerFactory : public DeserializerFactory {
       std::shared_ptr<facebook::velox::memory::MemoryPool> veloxPool,
       ShuffleWriterType shuffleWriterType);
 
-  std::unique_ptr<ColumnarBatchIterator> createDeserializer(std::shared_ptr<arrow::io::InputStream> in) override;
-
-  arrow::MemoryPool* getPool() override;
+  std::unique_ptr<ColumnarBatchIterator> createDeserializer(std::shared_ptr<arrow::io::InputStream> in);
 
-  int64_t getDecompressTime() override;
+  arrow::MemoryPool* getPool();
 
-  int64_t getDeserializeTime() override;
+  int64_t getDecompressTime();
 
-  ShuffleWriterType getShuffleWriterType() override;
+  int64_t getDeserializeTime();
 
  private:
   void initFromSchema();
@@ -180,6 +176,17 @@ class VeloxColumnarBatchDeserializerFactory : public DeserializerFactory {
 
 class VeloxShuffleReader final : public ShuffleReader {
  public:
-  VeloxShuffleReader(std::unique_ptr<DeserializerFactory> factory);
+  VeloxShuffleReader(std::unique_ptr<VeloxShuffleReaderDeserializerFactory> factory);
+
+  std::shared_ptr<ResultIterator> readStream(std::shared_ptr<arrow::io::InputStream> in) override;
+
+  int64_t getDecompressTime() const override;
+
+  int64_t getDeserializeTime() const override;
+
+  arrow::MemoryPool* getPool() const override;
+
+ private:
+  std::unique_ptr<VeloxShuffleReaderDeserializerFactory> factory_;
 };
 } // namespace gluten

cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h

@@ -360,7 +360,7 @@ class VeloxShuffleWriterTest : public ::testing::TestWithParam<ShuffleTestParams
       facebook::velox::serializer::presto::PrestoVectorSerde::registerVectorSerde();
     }
     // Set batchSize to a large value to make all batches are merged by reader.
-    auto deserializerFactory = std::make_unique<gluten::VeloxColumnarBatchDeserializerFactory>(
+    auto deserializerFactory = std::make_unique<gluten::VeloxShuffleReaderDeserializerFactory>(
         schema,
         std::move(codec),
         veloxCompressionType,