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 6a64c4f706a..e95bf788e4d 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 @@ -137,20 +137,22 @@ 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) @@ -158,7 +160,7 @@ private class ColumnarBatchSerializerInstanceImpl( 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 diff --git a/backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializerInstance.scala b/backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializerInstance.scala index fea2ca70d59..583f3b4c25f 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializerInstance.scala +++ b/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 diff --git a/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleReader.scala b/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleReader.scala index c22daa09a79..f82a64ec32a 100644 --- a/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleReader.scala +++ b/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 diff --git a/backends-velox/src/main/scala/org/apache/spark/shuffle/VeloxShuffleUtils.scala b/backends-velox/src/main/scala/org/apache/spark/shuffle/VeloxShuffleUtils.scala index fe083f83f50..56a0083ce6b 100644 --- a/backends-velox/src/main/scala/org/apache/spark/shuffle/VeloxShuffleUtils.scala +++ b/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 ) diff --git a/backends-velox/src/test/scala/org/apache/spark/shuffle/ColumnarShuffleReaderSuite.scala b/backends-velox/src/test/scala/org/apache/spark/shuffle/ColumnarShuffleReaderSuite.scala index 7209dab1780..dd7ae9d45ec 100644 --- a/backends-velox/src/test/scala/org/apache/spark/shuffle/ColumnarShuffleReaderSuite.scala +++ b/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 } diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc index 91b35c822e6..d1489d987ff 100644 --- a/cpp/core/jni/JniWrapper.cc +++ b/cpp/core/jni/JniWrapper.cc @@ -1117,7 +1117,8 @@ 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(shuffleReaderHandle); @@ -1125,7 +1126,7 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleReaderJniWrappe ShuffleReader::OutputType requiredOutputType = ShuffleReader::getOutputType(executionMode); auto streamReader = std::make_shared(env, jStreamReader); - auto outItr = reader->read(streamReader, requiredOutputType); + auto outItr = reader->read(streamReader, requiredOutputType, readerOrder); return ctx->saveObject(outItr); JNI_METHOD_END(kInvalidObjectHandle) } diff --git a/cpp/core/shuffle/ShuffleReader.h b/cpp/core/shuffle/ShuffleReader.h index 7dd031522f5..4ac7c79b7f9 100644 --- a/cpp/core/shuffle/ShuffleReader.h +++ b/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 read( const std::shared_ptr& streamReader, - const OutputType& outputType) = 0; + const OutputType& outputType, + int32_t priority) = 0; virtual int64_t getDecompressTime() const = 0; diff --git a/cpp/velox/benchmarks/GenericBenchmark.cc b/cpp/velox/benchmarks/GenericBenchmark.cc index 087bcb8cf08..8dcde30bf3a 100644 --- a/cpp/velox/benchmarks/GenericBenchmark.cc +++ b/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(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 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(), diff --git a/cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.cc b/cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.cc index 1b34f3473f9..ef26011b268 100644 --- a/cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.cc +++ b/cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.cc @@ -52,6 +52,7 @@ arrow::Result readBlockType(arrow::io::InputStream* inputStream) { } // namespace VeloxGpuAsyncHashShuffleReaderDeserializer::VeloxGpuAsyncHashShuffleReaderDeserializer( + int32_t priority, const std::shared_ptr& streamReader, const std::shared_ptr& schema, const std::shared_ptr& 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), diff --git a/cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.h b/cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.h index 35fb80bc19b..ad4b163cbe5 100644 --- a/cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.h +++ b/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, const std::shared_ptr& schema, const std::shared_ptr& codec, @@ -57,6 +58,8 @@ class VeloxGpuAsyncHashShuffleReaderDeserializer final : public ShuffleReaderDes bool isStopped() const; + int32_t priority_{0}; + std::shared_ptr streamReader_; std::shared_ptr schema_; std::shared_ptr 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_; diff --git a/cpp/velox/shuffle/VeloxShuffleReader.cc b/cpp/velox/shuffle/VeloxShuffleReader.cc index d10a0ca3054..ae8fc3a81f3 100644 --- a/cpp/velox/shuffle/VeloxShuffleReader.cc +++ b/cpp/velox/shuffle/VeloxShuffleReader.cc @@ -965,7 +965,8 @@ VeloxShuffleReader::VeloxShuffleReader( void VeloxShuffleReader::createDeserializer( const std::shared_ptr& streamReader, - const OutputType& outputType) { + const OutputType& outputType, + int32_t priority) { switch (options_->shuffleWriterType) { case ShuffleWriterType::kHashShuffle: { if (outputType == OutputType::kCudfTable) { @@ -973,6 +974,7 @@ void VeloxShuffleReader::createDeserializer( VELOX_CHECK(!hasComplexType_); if (options_->enableGpuAsyncReader) { deserializer_ = std::make_unique( + priority, streamReader, schema_, codec_, @@ -1066,9 +1068,9 @@ void VeloxShuffleReader::initFromSchema() { std::shared_ptr VeloxShuffleReader::read( const std::shared_ptr& 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(deserializer_->deserializeStreams()); } diff --git a/cpp/velox/shuffle/VeloxShuffleReader.h b/cpp/velox/shuffle/VeloxShuffleReader.h index 5daa0c00671..f4d8918f8e5 100644 --- a/cpp/velox/shuffle/VeloxShuffleReader.h +++ b/cpp/velox/shuffle/VeloxShuffleReader.h @@ -202,8 +202,8 @@ class VeloxShuffleReader final : public ShuffleReader { VeloxMemoryManager* memoryManager, const std::shared_ptr& options); - std::shared_ptr read(const std::shared_ptr& streamReader, const OutputType& outputType) - override; + std::shared_ptr + read(const std::shared_ptr& 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, const OutputType& outputType); + void + createDeserializer(const std::shared_ptr& streamReader, const OutputType& outputType, int32_t priority); std::shared_ptr schema_; VeloxMemoryManager* memoryManager_; diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc index 46bafbfa42c..16fce21db7f 100644 --- a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc +++ b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc @@ -151,22 +151,24 @@ std::vector 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(schema, getDefaultMemoryManager(), options); const auto iter = - reader->read(std::make_shared(std::move(in)), ShuffleReader::OutputType::kCudfTable); + reader->read(std::make_shared(std::move(in)), ShuffleReader::OutputType::kCudfTable, 0); while (iter->hasNext()) { auto cb = std::dynamic_pointer_cast(iter->next()); @@ -495,12 +497,7 @@ TEST_P(GpuHashPartitioningShuffleWriterTest, hashPart1Vector) { makeFlatVector({232, 34567235, 1212, 4567}), makeFlatVector( 4, [](vector_size_t row) { return row % 2; }, nullEvery(5), DATE()), - makeFlatVector( - 4, - [](vector_size_t row) { - return Timestamp{row % 2, 0}; - }, - nullEvery(5))}; + makeFlatVector(4, [](vector_size_t row) { return Timestamp{row % 2, 0}; }, nullEvery(5))}; const auto vector = makeRowVector(data); diff --git a/cpp/velox/tests/VeloxShuffleWriterTest.cc b/cpp/velox/tests/VeloxShuffleWriterTest.cc index 668c00888d5..0d157b41a61 100644 --- a/cpp/velox/tests/VeloxShuffleWriterTest.cc +++ b/cpp/velox/tests/VeloxShuffleWriterTest.cc @@ -115,23 +115,25 @@ std::vector getTestParams() { for (const auto diskWriteBufferSize : {4, 56, 32 * 1024}) { for (const bool useRadixSort : {true, false}) { for (const auto deserializerBufferSize : {static_cast(1L), kDefaultDeserializerBufferSize}) { - params.push_back(ShuffleTestParams{ - .shuffleWriterType = ShuffleWriterType::kSortShuffle, - .partitionWriterType = partitionWriterType, - .compressionType = compression, - .diskWriteBufferSize = diskWriteBufferSize, - .useRadixSort = useRadixSort, - .deserializerBufferSize = deserializerBufferSize}); + params.push_back( + ShuffleTestParams{ + .shuffleWriterType = ShuffleWriterType::kSortShuffle, + .partitionWriterType = partitionWriterType, + .compressionType = compression, + .diskWriteBufferSize = diskWriteBufferSize, + .useRadixSort = useRadixSort, + .deserializerBufferSize = deserializerBufferSize}); } } } } // Rss sort-based shuffle. - params.push_back(ShuffleTestParams{ - .shuffleWriterType = ShuffleWriterType::kRssSortShuffle, - .partitionWriterType = PartitionWriterType::kRss, - .compressionType = compression}); + params.push_back( + ShuffleTestParams{ + .shuffleWriterType = ShuffleWriterType::kRssSortShuffle, + .partitionWriterType = PartitionWriterType::kRss, + .compressionType = compression}); // Hash-based shuffle. for (const auto compressionThreshold : compressionThresholds) { @@ -139,24 +141,26 @@ std::vector getTestParams() { for (const auto mergeBufferSize : mergeBufferSizes) { for (const bool enableDictionary : {true, false}) { for (const bool enableTypeAwareCompress : {true, false}) { - params.push_back(ShuffleTestParams{ - .shuffleWriterType = ShuffleWriterType::kHashShuffle, - .partitionWriterType = PartitionWriterType::kLocal, - .compressionType = compression, - .compressionThreshold = compressionThreshold, - .mergeBufferSize = mergeBufferSize, - .enableDictionary = enableDictionary, - .enableTypeAwareCompress = enableTypeAwareCompress}); + params.push_back( + ShuffleTestParams{ + .shuffleWriterType = ShuffleWriterType::kHashShuffle, + .partitionWriterType = PartitionWriterType::kLocal, + .compressionType = compression, + .compressionThreshold = compressionThreshold, + .mergeBufferSize = mergeBufferSize, + .enableDictionary = enableDictionary, + .enableTypeAwareCompress = enableTypeAwareCompress}); } } } // Rss. - params.push_back(ShuffleTestParams{ - .shuffleWriterType = ShuffleWriterType::kHashShuffle, - .partitionWriterType = PartitionWriterType::kRss, - .compressionType = compression, - .compressionThreshold = compressionThreshold}); + params.push_back( + ShuffleTestParams{ + .shuffleWriterType = ShuffleWriterType::kHashShuffle, + .partitionWriterType = PartitionWriterType::kRss, + .compressionType = compression, + .compressionThreshold = compressionThreshold}); } } @@ -330,7 +334,7 @@ class VeloxShuffleWriterTest : public ::testing::TestWithParam(schema, getDefaultMemoryManager(), options); const auto iter = - reader->read(std::make_shared(std::move(in)), ShuffleReader::OutputType::kRowVector); + reader->read(std::make_shared(std::move(in)), ShuffleReader::OutputType::kRowVector, 0); while (iter->hasNext()) { auto vector = std::dynamic_pointer_cast(iter->next())->getRowVector(); vectors.emplace_back(vector); @@ -545,7 +549,7 @@ class VeloxShuffleReaderStreamMergeTest : public ::testing::Test, public VeloxSh const auto reader = std::make_shared(schema, getDefaultMemoryManager(), options); const auto iter = - reader->read(std::make_shared(std::move(streams)), ShuffleReader::OutputType::kRowVector); + reader->read(std::make_shared(std::move(streams)), ShuffleReader::OutputType::kRowVector, 0); std::vector output; while (iter->hasNext()) { @@ -726,7 +730,7 @@ TEST_F(VeloxShuffleReaderStreamMergeTest, hashReaderDoesNotReuseDictionaryAcross const auto reader = std::make_shared(schema, getDefaultMemoryManager(), options); const auto iter = - reader->read(std::make_shared(std::move(streams)), ShuffleReader::OutputType::kRowVector); + reader->read(std::make_shared(std::move(streams)), ShuffleReader::OutputType::kRowVector, 0); ASSERT_TRUE(iter->hasNext()); facebook::velox::test::assertEqualVectors( @@ -801,12 +805,7 @@ TEST_P(HashPartitioningShuffleWriterTest, hashPart1Vector) { makeFlatVector({232, 34567235, 1212, 4567}, DECIMAL(20, 4)), makeFlatVector( 4, [](vector_size_t row) { return row % 2; }, nullEvery(5), DATE()), - makeFlatVector( - 4, - [](vector_size_t row) { - return Timestamp{row % 2, 0}; - }, - nullEvery(5))}; + makeFlatVector(4, [](vector_size_t row) { return Timestamp{row % 2, 0}; }, nullEvery(5))}; const auto vector = makeRowVector(data); 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 65a320f72c6..871fe8d028b 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 @@ -71,10 +71,13 @@ public native long make( long gpuAsyncReaderMaxPrefetchBytes); public native long read( - long shuffleReaderHandle, ShuffleStreamReader streamReader, int executionMode); + long shuffleReaderHandle, + ShuffleStreamReader streamReader, + int executionMode, + int readerOrder); public long read(long shuffleReaderHandle, ShuffleStreamReader streamReader) { - return read(shuffleReaderHandle, streamReader, CPUStageMode.id()); + return read(shuffleReaderHandle, streamReader, CPUStageMode.id(), 0); } public native void populateMetrics(long shuffleReaderHandle, ShuffleReaderMetrics metrics); diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala b/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala index d7a39992176..7eb307814e3 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala @@ -26,7 +26,7 @@ import org.apache.gluten.substrait.SubstraitContext import org.apache.gluten.substrait.plan.{PlanBuilder, PlanNode} import org.apache.gluten.substrait.rel.{LocalFilesNode, RelNode, SplitInfo} import org.apache.gluten.substrait.rel.LocalFilesNode.ReadFileFormat -import org.apache.gluten.utils.SubstraitPlanPrinterUtil +import org.apache.gluten.utils.{ShuffleReaderOrderUtil, SubstraitPlanPrinterUtil} import org.apache.spark._ import org.apache.spark.rdd.RDD @@ -398,6 +398,11 @@ case class WholeStageTransformer(child: SparkPlan, materializeInput: Boolean = f override def doExecuteColumnar(): RDD[ColumnarBatch] = { assert(child.isInstanceOf[TransformSupport]) val pipelineTime: SQLMetric = longMetric("pipelineTime") + + if (offloadCuda) { + ShuffleReaderOrderUtil.assign(child) + } + // We should do transform first to make sure all subqueries are materialized val wsCtx = GlutenTimeMetric.withMillisTime { doWholeStageTransform() diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scala b/gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scala new file mode 100644 index 00000000000..05878306caa --- /dev/null +++ b/gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scala @@ -0,0 +1,147 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten.utils + +import org.apache.gluten.execution.{BroadcastHashJoinExecTransformerBase, ShuffledHashJoinExecTransformerBase} + +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.adaptive.ColumnarAQEShuffleReadExec + +import scala.collection.mutable + +/** + * Computes the order in which Velox will first pull data from each shuffle reader of a whole stage. + */ +object ShuffleReaderOrderUtil { + + /** + * Assigns each [[ColumnarAQEShuffleReadExec]] in the stage rooted at `stageRoot` the order in + * which Velox will first pull data from it at runtime (0-based, via `setReaderOrder`). + * + * Why the order is not simply the plan-tree order: + * + * When Gluten hands a whole stage to Velox, Velox's LocalPlanner chops the operator tree into + * "pipelines" (linear chains of operators driven by one source). The cut points are hash-join + * build sides: the probe (streamed) side of a join stays in the current pipeline, while the build + * side becomes a NEW pipeline, appended to the pipeline list at the moment the planner's + * depth-first, probe-side-first walk reaches it. Every pipeline therefore has one source (a + * shuffle read, a scan, or a broadcast input) and may depend on the build pipelines feeding the + * hash joins it contains. + * + * Gluten then executes the Velox task single-threaded via Task::next(). The single thread + * repeatedly walks the pipeline list from index 0 upward — each full top-to-bottom walk is called + * a "sweep" below. On each sweep, a pipeline whose hash-join build inputs are not all finished is + * blocked (HashProbe reports kWaitForJoinBuild before consuming any input, so even the pipeline's + * source is not touched) and gets skipped; an unblocked pipeline runs, pulling its source — that + * is the moment its shuffle reader is first invoked. A pipeline finishing mid-sweep is visible to + * later-indexed pipelines within the same sweep, but earlier-indexed pipelines only notice on the + * next sweep. + * + * Net effect: all independent build-side shuffle readers fire first (in pipeline-list order), + * then intermediate probe readers as their builds complete, and the top-level probe reader fires + * last. This method reproduces the pipeline list and replays the sweeps to compute each reader's + * first-read order without running anything. + */ + def assign(stageRoot: SparkPlan): Unit = { + // One entry per Velox pipeline: + // - source: the shuffle reader at the bottom of the pipeline's operator chain, if any + // (pipelines fed by scans or broadcast inputs have none and get no order number); + // - builds: indices of the build pipelines that must finish before this pipeline may run + // (one per hash join contained in this pipeline); + // - done: completion flag used while replaying the sweeps. + class PipelineSim { + var source: Option[ColumnarAQEShuffleReadExec] = None + val builds = mutable.ArrayBuffer.empty[Int] + var done = false + } + + // Ordered as Velox's LocalPlanner orders its driver factories; the index in this buffer is + // the pipeline index the executor sweeps over. + val pipelines = mutable.ArrayBuffer.empty[PipelineSim] + + // Walk the plan tree and rebuild the pipeline list, mirroring LocalPlanner: `pipelineIdx` + // is the pipeline the current node belongs to; non-join nodes just stay in it, joins split. + def planPipelines(plan: SparkPlan, pipelineIdx: Int): Unit = { + // The probe side is walked first and stays in the current pipeline (so any joins nested + // inside it append THEIR build pipelines before this one); then the build side is placed + // in a fresh pipeline and recorded as a prerequisite of the current pipeline. + def splitBuildPipeline(probe: SparkPlan, build: SparkPlan): Unit = { + planPipelines(probe, pipelineIdx) + val buildIdx = pipelines.size + pipelines += new PipelineSim + planPipelines(build, buildIdx) + pipelines(pipelineIdx).builds += buildIdx + } + plan match { + // Broadcast builds contain no shuffle reader, but still get a pipeline: a probe blocked + // only on a broadcast build must still wait one sweep, which can delay its shuffle read + // past readers that fire on the first sweep. + case bhj: BroadcastHashJoinExecTransformerBase => + bhj.joinBuildSide match { + case BuildLeft => splitBuildPipeline(bhj.right, bhj.left) + case BuildRight => splitBuildPipeline(bhj.left, bhj.right) + } + + // For BuildLeft, Gluten swaps the children when lowering to Velox's HashJoinNode (whose + // build side is always the right source), so the left child is the build pipeline here. + case shj: ShuffledHashJoinExecTransformerBase => + shj.joinBuildSide match { + case BuildLeft => splitBuildPipeline(shj.right, shj.left) + case BuildRight => splitBuildPipeline(shj.left, shj.right) + } + + case c: ColumnarAQEShuffleReadExec => + pipelines(pipelineIdx).source = Some(c) + + // Any other operator (project, filter, aggregate, input iterator, ...) is a single-input + // link in the current pipeline's chain. NOTE: operators that Velox would also split into + // extra pipelines (union/local exchange, nested-loop join, ...) are not modeled; if one + // appears in a stage, the computed order may not match the runtime order. + case other => + other.children.foreach(planPipelines(_, pipelineIdx)) + } + } + + // Pipeline 0 is the output pipeline: the chain from the stage root down its probe sides. + pipelines += new PipelineSim + planPipelines(stageRoot, 0) + + // Replay the serial executor: each `while` iteration is one sweep over the pipeline list. + // A pipeline runs once all its build prerequisites are done; running it assigns the next + // order number to its shuffle reader (its source is pulled to exhaustion at that point). + // Marking `done` mid-sweep lets later-indexed pipelines run in the same sweep, matching the + // executor's forward-only scan. `progressed` guards against a malformed dependency graph. + var readerOrder = 0 + var progressed = true + while (progressed && pipelines.exists(!_.done)) { + progressed = false + pipelines.foreach { + p => + if (!p.done && p.builds.forall(pipelines(_).done)) { + p.source.foreach { + reader => + reader.setReaderOrder(readerOrder) + readerOrder += 1 + } + p.done = true + progressed = true + } + } + } + } +} diff --git a/gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleReaderWrapper.scala b/gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleReaderWrapper.scala index 9728adafec0..9001a13d87d 100644 --- a/gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleReaderWrapper.scala +++ b/gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleReaderWrapper.scala @@ -28,5 +28,6 @@ case class GenShuffleReaderParameters[K, C]( blocksByAddress: Iterator[(BlockManagerId, collection.Seq[(BlockId, Long, Int)])], context: TaskContext, readMetrics: ShuffleReadMetricsReporter, - shouldBatchFetch: Boolean = false, - executionMode: StageExecutionMode) + shouldBatchFetch: Boolean, + executionMode: StageExecutionMode, + readerOrder: Option[Int]) diff --git a/gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleUtils.scala b/gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleUtils.scala index 563109d22be..082aa9e2b77 100644 --- a/gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleUtils.scala +++ b/gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleUtils.scala @@ -159,7 +159,8 @@ object GlutenShuffleUtils { endPartition: Int, context: TaskContext, metrics: ShuffleReadMetricsReporter, - executionMode: StageExecutionMode): ShuffleReader[K, C] = { + executionMode: StageExecutionMode, + readerOrder: Option[Int]): ShuffleReader[K, C] = { val (blocksByAddress, canEnableBatchFetch) = { getReaderParam(handle, startMapIndex, endMapIndex, startPartition, endPartition) } @@ -174,7 +175,8 @@ object GlutenShuffleUtils { context, metrics, shouldBatchFetch, - executionMode)) + executionMode, + readerOrder)) .shuffleReader } } diff --git a/gluten-substrait/src/main/scala/org/apache/spark/shuffle/sort/ColumnarShuffleManager.scala b/gluten-substrait/src/main/scala/org/apache/spark/shuffle/sort/ColumnarShuffleManager.scala index 3b4126dc7cd..54a2d2610a0 100644 --- a/gluten-substrait/src/main/scala/org/apache/spark/shuffle/sort/ColumnarShuffleManager.scala +++ b/gluten-substrait/src/main/scala/org/apache/spark/shuffle/sort/ColumnarShuffleManager.scala @@ -141,7 +141,8 @@ class ColumnarShuffleManager(conf: SparkConf) endPartition, context, metrics, - CPUStageMode) + CPUStageMode, + None) } def getReader[K, C]( @@ -152,7 +153,8 @@ class ColumnarShuffleManager(conf: SparkConf) endPartition: Int, context: TaskContext, metrics: ShuffleReadMetricsReporter, - executionMode: StageExecutionMode): ShuffleReader[K, C] = { + executionMode: StageExecutionMode, + readerOrder: Option[Int]): ShuffleReader[K, C] = { GlutenShuffleUtils.genColumnarShuffleReader( handle, startMapIndex, @@ -161,7 +163,8 @@ class ColumnarShuffleManager(conf: SparkConf) endPartition, context, metrics, - executionMode) + executionMode, + readerOrder) } /** Remove a shuffle's metadata from the ShuffleManager. */ diff --git a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala index 0fb50694587..7f280be4f1e 100644 --- a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala +++ b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala @@ -145,19 +145,22 @@ case class ColumnarShuffleExchangeExec( columnarShuffleDependency, readMetrics, partitionSpecs, - CPUStageMode) + CPUStageMode, + None) } // Called by ColumnarAQEShuffleReaderExec to create a ShuffleRDD with custom partition specs, // and reducer stage execution mode. def getShuffleRDD( partitionSpecs: Array[ShufflePartitionSpec], - reducerStageMode: StageExecutionMode): RDD[ColumnarBatch] = { + reducerStageMode: StageExecutionMode, + readerOrder: Int): RDD[ColumnarBatch] = { new ShuffledColumnarBatchRDD( columnarShuffleDependency, readMetrics, partitionSpecs, - reducerStageMode) + reducerStageMode, + Some(readerOrder)) } // super.stringArgs ++ Iterator(output.map(o => s"${o}#${o.dataType.simpleString}")) diff --git a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala index ce01ce7b3e0..81f5851246e 100644 --- a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala +++ b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala @@ -50,7 +50,8 @@ class ShuffledColumnarBatchRDD( var dependency: ShuffleDependency[Int, ColumnarBatch, ColumnarBatch], metrics: Map[String, SQLMetric], partitionSpecs: Array[ShufflePartitionSpec], - executionMode: StageExecutionMode) + executionMode: StageExecutionMode, + readerOrder: Option[Int]) extends RDD[ColumnarBatch](dependency.rdd.context, Nil) { override val partitioner: Option[Partitioner] = @@ -79,7 +80,8 @@ class ShuffledColumnarBatchRDD( dependency, metrics, Array.tabulate(dependency.partitioner.numPartitions)(i => CoalescedPartitionSpec(i, i + 1)), - executionMode + executionMode, + None ) } @@ -124,7 +126,8 @@ class ShuffledColumnarBatchRDD( endReducerIndex, context, sqlMetricsReporter, - executionMode) + executionMode, + readerOrder) case PartialReducerPartitionSpec(reducerIndex, startMapIndex, endMapIndex, _) => getReader( @@ -136,7 +139,9 @@ class ShuffledColumnarBatchRDD( reducerIndex + 1, context, sqlMetricsReporter, - executionMode) + executionMode, + readerOrder + ) case PartialMapperPartitionSpec(mapIndex, startReducerIndex, endReducerIndex) => getReader( @@ -148,7 +153,9 @@ class ShuffledColumnarBatchRDD( endReducerIndex, context, sqlMetricsReporter, - executionMode) + executionMode, + readerOrder + ) case CoalescedMapperPartitionSpec(startMapIndex, endMapIndex, numReducers) => getReader( @@ -160,7 +167,8 @@ class ShuffledColumnarBatchRDD( numReducers, context, sqlMetricsReporter, - executionMode) + executionMode, + readerOrder) } new ShuffleReaderWithMetricsIterator( reader.read().asInstanceOf[Iterator[Product2[Int, ColumnarBatch]]], @@ -183,7 +191,8 @@ object ShuffledColumnarBatchRDD { endPartition: Int, context: TaskContext, metrics: ShuffleReadMetricsReporter, - executionMode: StageExecutionMode): ShuffleReader[K, C] = { + executionMode: StageExecutionMode, + readerOrder: Option[Int]): ShuffleReader[K, C] = { shuffleManager match { case columnarShuffleManager: ColumnarShuffleManager => columnarShuffleManager.getReader( @@ -194,7 +203,8 @@ object ShuffledColumnarBatchRDD { endPartition, context, metrics, - executionMode) + executionMode, + readerOrder) case _ => shuffleManager.getReader( handle, @@ -214,7 +224,8 @@ object ShuffledColumnarBatchRDD { endPartition: Int, context: TaskContext, metrics: ShuffleReadMetricsReporter, - executionMode: StageExecutionMode): ShuffleReader[K, C] = { + executionMode: StageExecutionMode, + readerOrder: Option[Int]): ShuffleReader[K, C] = { getReader[K, C]( shuffleManager, handle, @@ -224,6 +235,7 @@ object ShuffledColumnarBatchRDD { endPartition, context, metrics, - executionMode) + executionMode, + readerOrder) } } diff --git a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala index ba59c2c06ae..fffe6b19a93 100644 --- a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala +++ b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala @@ -57,7 +57,7 @@ case class ColumnarAQEShuffleReadExec( case a: AQEShuffleReadExec => a.stringArgs case _ => super.stringArgs } - } + } ++ Iterator(s"[order=$readerOrder]") override protected def withNewChildInternal(newChild: SparkPlan): ColumnarAQEShuffleReadExec = { delegate match { @@ -81,9 +81,14 @@ case class ColumnarAQEShuffleReadExec( s"Cannot get aqeReader from delegate node ${delegate.nodeName}.") } } - @transient override lazy val metrics: Map[String, SQLMetric] = aqeReader.metrics + private var readerOrder: Int = 0 + + def setReaderOrder(readerOrder: Int): Unit = { + this.readerOrder = readerOrder + } + private def shuffleStage = { val method = classOf[AQEShuffleReadExec].getDeclaredMethod("shuffleStage") method.setAccessible(true) @@ -105,7 +110,10 @@ case class ColumnarAQEShuffleReadExec( } stage.shuffle match { case columnarShuffle: ColumnarShuffleExchangeExec => - columnarShuffle.getShuffleRDD(aqeReader.partitionSpecs.toArray, executionMode) + columnarShuffle.getShuffleRDD( + aqeReader.partitionSpecs.toArray, + executionMode, + readerOrder) case _ => throw new IllegalStateException("shuffle stage is not a ColumnarShuffleExchangeExec") }