[GLUTEN-11524][VL] Set shuffle reader order for gpu async shuffle read

backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala

@@ -137,28 +137,30 @@ private class ColumnarBatchSerializerInstanceImpl(
 
   // `deserializeStream` is currently still used by uniffle shuffle reader.
   override def deserializeStream(in: InputStream): DeserializationStream = {
-    new TaskDeserializationStream(Iterator((null, in)), None, CPUStageMode)
+    new TaskDeserializationStream(Iterator((null, in)), None, CPUStageMode, None)
   }
 
   def deserializeStreams(
       streams: Iterator[(BlockId, InputStream)],
       onComplete: () => Unit,
-      executionMode: StageExecutionMode = CPUStageMode): DeserializationStream = {
-    new TaskDeserializationStream(streams, Some(onComplete), executionMode)
+      executionMode: StageExecutionMode = CPUStageMode,
+      readerOrder: Option[Int] = None): DeserializationStream = {
+    new TaskDeserializationStream(streams, Some(onComplete), executionMode, readerOrder)
   }
 
   private class TaskDeserializationStream(
       streams: Iterator[(BlockId, InputStream)],
       onComplete: Option[() => Unit],
-      executionMode: StageExecutionMode)
+      executionMode: StageExecutionMode,
+      var readerOrder: Option[Int])
     extends DeserializationStream
     with TaskResource {
     private val streamReader = ShuffleStreamReader(streams)
 
     private val wrappedOut: ClosableIterator[ColumnarBatch] = new ColumnarBatchOutIterator(
       runtime,
       jniWrapper
-        .read(shuffleReaderHandle, streamReader, executionMode.id))
+        .read(shuffleReaderHandle, streamReader, executionMode.id, readerOrder.getOrElse(0)))
 
     private var cb: ColumnarBatch = _
 
@@ -188,6 +190,10 @@ private class ColumnarBatchSerializerInstanceImpl(
 
     @throws(classOf[EOFException])
     override def readValue[T: ClassTag](): T = {
+      if (readerOrder.isDefined) {
+        logWarning(s"Start reading reader order: ${readerOrder.get}")
+        readerOrder = None
+      }
       if (cb != null) {
         cb.close()
         cb = null

backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializerInstance.scala

@@ -33,7 +33,8 @@ abstract class ColumnarBatchSerializerInstance extends SerializerInstance {
   def deserializeStreams(
       streams: Iterator[(BlockId, InputStream)],
       onComplete: () => Unit,
-      executionMode: StageExecutionMode = CPUStageMode): DeserializationStream
+      executionMode: StageExecutionMode = CPUStageMode,
+      readerOrder: Option[Int] = None): DeserializationStream
 
   override def serialize[T: ClassTag](t: T): ByteBuffer = {
     throw new UnsupportedOperationException

backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleReader.scala

@@ -36,6 +36,7 @@ class ColumnarShuffleReader[K, C](
     context: TaskContext,
     readMetrics: ShuffleReadMetricsReporter,
     executionMode: StageExecutionMode,
+    readerOrder: Option[Int],
     serializerManager: SerializerManager = SparkEnv.get.serializerManager,
     blockManager: BlockManager = SparkEnv.get.blockManager,
     mapOutputTracker: MapOutputTracker = SparkEnv.get.mapOutputTracker,
@@ -108,7 +109,8 @@ class ColumnarShuffleReader[K, C](
           .deserializeStreams(
             shuffleBlockFetcherIterator,
             shuffleBlockFetcherIterator.onComplete,
-            executionMode)
+            executionMode,
+            readerOrder)
           .asKeyValueIterator
       case serializerInstance =>
         // The dependency's serializer is not Gluten's ColumnarBatchSerializerInstance. This

backends-velox/src/main/scala/org/apache/spark/shuffle/VeloxShuffleUtils.scala

@@ -39,6 +39,7 @@ object VeloxShuffleUtils {
         parameters.context,
         parameters.readMetrics,
         parameters.executionMode,
+        parameters.readerOrder,
         serializerManager = ColumnarShuffleManager.bypassDecompressionSerializerManger,
         shouldBatchFetch = parameters.shouldBatchFetch
       )

backends-velox/src/test/scala/org/apache/spark/shuffle/ColumnarShuffleReaderSuite.scala

@@ -68,7 +68,8 @@ class ColumnarShuffleReaderSuite extends SharedSparkSession {
       Iterator.empty,
       TaskContext.empty(),
       new TempShuffleReadMetrics(),
-      CPUStageMode
+      CPUStageMode,
+      None
     )
     reader.read().toSeq
   }

cpp/core/jni/JniWrapper.cc

@@ -1117,15 +1117,16 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleReaderJniWrappe
     jobject wrapper,
     jlong shuffleReaderHandle,
     jobject jStreamReader,
-    jint executionMode) {
+    jint executionMode,
+    jint readerOrder) {
   JNI_METHOD_START
   auto ctx = getRuntime(env, wrapper);
   auto reader = ObjectStore::retrieve<ShuffleReader>(shuffleReaderHandle);
 
   ShuffleReader::OutputType requiredOutputType = ShuffleReader::getOutputType(executionMode);
   auto streamReader = std::make_shared<ShuffleStreamReader>(env, jStreamReader);
 
-  auto outItr = reader->read(streamReader, requiredOutputType);
+  auto outItr = reader->read(streamReader, requiredOutputType, readerOrder);
   return ctx->saveObject(outItr);
   JNI_METHOD_END(kInvalidObjectHandle)
 }

cpp/core/shuffle/ShuffleReader.h

@@ -37,7 +37,8 @@ class ShuffleReader {
   // FIXME iterator should be unique_ptr or un-copyable singleton
   virtual std::shared_ptr<ResultIterator> read(
       const std::shared_ptr<StreamReader>& streamReader,
-      const OutputType& outputType) = 0;
+      const OutputType& outputType,
+      int32_t priority) = 0;
 
   virtual int64_t getDecompressTime() const = 0;
 

cpp/velox/benchmarks/GenericBenchmark.cc

@@ -308,7 +308,7 @@ void runShuffle(
     GLUTEN_ASSIGN_OR_THROW(auto in, arrow::io::ReadableFile::Open(dataFile));
     auto streamReader = std::make_shared<TestStreamReader>(std::move(in));
     // Read all partitions.
-    auto iter = reader->read(streamReader, ShuffleReader::OutputType::kRowVector);
+    auto iter = reader->read(streamReader, ShuffleReader::OutputType::kRowVector, 0);
     while (iter->hasNext()) {
       // Read and discard.
       auto cb = iter->next();
@@ -435,8 +435,9 @@ auto BM_Generic = [](::benchmark::State& state,
       std::vector<FileReaderIterator*> inputItersRaw;
       if (!dataFiles.empty()) {
         for (const auto& input : dataFiles) {
-          inputIters.push_back(FileReaderIterator::getInputIteratorFromFileReader(
-              readerType, input, FLAGS_batch_size, runtime->memoryManager()->getLeafMemoryPool()));
+          inputIters.push_back(
+              FileReaderIterator::getInputIteratorFromFileReader(
+                  readerType, input, FLAGS_batch_size, runtime->memoryManager()->getLeafMemoryPool()));
         }
         std::transform(
             inputIters.begin(),

cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.cc

@@ -52,6 +52,7 @@ arrow::Result<BlockType> readBlockType(arrow::io::InputStream* inputStream) {
 } // namespace
 
 VeloxGpuAsyncHashShuffleReaderDeserializer::VeloxGpuAsyncHashShuffleReaderDeserializer(
+    int32_t priority,
     const std::shared_ptr<StreamReader>& streamReader,
     const std::shared_ptr<arrow::Schema>& schema,
     const std::shared_ptr<arrow::util::Codec>& codec,
@@ -61,7 +62,8 @@ VeloxGpuAsyncHashShuffleReaderDeserializer::VeloxGpuAsyncHashShuffleReaderDeseri
     VeloxMemoryManager* memoryManager,
     int64_t& deserializeTime,
     int64_t& decompressTime)
-    : streamReader_(streamReader),
+    : priority_(priority),
+      streamReader_(streamReader),
       schema_(schema),
       codec_(codec),
       rowType_(rowType),

cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.h

@@ -35,6 +35,7 @@ namespace gluten {
 class VeloxGpuAsyncHashShuffleReaderDeserializer final : public ShuffleReaderDeserializer {
  public:
   VeloxGpuAsyncHashShuffleReaderDeserializer(
+      int32_t priority,
       const std::shared_ptr<StreamReader>& streamReader,
       const std::shared_ptr<arrow::Schema>& schema,
       const std::shared_ptr<arrow::util::Codec>& codec,
@@ -57,6 +58,8 @@ class VeloxGpuAsyncHashShuffleReaderDeserializer final : public ShuffleReaderDes
 
   bool isStopped() const;
 
+  int32_t priority_{0};
+
   std::shared_ptr<StreamReader> streamReader_;
   std::shared_ptr<arrow::Schema> schema_;
   std::shared_ptr<arrow::util::Codec> codec_;
@@ -65,8 +68,6 @@ class VeloxGpuAsyncHashShuffleReaderDeserializer final : public ShuffleReaderDes
   int64_t maxPrefetchBytes_;
   VeloxMemoryManager* memoryManager_;
 
-  int32_t priority_{0};
-
   int64_t& deserializeTime_;
   int64_t& decompressTime_;
 

cpp/velox/shuffle/VeloxShuffleReader.cc

@@ -965,14 +965,16 @@ VeloxShuffleReader::VeloxShuffleReader(
 
 void VeloxShuffleReader::createDeserializer(
     const std::shared_ptr<StreamReader>& streamReader,
-    const OutputType& outputType) {
+    const OutputType& outputType,
+    int32_t priority) {
   switch (options_->shuffleWriterType) {
     case ShuffleWriterType::kHashShuffle: {
       if (outputType == OutputType::kCudfTable) {
 #ifdef GLUTEN_ENABLE_GPU
         VELOX_CHECK(!hasComplexType_);
         if (options_->enableGpuAsyncReader) {
           deserializer_ = std::make_unique<VeloxGpuAsyncHashShuffleReaderDeserializer>(
+              priority,
               streamReader,
               schema_,
               codec_,
@@ -1066,9 +1068,9 @@ void VeloxShuffleReader::initFromSchema() {
 
 std::shared_ptr<ResultIterator> VeloxShuffleReader::read(
     const std::shared_ptr<StreamReader>& streamReader,
-    const OutputType& outputType) {
-  // TODO: Support reader priority for async reader.
-  createDeserializer(streamReader, outputType);
+    const OutputType& outputType,
+    int32_t priority) {
+  createDeserializer(streamReader, outputType, priority);
   return std::make_shared<ResultIterator>(deserializer_->deserializeStreams());
 }
 

cpp/velox/shuffle/VeloxShuffleReader.h

@@ -202,8 +202,8 @@ class VeloxShuffleReader final : public ShuffleReader {
       VeloxMemoryManager* memoryManager,
       const std::shared_ptr<ShuffleReaderOptions>& options);
 
-  std::shared_ptr<ResultIterator> read(const std::shared_ptr<StreamReader>& streamReader, const OutputType& outputType)
-      override;
+  std::shared_ptr<ResultIterator>
+  read(const std::shared_ptr<StreamReader>& streamReader, const OutputType& outputType, int32_t priority) override;
 
   int64_t getDecompressTime() const override;
 
@@ -214,7 +214,8 @@ class VeloxShuffleReader final : public ShuffleReader {
  private:
   void initFromSchema();
 
-  void createDeserializer(const std::shared_ptr<StreamReader>& streamReader, const OutputType& outputType);
+  void
+  createDeserializer(const std::shared_ptr<StreamReader>& streamReader, const OutputType& outputType, int32_t priority);
 
   std::shared_ptr<arrow::Schema> schema_;
   VeloxMemoryManager* memoryManager_;

cpp/velox/tests/VeloxGpuShuffleWriterTest.cc

@@ -151,22 +151,24 @@ std::vector<GpuShuffleTestParams> getTestParams() {
       // Local.
       for (const auto mergeBufferSize : mergeBufferSizes) {
         for (const auto enableGpuAsyncReader : {false, true}) {
-          params.push_back(GpuShuffleTestParams{
-              .shuffleWriterType = ShuffleWriterType::kHashShuffle,
-              .partitionWriterType = PartitionWriterType::kLocal,
-              .compressionType = compression,
-              .compressionThreshold = compressionThreshold,
-              .mergeBufferSize = mergeBufferSize,
-              .enableGpuAsyncReader = enableGpuAsyncReader});
+          params.push_back(
+              GpuShuffleTestParams{
+                  .shuffleWriterType = ShuffleWriterType::kHashShuffle,
+                  .partitionWriterType = PartitionWriterType::kLocal,
+                  .compressionType = compression,
+                  .compressionThreshold = compressionThreshold,
+                  .mergeBufferSize = mergeBufferSize,
+                  .enableGpuAsyncReader = enableGpuAsyncReader});
         }
       }
 
       // Rss.
-      params.push_back(GpuShuffleTestParams{
-          .shuffleWriterType = ShuffleWriterType::kHashShuffle,
-          .partitionWriterType = PartitionWriterType::kRss,
-          .compressionType = compression,
-          .compressionThreshold = compressionThreshold});
+      params.push_back(
+          GpuShuffleTestParams{
+              .shuffleWriterType = ShuffleWriterType::kHashShuffle,
+              .partitionWriterType = PartitionWriterType::kRss,
+              .compressionType = compression,
+              .compressionThreshold = compressionThreshold});
     }
   }
 
@@ -314,7 +316,7 @@ class GpuVeloxShuffleWriterTest : public ::testing::TestWithParam<GpuShuffleTest
     const auto reader = std::make_shared<gluten::VeloxShuffleReader>(schema, getDefaultMemoryManager(), options);
 
     const auto iter =
-        reader->read(std::make_shared<TestStreamReader>(std::move(in)), ShuffleReader::OutputType::kCudfTable);
+        reader->read(std::make_shared<TestStreamReader>(std::move(in)), ShuffleReader::OutputType::kCudfTable, 0);
 
     while (iter->hasNext()) {
       auto cb = std::dynamic_pointer_cast<GpuBufferColumnarBatch>(iter->next());
@@ -495,12 +497,7 @@ TEST_P(GpuHashPartitioningShuffleWriterTest, hashPart1Vector) {
         makeFlatVector<int32_t>({232, 34567235, 1212, 4567}),
         makeFlatVector<int32_t>(
             4, [](vector_size_t row) { return row % 2; }, nullEvery(5), DATE()),
-        makeFlatVector<Timestamp>(
-            4,
-            [](vector_size_t row) {
-              return Timestamp{row % 2, 0};
-            },
-            nullEvery(5))};
+        makeFlatVector<Timestamp>(4, [](vector_size_t row) { return Timestamp{row % 2, 0}; }, nullEvery(5))};
 
     const auto vector = makeRowVector(data);