From 89c5c6a20d36363f95d9eae331b3abdb2cc34ea6 Mon Sep 17 00:00:00 2001 From: Chengcheng Jin Date: Fri, 21 Nov 2025 13:49:45 +0000 Subject: [PATCH 1/3] [GLUTEN-10933][VL] Support resize GPU batch --- .github/workflows/velox_backend_x86.yml | 2 +- .../GpuBufferBatchResizerJniWrapper.java | 40 + .../backendsapi/velox/VeloxMetricsApi.scala | 4 +- .../backendsapi/velox/VeloxRuleApi.scala | 2 + .../apache/gluten/config/VeloxConfig.scala | 8 + .../GpuResizeBufferColumnarBatchExec.scala | 68 ++ ...dBatchResizeForShuffleInputAndOutput.scala | 3 + .../extension/CudfNodeValidationRule.scala | 43 +- ...fferBatchResizeForShuffleInputOutput.scala | 80 ++ .../spark/shuffle/ColumnarShuffleWriter.scala | 15 +- cpp/core/jni/JniWrapper.cc | 28 + cpp/core/shuffle/Options.h | 32 + cpp/velox/CMakeLists.txt | 31 +- cpp/velox/compute/VeloxBackend.cc | 2 + cpp/velox/compute/VeloxPlanConverter.cc | 1 - cpp/velox/cudf/CudfPlanValidator.cc | 4 - cpp/velox/cudf/GpuLock.cc | 8 +- cpp/velox/jni/VeloxJniWrapper.cc | 19 + cpp/velox/memory/GpuBufferColumnarBatch.cc | 169 ++++ cpp/velox/memory/GpuBufferColumnarBatch.h | 73 ++ .../operators/plannodes/CudfVectorStream.h | 153 ++++ .../operators/plannodes/RowVectorStream.cc | 95 +++ .../operators/plannodes/RowVectorStream.h | 73 +- cpp/velox/shuffle/GpuShuffleReader.cc | 393 ---------- cpp/velox/shuffle/VeloxGpuShuffleReader.cc | 124 +++ ...huffleReader.h => VeloxGpuShuffleReader.h} | 15 +- cpp/velox/shuffle/VeloxGpuShuffleWriter.cc | 72 ++ cpp/velox/shuffle/VeloxGpuShuffleWriter.h | 57 ++ cpp/velox/shuffle/VeloxHashShuffleWriter.cc | 27 +- cpp/velox/shuffle/VeloxHashShuffleWriter.h | 29 +- cpp/velox/shuffle/VeloxShuffleReader.cc | 8 +- cpp/velox/shuffle/VeloxShuffleWriter.cc | 10 +- cpp/velox/substrait/SubstraitToVeloxPlan.cc | 31 +- cpp/velox/substrait/SubstraitToVeloxPlan.h | 1 + cpp/velox/tests/CMakeLists.txt | 78 ++ cpp/velox/tests/VeloxGpuShuffleWriterTest.cc | 741 ++++++++++++++++++ cpp/velox/utils/GpuBufferBatchResizer.cc | 214 +++++ cpp/velox/utils/GpuBufferBatchResizer.h | 44 ++ dev/docker/cudf/Dockerfile | 2 +- docs/velox-configuration.md | 1 + .../vectorized/ShuffleWriterJniWrapper.java | 8 + .../ColumnarShuffleExchangeExec.scala | 8 - .../ColumnarShuffleExchangeExecBase.scala | 12 +- .../GPUColumnarShuffleExchangeExec.scala | 14 +- 44 files changed, 2283 insertions(+), 559 deletions(-) create mode 100644 backends-velox/src/main/java/org/apache/gluten/utils/GpuBufferBatchResizerJniWrapper.java create mode 100644 backends-velox/src/main/scala/org/apache/gluten/execution/GpuResizeBufferColumnarBatchExec.scala create mode 100644 backends-velox/src/main/scala/org/apache/gluten/extension/GpuBufferBatchResizeForShuffleInputOutput.scala create mode 100644 cpp/velox/memory/GpuBufferColumnarBatch.cc create mode 100644 cpp/velox/memory/GpuBufferColumnarBatch.h create mode 100644 cpp/velox/operators/plannodes/CudfVectorStream.h create mode 100644 cpp/velox/operators/plannodes/RowVectorStream.cc delete mode 100644 cpp/velox/shuffle/GpuShuffleReader.cc create mode 100644 cpp/velox/shuffle/VeloxGpuShuffleReader.cc rename cpp/velox/shuffle/{GpuShuffleReader.h => VeloxGpuShuffleReader.h} (84%) create mode 100644 cpp/velox/shuffle/VeloxGpuShuffleWriter.cc create mode 100644 cpp/velox/shuffle/VeloxGpuShuffleWriter.h create mode 100644 cpp/velox/tests/VeloxGpuShuffleWriterTest.cc create mode 100644 cpp/velox/utils/GpuBufferBatchResizer.cc create mode 100644 cpp/velox/utils/GpuBufferBatchResizer.h diff --git a/.github/workflows/velox_backend_x86.yml b/.github/workflows/velox_backend_x86.yml index c200b79a622..bd640b6d1e6 100644 --- a/.github/workflows/velox_backend_x86.yml +++ b/.github/workflows/velox_backend_x86.yml @@ -1304,7 +1304,7 @@ jobs: dnf config-manager --add-repo "$repo_url" dnf install -y libnvjitlink-devel-12-8 df -a - bash dev/buildbundle-veloxbe.sh --run_setup_script=OFF --build_arrow=OFF --spark_version=3.4 --enable_gpu=ON + bash dev/buildbundle-veloxbe.sh --run_setup_script=OFF --build_arrow=OFF --spark_version=3.4 --build_tests=ON --build_benchmarks=ON --enable_gpu=ON ccache -s spark-test-spark40: diff --git a/backends-velox/src/main/java/org/apache/gluten/utils/GpuBufferBatchResizerJniWrapper.java b/backends-velox/src/main/java/org/apache/gluten/utils/GpuBufferBatchResizerJniWrapper.java new file mode 100644 index 00000000000..01fb7bf390a --- /dev/null +++ b/backends-velox/src/main/java/org/apache/gluten/utils/GpuBufferBatchResizerJniWrapper.java @@ -0,0 +1,40 @@ +/* + * 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.runtime.Runtime; +import org.apache.gluten.runtime.RuntimeAware; +import org.apache.gluten.vectorized.ColumnarBatchInIterator; + +public class GpuBufferBatchResizerJniWrapper implements RuntimeAware { + private final Runtime runtime; + + private GpuBufferBatchResizerJniWrapper(Runtime runtime) { + this.runtime = runtime; + } + + public static GpuBufferBatchResizerJniWrapper create(Runtime runtime) { + return new GpuBufferBatchResizerJniWrapper(runtime); + } + + @Override + public long rtHandle() { + return runtime.getHandle(); + } + + public native long create(int minOutputBatchSize, ColumnarBatchInIterator itr); +} diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala index 70597442daa..80afe9f19f7 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala @@ -17,7 +17,7 @@ package org.apache.gluten.backendsapi.velox import org.apache.gluten.backendsapi.MetricsApi -import org.apache.gluten.config.{HashShuffleWriterType, RssSortShuffleWriterType, ShuffleWriterType, SortShuffleWriterType} +import org.apache.gluten.config.{GpuHashShuffleWriterType, HashShuffleWriterType, RssSortShuffleWriterType, ShuffleWriterType, SortShuffleWriterType} import org.apache.gluten.metrics._ import org.apache.gluten.substrait.{AggregationParams, JoinParams} @@ -370,7 +370,7 @@ class VeloxMetricsApi extends MetricsApi with Logging { "peakBytes" -> SQLMetrics.createSizeMetric(sparkContext, "peak bytes allocated") ) shuffleWriterType match { - case HashShuffleWriterType => + case HashShuffleWriterType | GpuHashShuffleWriterType => baseMetrics ++ Map( "splitTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "time to split"), "avgDictionaryFields" -> SQLMetrics diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala index 485bf48f152..496f18811c1 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala @@ -109,6 +109,7 @@ object VeloxRuleApi { // Legacy: Post-transform rules. injector.injectPostTransform(_ => AppendBatchResizeForShuffleInputAndOutput()) + injector.injectPostTransform(_ => GpuBufferBatchResizeForShuffleInputOutput()) injector.injectPostTransform(_ => UnionTransformerRule()) injector.injectPostTransform(c => PartialProjectRule.apply(c.session)) injector.injectPostTransform(_ => PartialGenerateRule()) @@ -209,6 +210,7 @@ object VeloxRuleApi { // Gluten RAS: Post rules. injector.injectPostTransform(_ => DistinguishIdenticalScans) injector.injectPostTransform(_ => AppendBatchResizeForShuffleInputAndOutput()) + injector.injectPostTransform(_ => GpuBufferBatchResizeForShuffleInputOutput()) injector.injectPostTransform(_ => RemoveTransitions) injector.injectPostTransform(_ => UnionTransformerRule()) injector.injectPostTransform(c => PartialProjectRule.apply(c.session)) diff --git a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala index f68a23c3507..6ede352f53e 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala @@ -80,6 +80,8 @@ class VeloxConfig(conf: SQLConf) extends GlutenConfig(conf) { def cudfEnableValidation: Boolean = getConf(CUDF_ENABLE_VALIDATION) + def cudfBatchSize: Int = getConf(CUDF_BATCH_SIZE) + def orcUseColumnNames: Boolean = getConf(ORC_USE_COLUMN_NAMES) def parquetUseColumnNames: Boolean = getConf(PARQUET_USE_COLUMN_NAMES) @@ -634,6 +636,12 @@ object VeloxConfig extends ConfigRegistry { .booleanConf .createWithDefault(true) + val CUDF_BATCH_SIZE = + buildConf("spark.gluten.sql.columnar.backend.velox.cudf.batchSize") + .doc("Cudf input batch size after shuffle reader") + .intConf + .createWithDefault(Integer.MAX_VALUE) + val MEMORY_DUMP_ON_EXIT = buildConf("spark.gluten.monitor.memoryDumpOnExit") .internal() diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/GpuResizeBufferColumnarBatchExec.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/GpuResizeBufferColumnarBatchExec.scala new file mode 100644 index 00000000000..84768d71bc1 --- /dev/null +++ b/backends-velox/src/main/scala/org/apache/gluten/execution/GpuResizeBufferColumnarBatchExec.scala @@ -0,0 +1,68 @@ +/* + * 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.execution + +import org.apache.gluten.backendsapi.BackendsApiManager +import org.apache.gluten.backendsapi.velox.VeloxBatchType +import org.apache.gluten.extension.columnar.transition.Convention +import org.apache.gluten.iterator.ClosableIterator +import org.apache.gluten.runtime.Runtimes +import org.apache.gluten.utils.GpuBufferBatchResizerJniWrapper +import org.apache.gluten.vectorized.{ColumnarBatchInIterator, ColumnarBatchOutIterator} + +import org.apache.spark.sql.catalyst.expressions.SortOrder +import org.apache.spark.sql.catalyst.plans.physical.Partitioning +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.vectorized.ColumnarBatch + +import scala.collection.JavaConverters._ + +/** + * An operator to resize input BufferBatches generated by shuffle reader, and convert to cudf table. + */ +case class GpuResizeBufferColumnarBatchExec(override val child: SparkPlan, minOutputBatchSize: Int) + extends ColumnarToColumnarExec(child) { + + override protected def mapIterator(in: Iterator[ColumnarBatch]): Iterator[ColumnarBatch] = { + val runtime = + Runtimes.contextInstance(BackendsApiManager.getBackendName, "GpuBufferColumnarBatchResizer") + val outHandle = GpuBufferBatchResizerJniWrapper + .create(runtime) + .create( + minOutputBatchSize, + new ColumnarBatchInIterator(BackendsApiManager.getBackendName, in.asJava)) + new ColumnarBatchOutIterator(runtime, outHandle).asScala + } + + override protected def closeIterator(out: Iterator[ColumnarBatch]): Unit = { + out.asJava match { + case c: ClosableIterator[ColumnarBatch] => c.close() + case _ => + } + } + + override protected def needRecyclePayload: Boolean = true + + override def outputPartitioning: Partitioning = child.outputPartitioning + override def outputOrdering: Seq[SortOrder] = child.outputOrdering + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) + + override def batchType(): Convention.BatchType = VeloxBatchType + + override def rowType0(): Convention.RowType = Convention.RowType.None +} diff --git a/backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala b/backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala index 92519eecf77..01340c3a0be 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala @@ -30,6 +30,9 @@ import org.apache.spark.sql.execution.exchange.ReusedExchangeExec */ case class AppendBatchResizeForShuffleInputAndOutput() extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { + if (VeloxConfig.get.enableColumnarCudf) { + return plan + } val resizeBatchesShuffleInputEnabled = VeloxConfig.get.veloxResizeBatchesShuffleInput val resizeBatchesShuffleOutputEnabled = VeloxConfig.get.veloxResizeBatchesShuffleOutput if (!resizeBatchesShuffleInputEnabled && !resizeBatchesShuffleOutputEnabled) { diff --git a/backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala b/backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala index 97419ece7f2..29c85e0f5f9 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala @@ -18,8 +18,9 @@ package org.apache.gluten.extension import org.apache.gluten.config.{GlutenConfig, VeloxConfig} import org.apache.gluten.cudf.VeloxCudfPlanValidatorJniWrapper -import org.apache.gluten.execution.{CudfTag, LeafTransformSupport, TransformSupport, VeloxResizeBatchesExec, WholeStageTransformer} -import org.apache.gluten.extension.CudfNodeValidationRule.setTagForWholeStageTransformer +import org.apache.gluten.exception.GlutenNotSupportException +import org.apache.gluten.execution._ +import org.apache.gluten.extension.CudfNodeValidationRule.{createGPUColumnarExchange, setTagForWholeStageTransformer} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.execution.{ColumnarShuffleExchangeExec, GPUColumnarShuffleExchangeExec, SparkPlan} @@ -31,37 +32,23 @@ case class CudfNodeValidationRule(glutenConf: GlutenConfig) extends Rule[SparkPl if (!glutenConf.enableColumnarCudf) { return plan } - plan.transformUp { + val transformedPlan = plan.transformUp { case shuffle @ ColumnarShuffleExchangeExec( _, - v @ VeloxResizeBatchesExec(w: WholeStageTransformer, _, _), + VeloxResizeBatchesExec(w: WholeStageTransformer, _, _), _, _, _) => setTagForWholeStageTransformer(w) - if (w.isCudf) { - log.info("VeloxResizeBatchesExec is not supported in GPU") - } - GPUColumnarShuffleExchangeExec( - shuffle.outputPartitioning, - w, - shuffle.shuffleOrigin, - shuffle.projectOutputAttributes, - shuffle.advisoryPartitionSize) - + createGPUColumnarExchange(shuffle) case shuffle @ ColumnarShuffleExchangeExec(_, w: WholeStageTransformer, _, _, _) => setTagForWholeStageTransformer(w) - GPUColumnarShuffleExchangeExec( - shuffle.outputPartitioning, - w, - shuffle.shuffleOrigin, - shuffle.projectOutputAttributes, - shuffle.advisoryPartitionSize) - + createGPUColumnarExchange(shuffle) case transformer: WholeStageTransformer => setTagForWholeStageTransformer(transformer) transformer } + transformedPlan } } @@ -93,4 +80,18 @@ object CudfNodeValidationRule { transformer.setTagValue(CudfTag.CudfTag, true) } } + + def createGPUColumnarExchange(shuffle: ColumnarShuffleExchangeExec): SparkPlan = { + val exec = GPUColumnarShuffleExchangeExec( + shuffle.outputPartitioning, + shuffle.child, + shuffle.shuffleOrigin, + shuffle.projectOutputAttributes, + shuffle.advisoryPartitionSize) + val res = exec.doValidate() + if (!res.ok()) { + throw new GlutenNotSupportException(res.reason()) + } + exec + } } diff --git a/backends-velox/src/main/scala/org/apache/gluten/extension/GpuBufferBatchResizeForShuffleInputOutput.scala b/backends-velox/src/main/scala/org/apache/gluten/extension/GpuBufferBatchResizeForShuffleInputOutput.scala new file mode 100644 index 00000000000..d7a70e73295 --- /dev/null +++ b/backends-velox/src/main/scala/org/apache/gluten/extension/GpuBufferBatchResizeForShuffleInputOutput.scala @@ -0,0 +1,80 @@ +/* + * 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.extension + +import org.apache.gluten.config.{HashShuffleWriterType, VeloxConfig} +import org.apache.gluten.execution.{GpuResizeBufferColumnarBatchExec, VeloxResizeBatchesExec} + +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{ColumnarShuffleExchangeExec, ColumnarShuffleExchangeExecBase, SparkPlan} +import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.exchange.ReusedExchangeExec + +/** + * Try to append [[GpuBufferBatchResizeForShuffleInputOutput]] for shuffle input and output to make + * the batch sizes in good shape. + */ +case class GpuBufferBatchResizeForShuffleInputOutput() extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + if (!VeloxConfig.get.enableColumnarCudf) { + return plan + } + val range = VeloxConfig.get.veloxResizeBatchesShuffleInputOutputRange + val batchSize = VeloxConfig.get.cudfBatchSize + plan.transformUp { + case shuffle: ColumnarShuffleExchangeExec + if shuffle.shuffleWriterType == HashShuffleWriterType && + VeloxConfig.get.veloxResizeBatchesShuffleInput => + val appendBatches = + VeloxResizeBatchesExec(shuffle.child, range.min, range.max) + shuffle.withNewChildren(Seq(appendBatches)) + case a @ AQEShuffleReadExec( + ShuffleQueryStageExec(_, _: ColumnarShuffleExchangeExecBase, _), + _) => + GpuResizeBufferColumnarBatchExec(a, batchSize) + case a @ AQEShuffleReadExec( + ShuffleQueryStageExec(_, ReusedExchangeExec(_, _: ColumnarShuffleExchangeExecBase), _), + _) => + GpuResizeBufferColumnarBatchExec(a, batchSize) + // Since it's transformed in a bottom to up order, so we may first encounter + // ShuffeQueryStageExec, which is transformed to VeloxResizeBatchesExec(ShuffeQueryStageExec), + // then we see AQEShuffleReadExec + case a @ AQEShuffleReadExec( + GpuResizeBufferColumnarBatchExec( + s @ ShuffleQueryStageExec(_, _: ColumnarShuffleExchangeExecBase, _), + _), + _) => + GpuResizeBufferColumnarBatchExec(a.copy(child = s), batchSize) + case a @ AQEShuffleReadExec( + GpuResizeBufferColumnarBatchExec( + s @ ShuffleQueryStageExec( + _, + ReusedExchangeExec(_, _: ColumnarShuffleExchangeExecBase), + _), + _), + _) => + GpuResizeBufferColumnarBatchExec(a.copy(child = s), batchSize) + case s @ ShuffleQueryStageExec(_, _: ColumnarShuffleExchangeExecBase, _) => + GpuResizeBufferColumnarBatchExec(s, batchSize) + case s @ ShuffleQueryStageExec( + _, + ReusedExchangeExec(_, _: ColumnarShuffleExchangeExecBase), + _) => + GpuResizeBufferColumnarBatchExec(s, batchSize) + } + } +} diff --git a/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleWriter.scala b/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleWriter.scala index 94735e5b70a..f78a04e69fc 100644 --- a/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleWriter.scala +++ b/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleWriter.scala @@ -18,7 +18,7 @@ package org.apache.spark.shuffle import org.apache.gluten.backendsapi.BackendsApiManager import org.apache.gluten.columnarbatch.ColumnarBatches -import org.apache.gluten.config.{GlutenConfig, HashShuffleWriterType, SortShuffleWriterType} +import org.apache.gluten.config.{GlutenConfig, GpuHashShuffleWriterType, HashShuffleWriterType, SortShuffleWriterType} import org.apache.gluten.memory.memtarget.{MemoryTarget, Spiller} import org.apache.gluten.runtime.Runtimes import org.apache.gluten.vectorized._ @@ -44,7 +44,7 @@ class ColumnarShuffleWriter[K, V]( private val dep = handle.dependency.asInstanceOf[ColumnarShuffleDependency[K, V, V]] dep.shuffleWriterType match { - case HashShuffleWriterType | SortShuffleWriterType => + case HashShuffleWriterType | SortShuffleWriterType | GpuHashShuffleWriterType => // Valid shuffle writer types case _ => throw new IllegalArgumentException( @@ -171,6 +171,17 @@ class ColumnarShuffleWriter[K, V]( conf.get(SHUFFLE_SORT_USE_RADIXSORT), partitionWriterHandle ) + } else if (dep.shuffleWriterType == GpuHashShuffleWriterType) { + shuffleWriterJniWrapper.createGpuHashShuffleWriter( + numPartitions, + dep.nativePartitioning.getShortName, + GlutenShuffleUtils.getStartPartitionId( + dep.nativePartitioning, + taskContext.partitionId), + nativeBufferSize, + reallocThreshold, + partitionWriterHandle + ) } else { shuffleWriterJniWrapper.createHashShuffleWriter( numPartitions, diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc index f3b86f8ff6a..4129e53ac7b 100644 --- a/cpp/core/jni/JniWrapper.cc +++ b/cpp/core/jni/JniWrapper.cc @@ -916,6 +916,34 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrappe JNI_METHOD_END(kInvalidObjectHandle) } +JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrapper_createGpuHashShuffleWriter( + JNIEnv* env, + jobject wrapper, + jint numPartitions, + jstring partitioningNameJstr, + jint startPartitionId, + jint splitBufferSize, + jdouble splitBufferReallocThreshold, + jlong partitionWriterHandle) { + JNI_METHOD_START + + const auto ctx = getRuntime(env, wrapper); + + auto partitionWriter = ObjectStore::retrieve(partitionWriterHandle); + if (partitionWriter == nullptr) { + throw GlutenException("Partition writer handle is invalid: " + std::to_string(partitionWriterHandle)); + } + ObjectStore::release(partitionWriterHandle); + + auto shuffleWriterOptions = std::make_shared( + toPartitioning(jStringToCString(env, partitioningNameJstr)), + startPartitionId, + splitBufferSize, + splitBufferReallocThreshold); + return ctx->saveObject(ctx->createShuffleWriter(numPartitions, partitionWriter, shuffleWriterOptions)); + JNI_METHOD_END(kInvalidObjectHandle) +} + JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrapper_createSortShuffleWriter( JNIEnv* env, jobject wrapper, diff --git a/cpp/core/shuffle/Options.h b/cpp/core/shuffle/Options.h index aaf653f59ce..2139c6e9d72 100644 --- a/cpp/core/shuffle/Options.h +++ b/cpp/core/shuffle/Options.h @@ -91,6 +91,19 @@ struct HashShuffleWriterOptions : ShuffleWriterOptions { : ShuffleWriterOptions(ShuffleWriterType::kHashShuffle, partitioning, startPartitionId), splitBufferSize(partitionBufferSize), splitBufferReallocThreshold(partitionBufferReallocThreshold) {} + + protected: + HashShuffleWriterOptions(ShuffleWriterType shuffleWriterType) : ShuffleWriterOptions(shuffleWriterType) {} + + HashShuffleWriterOptions( + ShuffleWriterType shuffleWriterType, + Partitioning partitioning, + int32_t startPartitionId, + int32_t partitionBufferSize, + double partitionBufferReallocThreshold) + : ShuffleWriterOptions(shuffleWriterType, partitioning, startPartitionId), + splitBufferSize(partitionBufferSize), + splitBufferReallocThreshold(partitionBufferReallocThreshold) {} }; struct SortShuffleWriterOptions : ShuffleWriterOptions { @@ -131,6 +144,25 @@ struct RssSortShuffleWriterOptions : ShuffleWriterOptions { compressionType(compressionType) {} }; +struct GpuHashShuffleWriterOptions : HashShuffleWriterOptions { + int32_t splitBufferSize = kDefaultShuffleWriterBufferSize; + double splitBufferReallocThreshold = kDefaultSplitBufferReallocThreshold; + + GpuHashShuffleWriterOptions() : HashShuffleWriterOptions(ShuffleWriterType::kGpuHashShuffle) {} + + GpuHashShuffleWriterOptions( + Partitioning partitioning, + int32_t startPartitionId, + int32_t partitionBufferSize, + double partitionBufferReallocThreshold) + : HashShuffleWriterOptions( + ShuffleWriterType::kGpuHashShuffle, + partitioning, + startPartitionId, + partitionBufferSize, + partitionBufferReallocThreshold) {} +}; + struct LocalPartitionWriterOptions { int64_t shuffleFileBufferSize = kDefaultShuffleFileBufferSize; // spark.shuffle.file.buffer int32_t compressionBufferSize = diff --git a/cpp/velox/CMakeLists.txt b/cpp/velox/CMakeLists.txt index 3b9c8261c8e..e389113c4e0 100644 --- a/cpp/velox/CMakeLists.txt +++ b/cpp/velox/CMakeLists.txt @@ -163,6 +163,7 @@ set(VELOX_SRCS operators/functions/RegistrationAllFunctions.cc operators/functions/RowConstructorWithNull.cc operators/functions/SparkExprToSubfieldFilterParser.cc + operators/plannodes/RowVectorStream.cc operators/reader/FileReaderIterator.cc operators/reader/ParquetReaderIterator.cc operators/serializer/VeloxColumnarBatchSerializer.cc @@ -199,8 +200,15 @@ if(ENABLE_S3) endif() if(ENABLE_GPU) - list(APPEND VELOX_SRCS cudf/CudfPlanValidator.cc cudf/GpuLock.cc - shuffle/GpuShuffleReader.cc) + list( + APPEND + VELOX_SRCS + cudf/CudfPlanValidator.cc + cudf/GpuLock.cc + shuffle/VeloxGpuShuffleReader.cc + shuffle/VeloxGpuShuffleWriter.cc + utils/GpuBufferBatchResizer.cc + memory/GpuBufferColumnarBatch.cc) endif() if(ENABLE_ENHANCED_FEATURES) @@ -409,6 +417,10 @@ if(BUILD_EXAMPLES) endif() if(ENABLE_GPU) + import_library( + facebook::velox::velox_cudf_expression + ${VELOX_BUILD_PATH}/velox/experimental/cudf/expression/libvelox_cudf_expression.a + ) import_library( facebook::velox::velox_cudf_exec ${VELOX_BUILD_PATH}/velox/experimental/cudf/exec/libvelox_cudf_exec.a) @@ -435,9 +447,18 @@ if(ENABLE_GPU) target_link_libraries( velox PUBLIC facebook::velox::velox_cudf_exec facebook::velox::velox_cudf_vector - facebook::velox::velox_cudf_hive_connector) - target_link_libraries(velox - PRIVATE ${VELOX_BUILD_PATH}/_deps/cudf-build/libcudf.so) + facebook::velox::velox_cudf_hive_connector + facebook::velox::velox_cudf_expression) + target_link_libraries( + velox + PRIVATE + ${VELOX_BUILD_PATH}/_deps/cudf-build/libcudf.so + ${VELOX_BUILD_PATH}/_deps/rmm-build/librmm.so + ${VELOX_BUILD_PATH}/_deps/kvikio-build/libkvikio.so + ${VELOX_BUILD_PATH}/_deps/nvcomp_proprietary_binary-src/lib64/libnvcomp.so + ${VELOX_BUILD_PATH}/_deps/nvcomp_proprietary_binary-src/lib64/libnvcomp_cpu.so + ${VELOX_BUILD_PATH}/_deps/rapids_logger-build/librapids_logger.so + /usr/local/cuda-12.8/lib64/libcudart.so.12) endif() add_custom_command( diff --git a/cpp/velox/compute/VeloxBackend.cc b/cpp/velox/compute/VeloxBackend.cc index 8e54762fc01..a7ab9cef7c3 100644 --- a/cpp/velox/compute/VeloxBackend.cc +++ b/cpp/velox/compute/VeloxBackend.cc @@ -31,6 +31,7 @@ #include "velox/experimental/cudf/CudfConfig.h" #include "velox/experimental/cudf/connectors/hive/CudfHiveConnector.h" #include "velox/experimental/cudf/exec/ToCudf.h" +#include "operators/plannodes/CudfVectorStream.h" #endif #include "compute/VeloxRuntime.h" @@ -177,6 +178,7 @@ void VeloxBackend::init( auto& cudfConfig = velox::cudf_velox::CudfConfig::getInstance(); cudfConfig.initialize(std::move(options)); velox::cudf_velox::registerCudf(); + velox::exec::Operator::registerOperator(std::make_unique()); } #endif diff --git a/cpp/velox/compute/VeloxPlanConverter.cc b/cpp/velox/compute/VeloxPlanConverter.cc index 88eceb3a74e..05e78eb1bad 100644 --- a/cpp/velox/compute/VeloxPlanConverter.cc +++ b/cpp/velox/compute/VeloxPlanConverter.cc @@ -21,7 +21,6 @@ #include "compute/ResultIterator.h" #include "config/GlutenConfig.h" #include "iceberg/IcebergPlanConverter.h" -#include "operators/plannodes/RowVectorStream.h" #include "velox/common/file/FileSystems.h" namespace gluten { diff --git a/cpp/velox/cudf/CudfPlanValidator.cc b/cpp/velox/cudf/CudfPlanValidator.cc index 49949ca1007..a7efb3108e4 100644 --- a/cpp/velox/cudf/CudfPlanValidator.cc +++ b/cpp/velox/cudf/CudfPlanValidator.cc @@ -20,7 +20,6 @@ #include "compute/ResultIterator.h" #include "compute/VeloxBackend.h" #include "compute/VeloxPlanConverter.h" -#include "operators/plannodes/RowVectorStream.h" #include "velox/core/PlanNode.h" #include "velox/exec/Task.h" #include "velox/exec/TableScan.h" @@ -77,9 +76,6 @@ bool CudfPlanValidator::validate(const ::substrait::Plan& substraitPlan) { if (isCudfOperator(op)) { continue; } - if (dynamic_cast(op) != nullptr) { - continue; - } LOG(INFO) << "Operator " << op->operatorType() << " is not supported in cudf"; task->requestCancel().wait(); return false; diff --git a/cpp/velox/cudf/GpuLock.cc b/cpp/velox/cudf/GpuLock.cc index 47d383d9400..f28bbbe8a51 100644 --- a/cpp/velox/cudf/GpuLock.cc +++ b/cpp/velox/cudf/GpuLock.cc @@ -20,6 +20,7 @@ #include #include #include +#include namespace gluten { @@ -57,8 +58,13 @@ void lockGpu() { void unlockGpu() { std::thread::id tid = std::this_thread::get_id(); std::unique_lock lock(getGpuLockState().gGpuMutex); + if (!getGpuLockState().gGpuOwner.has_value()) { + LOG(INFO) <<"unlockGpu() called by non-owner thread!"<< std::endl; + return; + } + if (!getGpuLockState().gGpuOwner.has_value() || getGpuLockState().gGpuOwner != tid) { - throw std::runtime_error("unlockGpu() called by non-owner thread!"); + throw std::runtime_error("unlockGpu() called by other-owner thread!"); } // Release ownership diff --git a/cpp/velox/jni/VeloxJniWrapper.cc b/cpp/velox/jni/VeloxJniWrapper.cc index c18832f9be8..3993eeecdd4 100644 --- a/cpp/velox/jni/VeloxJniWrapper.cc +++ b/cpp/velox/jni/VeloxJniWrapper.cc @@ -41,6 +41,7 @@ #ifdef GLUTEN_ENABLE_GPU #include "cudf/CudfPlanValidator.h" +#include "utils/GpuBufferBatchResizer.h" #endif #ifdef GLUTEN_ENABLE_ENHANCED_FEATURES @@ -441,6 +442,24 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_utils_VeloxBatchResizerJniWrapper JNI_METHOD_END(kInvalidObjectHandle) } +#ifdef GLUTEN_ENABLE_GPU +JNIEXPORT jlong JNICALL Java_org_apache_gluten_utils_GpuBufferBatchResizerJniWrapper_create( // NOLINT + JNIEnv* env, + jobject wrapper, + jint minOutputBatchSize, + jobject jIter) { + JNI_METHOD_START + auto ctx = getRuntime(env, wrapper); + auto arrowPool = dynamic_cast(ctx->memoryManager())->defaultArrowMemoryPool(); + auto pool = dynamic_cast(ctx->memoryManager())->getLeafMemoryPool(); + auto iter = makeJniColumnarBatchIterator(env, jIter, ctx); + auto appender = std::make_shared( + std::make_unique(arrowPool, pool.get(), minOutputBatchSize, std::move(iter))); + return ctx->saveObject(appender); + JNI_METHOD_END(kInvalidObjectHandle) +} +#endif + JNIEXPORT jboolean JNICALL Java_org_apache_gluten_utils_VeloxFileSystemValidationJniWrapper_allSupportedByRegisteredFileSystems( // NOLINT JNIEnv* env, diff --git a/cpp/velox/memory/GpuBufferColumnarBatch.cc b/cpp/velox/memory/GpuBufferColumnarBatch.cc new file mode 100644 index 00000000000..649aadfb796 --- /dev/null +++ b/cpp/velox/memory/GpuBufferColumnarBatch.cc @@ -0,0 +1,169 @@ +/* + * 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. + */ +#include "GpuBufferColumnarBatch.h" +#include "compute/VeloxRuntime.h" +#include "utils/Timer.h" +#include "utils/VeloxArrowUtils.h" +#include "velox/row/UnsafeRowFast.h" +#include "velox/type/Type.h" +#include "velox/vector/FlatVector.h" + +#include +#include +#include + +namespace gluten { + namespace { + enum class BufferType{ + kNull, + kLength, + kValue + }; + } + +using namespace facebook; +using namespace facebook::velox; + +std::shared_ptr GpuBufferColumnarBatch::exportArrowSchema() { + throw GlutenException("GpuBufferColumnarBatch does not support exportArrowSchema"); +} + +std::shared_ptr GpuBufferColumnarBatch::exportArrowArray() { + throw GlutenException("GpuBufferColumnarBatch does not support exportArrowArray"); +} + +std::vector GpuBufferColumnarBatch::toUnsafeRow(int32_t rowId) const { + throw GlutenException("GpuBufferColumnarBatch does not support toUnsafeRow"); +} + +int64_t GpuBufferColumnarBatch::numBytes() { + int64_t numBytes = 0; + for (const auto& buffer : buffers_) { + numBytes += buffer->size(); + } + return numBytes; +} + +// Optimize to release the previous buffer after merge it. +std::shared_ptr GpuBufferColumnarBatch::compose( + arrow::MemoryPool* pool, + const std::vector>& batches, + int32_t numRows) { + GLUTEN_CHECK(!batches.empty(), "No batches to compose"); + // Compute the returned GpuBufferColumnarBatch buffers. + auto& type = batches[0]->getRowType(); + const auto bufferSize = batches[0]->buffers().size(); + std::vector bufferSizes; + bufferSizes.resize(bufferSize); + std::vector bufferTypes; + bufferTypes.reserve(bufferSize); + + for (const auto& colType : type->children()) { + bufferSizes[bufferTypes.size()] = arrow::bit_util::BytesForBits(numRows); + bufferTypes.push_back(BufferType::kNull); + if (colType->isFixedWidth()) { + bufferTypes.push_back(BufferType::kValue); + } else { + // Add the first offset 0. + // offsetBuffer length is more than original length. + bufferSizes[bufferTypes.size()] = sizeof(int32_t) * (numRows + 1); + bufferTypes.push_back(BufferType::kLength); + bufferTypes.push_back(BufferType::kValue); + } + } + VELOX_CHECK_EQ(bufferTypes.size(), bufferSize); + // This buffer may be more than the actual reauired buffer for null buffer. + for (const auto& batch : batches) { + if (batch->numRows() == 0) { + continue; + } + for (auto i = 0; i < bufferSize; ++i) { + // The null buffer may be null or length = 0. + // Maybe optimize later, detect if the null buffer is all true. And set the return null buffer to 0. + if (bufferTypes[i] == BufferType::kNull || bufferTypes[i] == BufferType::kLength) { + continue; + } + auto& buffer = batch->bufferAt(i); + VELOX_CHECK_NOT_NULL(buffer); + bufferSizes[i] += buffer->size(); + } + } + + std::vector> returnBuffers; + returnBuffers.reserve(bufferSize); + for (auto size : bufferSizes) { + std::shared_ptr buffer; + // May optimize to reuse the first batch buffer. + GLUTEN_ASSIGN_OR_THROW(buffer, arrow::AllocateResizableBuffer(size, pool)); + returnBuffers.emplace_back(std::move(buffer)); + } + + int32_t bufferIdx = 0; + for (const auto& colType : type->children()) { + size_t rowNumber = 0; + // Also records the value buffer offset. + size_t stringOffset = 0; + for (auto i = 0; i < batches.size(); ++i) { + const auto& batch = batches[i]; + if (batch->numRows() == 0) { + continue; + } + // Combine the null buffer + // The last byte may still have space to write when nullBitsRemainder != 0. + auto* dst = returnBuffers[bufferIdx]->mutable_data(); + if (batch->bufferAt(bufferIdx) == nullptr) { + arrow::bit_util::SetBitsTo(dst, rowNumber, batch->numRows(), true); + } else { + arrow::internal::CopyBitmap(batch->bufferAt(bufferIdx)->data(), 0, batch->numRows(), dst, rowNumber); + } + + if (colType->isFixedWidth()) { + // The buffer is values. + const auto bufferSize = batch->bufferAt(bufferIdx + 1)->size(); + VELOX_CHECK_LE(stringOffset + bufferSize, returnBuffers[bufferIdx + 1]->size()); + memcpy( + returnBuffers[bufferIdx + 1]->mutable_data() + stringOffset, + batch->bufferAt(bufferIdx + 1)->data(), + bufferSize); + stringOffset += bufferSize; + } else { + // String, lengths, values + memcpy( + returnBuffers[bufferIdx + 2]->mutable_data() + stringOffset, + batch->bufferAt(bufferIdx + 2)->data(), + batch->bufferAt(bufferIdx + 2)->size()); + const auto* lengths = reinterpret_cast(batch->bufferAt(bufferIdx + 1)->data()); + auto* offsetBuffer = reinterpret_cast(returnBuffers[bufferIdx + 1]->mutable_data()); + for (auto j = 0; j < batch->numRows(); ++j) { + offsetBuffer[rowNumber + j] = stringOffset; + stringOffset += lengths[j]; + } + offsetBuffer[rowNumber + batch->numRows()] = stringOffset; + } + rowNumber += batch->numRows(); + } + if (colType->isFixedWidth()) { + bufferIdx += 2; + } else { + bufferIdx += 3; + } + } + + return std::make_shared(batches[0]->getRowType(), std::move(returnBuffers), numRows); +} + +} // namespace gluten diff --git a/cpp/velox/memory/GpuBufferColumnarBatch.h b/cpp/velox/memory/GpuBufferColumnarBatch.h new file mode 100644 index 00000000000..40ff4a8eb8a --- /dev/null +++ b/cpp/velox/memory/GpuBufferColumnarBatch.h @@ -0,0 +1,73 @@ +/* + * 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. + */ + +#pragma once + +#include "memory/ColumnarBatch.h" +#include "memory/VeloxMemoryManager.h" +#include "velox/vector/ComplexVector.h" +#include "velox/vector/arrow/Bridge.h" + +#include +#include +#include +#include + +namespace gluten { + +class GpuBufferColumnarBatch final : public ColumnarBatch { + public: + GpuBufferColumnarBatch( + facebook::velox::RowTypePtr rowType, + std::vector>&& buffers, + int32_t numRows) + : ColumnarBatch(rowType->children().size(), numRows), rowType_(rowType), buffers_(std::move(buffers)) {} + + std::string getType() const override { + return kType; + } + + const facebook::velox::RowTypePtr& getRowType() const { + return rowType_; + } + + const std::vector>& buffers() const { + return buffers_; + } + + const std::shared_ptr& bufferAt(size_t i) const { + return buffers_[i]; + } + + static std::shared_ptr compose( + arrow::MemoryPool* pool, + const std::vector>& batches, + int32_t numRows); + + int64_t numBytes() override; + + std::shared_ptr exportArrowSchema() override; + std::shared_ptr exportArrowArray() override; + std::vector toUnsafeRow(int32_t rowId) const override; + + private: + inline static const std::string kType{"gpu"}; + facebook::velox::RowTypePtr rowType_; + std::vector> buffers_; +}; + +} // namespace gluten diff --git a/cpp/velox/operators/plannodes/CudfVectorStream.h b/cpp/velox/operators/plannodes/CudfVectorStream.h new file mode 100644 index 00000000000..4707b8789d5 --- /dev/null +++ b/cpp/velox/operators/plannodes/CudfVectorStream.h @@ -0,0 +1,153 @@ +/* + * 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. + */ + +#pragma once + +#include "CudfVectorStream.h" +#include "velox/experimental/cudf/exec/NvtxHelper.h" +#include "velox/experimental/cudf/exec/Utilities.h" +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" +#include "velox/experimental/cudf/vector/CudfVector.h" + +namespace gluten { + +class CudfVectorStream : public RowVectorStream { + public: + CudfVectorStream( + facebook::velox::exec::DriverCtx* driverCtx, + facebook::velox::memory::MemoryPool* pool, + ResultIterator* iterator, + const facebook::velox::RowTypePtr& outputType) + : RowVectorStream(driverCtx, pool, iterator, outputType) {} + + // Convert arrow batch to row vector and use new output columns + facebook::velox::RowVectorPtr next() override { + auto cb = nextInternal(); + if (cb == nullptr) { + return nullptr; + } + auto vb = std::dynamic_pointer_cast(cb); + VELOX_CHECK_NOT_NULL(vb); + auto vp = vb->getRowVector(); + VELOX_DCHECK(vp != nullptr); + auto cudfVector = std::dynamic_pointer_cast(vp); + if (cudfVector == nullptr) { + // The vector may comes from BroadcastExchange, in this case, it's not a CudfVector. + vp->setType(outputType_); + return vp; + } + VELOX_CHECK_NOT_NULL(cudfVector); + return std::make_shared( + vp->pool(), outputType_, vp->size(), cudfVector->release(), cudfVector->stream()); + } +}; + +// To avoid plan translator uses false node, this one cannot inherit ValueStreamNode. +class CudfValueStreamNode final : public facebook::velox::core::PlanNode { + public: + CudfValueStreamNode( + const facebook::velox::core::PlanNodeId& id, + const facebook::velox::RowTypePtr& outputType, + std::shared_ptr iterator) + : facebook::velox::core::PlanNode(id), outputType_(outputType), iterator_(std::move(iterator)) {} + + const facebook::velox::RowTypePtr& outputType() const override { + return outputType_; + } + + const std::vector& sources() const override { + return kEmptySources_; + }; + + ResultIterator* iterator() const { + return iterator_.get(); + } + + std::string_view name() const override { + return "CudfValueStream"; + } + + folly::dynamic serialize() const override { + VELOX_UNSUPPORTED("CudfValueStream plan node is not serializable"); + } + + private: + void addDetails(std::stringstream& stream) const override{}; + + const facebook::velox::RowTypePtr outputType_; + std::shared_ptr iterator_; + const std::vector kEmptySources_; +}; + +// Extends NvtxHelper to identify it as GPU node, so not add CudfFormVelox operator. +class CudfValueStream : public facebook::velox::exec::SourceOperator, public facebook::velox::cudf_velox::NvtxHelper { + public: + CudfValueStream( + int32_t operatorId, + facebook::velox::exec::DriverCtx* driverCtx, + std::shared_ptr valueStreamNode) + : facebook::velox::exec::SourceOperator( + driverCtx, + valueStreamNode->outputType(), + operatorId, + valueStreamNode->id(), + valueStreamNode->name().data()), + facebook::velox::cudf_velox::NvtxHelper( + nvtx3::rgb{160, 82, 45}, // Sienna + operatorId, + fmt::format("[{}]", valueStreamNode->id())) { + ResultIterator* itr = valueStreamNode->iterator(); + rvStream_ = std::make_unique(driverCtx, pool(), itr, outputType_); + } + + facebook::velox::RowVectorPtr getOutput() override { + if (finished_) { + return nullptr; + } + if (rvStream_->hasNext()) { + return rvStream_->next(); + } else { + finished_ = true; + return nullptr; + } + } + + facebook::velox::exec::BlockingReason isBlocked(facebook::velox::ContinueFuture* /* unused */) override { + return facebook::velox::exec::BlockingReason::kNotBlocked; + } + + bool isFinished() override { + return finished_; + } + + private: + bool finished_ = false; + std::unique_ptr rvStream_; +}; + +class CudfVectorStreamOperatorTranslator : public facebook::velox::exec::Operator::PlanNodeTranslator { + std::unique_ptr toOperator( + facebook::velox::exec::DriverCtx* ctx, + int32_t id, + const facebook::velox::core::PlanNodePtr& node) override { + if (auto valueStreamNode = std::dynamic_pointer_cast(node)) { + return std::make_unique(id, ctx, valueStreamNode); + } + return nullptr; + } +}; +} // namespace gluten diff --git a/cpp/velox/operators/plannodes/RowVectorStream.cc b/cpp/velox/operators/plannodes/RowVectorStream.cc new file mode 100644 index 00000000000..7121fe93ed1 --- /dev/null +++ b/cpp/velox/operators/plannodes/RowVectorStream.cc @@ -0,0 +1,95 @@ +/* + * 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. + */ + +#include "RowVectorStream.h" +#include "memory/VeloxColumnarBatch.h" +#include "velox/exec/Driver.h" +#include "velox/exec/Operator.h" +#include "velox/exec/Task.h" + +namespace { + +class SuspendedSection { + public: + explicit SuspendedSection(facebook::velox::exec::Driver* driver) : driver_(driver) { + if (driver_->task()->enterSuspended(driver->state()) != facebook::velox::exec::StopReason::kNone) { + VELOX_FAIL("Terminate detected when entering suspended section"); + } + } + + virtual ~SuspendedSection() { + if (driver_->task()->leaveSuspended(driver_->state()) != facebook::velox::exec::StopReason::kNone) { + LOG(WARNING) << "Terminate detected when leaving suspended section for driver " << driver_->driverCtx()->driverId + << " from task " << driver_->task()->taskId(); + } + } + + private: + facebook::velox::exec::Driver* const driver_; +}; + +} // namespace + +namespace gluten { +bool RowVectorStream::hasNext() { + if (finished_) { + return false; + } + VELOX_DCHECK_NOT_NULL(iterator_); + + bool hasNext; + { + // We are leaving Velox task execution and are probably entering Spark code through JNI. Suspend the current + // driver to make the current task open to spilling. + // + // When a task is getting spilled, it should have been suspended so has zero running threads, otherwise there's + // possibility that this spill call hangs. See https://github.com/apache/incubator-gluten/issues/7243. + // As of now, non-zero running threads usually happens when: + // 1. Task A spills task B; + // 2. Task A tries to grow buffers created by task B, during which spill is requested on task A again. + SuspendedSection ss(driverCtx_->driver); + hasNext = iterator_->hasNext(); + } + if (!hasNext) { + finished_ = true; + } + return hasNext; +} + +std::shared_ptr RowVectorStream::nextInternal() { + if (finished_) { + return nullptr; + } + std::shared_ptr cb; + { + // We are leaving Velox task execution and are probably entering Spark code through JNI. Suspend the current + // driver to make the current task open to spilling. + SuspendedSection ss(driverCtx_->driver); + cb = iterator_->next(); + } + return cb; +} + +facebook::velox::RowVectorPtr RowVectorStream::next() { + auto cb = nextInternal(); + const std::shared_ptr& vb = VeloxColumnarBatch::from(pool_, cb); + auto vp = vb->getRowVector(); + VELOX_DCHECK(vp != nullptr); + return std::make_shared( + vp->pool(), outputType_, facebook::velox::BufferPtr(0), vp->size(), vp->children()); +} +} // namespace gluten diff --git a/cpp/velox/operators/plannodes/RowVectorStream.h b/cpp/velox/operators/plannodes/RowVectorStream.h index 0c561124509..78679c20a41 100644 --- a/cpp/velox/operators/plannodes/RowVectorStream.h +++ b/cpp/velox/operators/plannodes/RowVectorStream.h @@ -23,29 +23,6 @@ #include "velox/exec/Operator.h" #include "velox/exec/Task.h" -namespace { - -class SuspendedSection { - public: - explicit SuspendedSection(facebook::velox::exec::Driver* driver) : driver_(driver) { - if (driver_->task()->enterSuspended(driver->state()) != facebook::velox::exec::StopReason::kNone) { - VELOX_FAIL("Terminate detected when entering suspended section"); - } - } - - virtual ~SuspendedSection() { - if (driver_->task()->leaveSuspended(driver_->state()) != facebook::velox::exec::StopReason::kNone) { - LOG(WARNING) << "Terminate detected when leaving suspended section for driver " << driver_->driverCtx()->driverId - << " from task " << driver_->task()->taskId(); - } - } - - private: - facebook::velox::exec::Driver* const driver_; -}; - -} // namespace - namespace gluten { class RowVectorStream { @@ -57,51 +34,15 @@ class RowVectorStream { const facebook::velox::RowTypePtr& outputType) : driverCtx_(driverCtx), pool_(pool), outputType_(outputType), iterator_(iterator) {} - bool hasNext() { - if (finished_) { - return false; - } - VELOX_DCHECK_NOT_NULL(iterator_); - - bool hasNext; - { - // We are leaving Velox task execution and are probably entering Spark code through JNI. Suspend the current - // driver to make the current task open to spilling. - // - // When a task is getting spilled, it should have been suspended so has zero running threads, otherwise there's - // possibility that this spill call hangs. See https://github.com/apache/incubator-gluten/issues/7243. - // As of now, non-zero running threads usually happens when: - // 1. Task A spills task B; - // 2. Task A tries to grow buffers created by task B, during which spill is requested on task A again. - SuspendedSection ss(driverCtx_->driver); - hasNext = iterator_->hasNext(); - } - if (!hasNext) { - finished_ = true; - } - return hasNext; - } + bool hasNext(); - // Convert arrow batch to row vector and use new output columns - facebook::velox::RowVectorPtr next() { - if (finished_) { - return nullptr; - } - std::shared_ptr cb; - { - // We are leaving Velox task execution and are probably entering Spark code through JNI. Suspend the current - // driver to make the current task open to spilling. - SuspendedSection ss(driverCtx_->driver); - cb = iterator_->next(); - } - const std::shared_ptr& vb = VeloxColumnarBatch::from(pool_, cb); - auto vp = vb->getRowVector(); - VELOX_DCHECK(vp != nullptr); - return std::make_shared( - vp->pool(), outputType_, facebook::velox::BufferPtr(0), vp->size(), vp->children()); - } + // Convert arrow batch to row vector, construct the new Rowvector with new outputType. + virtual facebook::velox::RowVectorPtr next(); + + protected: + // Get the next batch from iterator_. + std::shared_ptr nextInternal(); - private: facebook::velox::exec::DriverCtx* driverCtx_; facebook::velox::memory::MemoryPool* pool_; const facebook::velox::RowTypePtr outputType_; diff --git a/cpp/velox/shuffle/GpuShuffleReader.cc b/cpp/velox/shuffle/GpuShuffleReader.cc deleted file mode 100644 index fefa473bf14..00000000000 --- a/cpp/velox/shuffle/GpuShuffleReader.cc +++ /dev/null @@ -1,393 +0,0 @@ -/* - * 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. - */ - -#include "shuffle/GpuShuffleReader.h" - -#include -#include - -#include "memory/VeloxColumnarBatch.h" -#include "shuffle/Payload.h" -#include "shuffle/Utils.h" -#include "utils/Common.h" -#include "utils/Macros.h" -#include "utils/Timer.h" -#include "utils/VeloxArrowUtils.h" -#include "velox/common/caching/AsyncDataCache.h" -#include "velox/serializers/PrestoSerializer.h" -#include "velox/vector/ComplexVector.h" -#include "velox/vector/FlatVector.h" -#include "velox/vector/arrow/Bridge.h" - -#include - -#include "cudf/GpuLock.h" - -using namespace facebook::velox; - -namespace gluten { -namespace { - -arrow::Result readBlockType(arrow::io::InputStream* inputStream) { - BlockType type; - ARROW_ASSIGN_OR_RAISE(auto bytes, inputStream->Read(sizeof(BlockType), &type)); - if (bytes == 0) { - // Reach EOS. - return BlockType::kEndOfStream; - } - return type; -} - -struct BufferViewReleaser { - BufferViewReleaser() : BufferViewReleaser(nullptr) {} - - BufferViewReleaser(std::shared_ptr arrowBuffer) : bufferReleaser_(std::move(arrowBuffer)) {} - - void addRef() const {} - - void release() const {} - - private: - const std::shared_ptr bufferReleaser_; -}; - -BufferPtr wrapInBufferViewAsOwner(const void* buffer, size_t length, std::shared_ptr bufferReleaser) { - return BufferView::create( - static_cast(buffer), length, {std::move(bufferReleaser)}); -} - -BufferPtr convertToVeloxBuffer(std::shared_ptr buffer) { - if (buffer == nullptr) { - return nullptr; - } - return wrapInBufferViewAsOwner(buffer->data(), buffer->size(), buffer); -} - -template ::NativeType> -VectorPtr readFlatVector( - std::vector& buffers, - int32_t& bufferIdx, - uint32_t length, - std::shared_ptr type, - const VectorPtr& dictionary, - memory::MemoryPool* pool) { - auto nulls = buffers[bufferIdx++]; - auto valuesOrIndices = buffers[bufferIdx++]; - - nulls = nulls == nullptr || nulls->size() == 0 ? BufferPtr(nullptr) : nulls; - - if (dictionary != nullptr) { - return BaseVector::wrapInDictionary(nulls, valuesOrIndices, length, dictionary); - } - - return std::make_shared>( - pool, type, nulls, length, std::move(valuesOrIndices), std::vector{}); -} - -template <> -VectorPtr readFlatVector( - std::vector& buffers, - int32_t& bufferIdx, - uint32_t length, - std::shared_ptr type, - const VectorPtr& dictionary, - memory::MemoryPool* pool) { - return BaseVector::createNullConstant(type, length, pool); -} - -template <> -VectorPtr readFlatVector( - std::vector& buffers, - int32_t& bufferIdx, - uint32_t length, - std::shared_ptr type, - const VectorPtr& dictionary, - memory::MemoryPool* pool) { - auto nulls = buffers[bufferIdx++]; - auto valuesOrIndices = buffers[bufferIdx++]; - - // Because if buffer does not compress, it will get from netty, the address maynot aligned 16B, which will cause - // int128_t = xxx coredump by instruction movdqa - const auto* addr = valuesOrIndices->as(); - if ((reinterpret_cast(addr) & 0xf) != 0) { - auto alignedBuffer = AlignedBuffer::allocate(valuesOrIndices->size(), pool); - fastCopy(alignedBuffer->asMutable(), valuesOrIndices->as(), valuesOrIndices->size()); - valuesOrIndices = alignedBuffer; - } - - nulls = nulls == nullptr || nulls->size() == 0 ? BufferPtr(nullptr) : nulls; - - if (dictionary != nullptr) { - return BaseVector::wrapInDictionary(nulls, valuesOrIndices, length, dictionary); - } - - return std::make_shared>( - pool, type, nulls, length, std::move(valuesOrIndices), std::vector{}); -} - -VectorPtr readFlatVectorStringView( - std::vector& buffers, - int32_t& bufferIdx, - uint32_t length, - std::shared_ptr type, - const VectorPtr& dictionary, - memory::MemoryPool* pool) { - auto nulls = buffers[bufferIdx++]; - auto lengthOrIndices = buffers[bufferIdx++]; - - nulls = nulls == nullptr || nulls->size() == 0 ? BufferPtr(nullptr) : nulls; - - if (dictionary != nullptr) { - return BaseVector::wrapInDictionary(nulls, lengthOrIndices, length, dictionary); - } - - auto valueBuffer = buffers[bufferIdx++]; - - const auto* rawLength = lengthOrIndices->as(); - const auto* valueBufferPtr = valueBuffer->as(); - - auto values = AlignedBuffer::allocate(sizeof(StringView) * length, pool); - auto* rawValues = values->asMutable(); - - uint64_t offset = 0; - for (int32_t i = 0; i < length; ++i) { - rawValues[i] = StringView(valueBufferPtr + offset, rawLength[i]); - offset += rawLength[i]; - } - - std::vector stringBuffers; - stringBuffers.emplace_back(valueBuffer); - - return std::make_shared>( - pool, type, nulls, length, std::move(values), std::move(stringBuffers)); -} - -template <> -VectorPtr readFlatVector( - std::vector& buffers, - int32_t& bufferIdx, - uint32_t length, - std::shared_ptr type, - const VectorPtr& dictionary, - memory::MemoryPool* pool) { - return readFlatVectorStringView(buffers, bufferIdx, length, type, dictionary, pool); -} - -template <> -VectorPtr readFlatVector( - std::vector& buffers, - int32_t& bufferIdx, - uint32_t length, - std::shared_ptr type, - const VectorPtr& dictionary, - memory::MemoryPool* pool) { - return readFlatVectorStringView(buffers, bufferIdx, length, type, dictionary, pool); -} - -std::unique_ptr toByteStream(uint8_t* data, int32_t size) { - std::vector byteRanges; - byteRanges.push_back(ByteRange{data, size, 0}); - auto byteStream = std::make_unique(byteRanges); - return byteStream; -} - -RowVectorPtr readComplexType(BufferPtr buffer, RowTypePtr& rowType, memory::MemoryPool* pool) { - RowVectorPtr result; - auto byteStream = toByteStream(const_cast(buffer->as()), buffer->size()); - auto serde = std::make_unique(); - serializer::presto::PrestoVectorSerde::PrestoOptions options; - options.useLosslessTimestamp = true; - serde->deserialize(byteStream.get(), pool, rowType, &result, &options); - return result; -} - -RowTypePtr getComplexWriteType(const std::vector& types) { - std::vector complexTypeColNames; - std::vector complexTypeChildrens; - for (int32_t i = 0; i < types.size(); ++i) { - auto kind = types[i]->kind(); - switch (kind) { - case TypeKind::ROW: - case TypeKind::MAP: - case TypeKind::ARRAY: { - complexTypeColNames.emplace_back(types[i]->name()); - complexTypeChildrens.emplace_back(types[i]); - } break; - default: - break; - } - } - return std::make_shared(std::move(complexTypeColNames), std::move(complexTypeChildrens)); -} - -RowVectorPtr deserialize( - RowTypePtr type, - uint32_t numRows, - std::vector& buffers, - const std::vector& dictionaryFields, - const std::vector& dictionaries, - memory::MemoryPool* pool) { - std::vector children; - auto types = type->as().children(); - - std::vector complexChildren; - auto complexRowType = getComplexWriteType(types); - if (complexRowType->children().size() > 0) { - complexChildren = readComplexType(buffers[buffers.size() - 1], complexRowType, pool)->children(); - } - - int32_t bufferIdx = 0; - int32_t complexIdx = 0; - int32_t dictionaryIdx = 0; - for (size_t i = 0; i < types.size(); ++i) { - const auto kind = types[i]->kind(); - switch (kind) { - case TypeKind::ROW: - case TypeKind::MAP: - case TypeKind::ARRAY: { - children.emplace_back(std::move(complexChildren[complexIdx])); - complexIdx++; - } break; - default: { - VectorPtr dictionary{nullptr}; - if (!dictionaryFields.empty() && dictionaryIdx < dictionaryFields.size() && - dictionaryFields[dictionaryIdx] == i) { - dictionary = dictionaries[dictionaryIdx++]; - } - auto res = VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH_ALL( - readFlatVector, kind, buffers, bufferIdx, numRows, types[i], dictionary, pool); - children.emplace_back(std::move(res)); - } break; - } - } - - return std::make_shared(pool, type, BufferPtr(nullptr), numRows, children); -} - -std::shared_ptr makeColumnarBatch( - RowTypePtr type, - uint32_t numRows, - std::vector> arrowBuffers, - const std::vector& dictionaryFields, - const std::vector& dictionaries, - memory::MemoryPool* pool, - int64_t& deserializeTime) { - ScopedTimer timer(&deserializeTime); - std::vector veloxBuffers; - veloxBuffers.reserve(arrowBuffers.size()); - for (auto& buffer : arrowBuffers) { - veloxBuffers.push_back(convertToVeloxBuffer(std::move(buffer))); - } - auto rowVector = deserialize(type, numRows, veloxBuffers, dictionaryFields, dictionaries, pool); - return std::make_shared(std::move(rowVector)); -} - -} // namespace - - -GpuHashShuffleReaderDeserializer::GpuHashShuffleReaderDeserializer( - const std::shared_ptr& streamReader, - const std::shared_ptr& schema, - const std::shared_ptr& codec, - const facebook::velox::RowTypePtr& rowType, - int32_t batchSize, - int64_t readerBufferSize, - VeloxMemoryManager* memoryManager, - std::vector* isValidityBuffer, - bool hasComplexType, - int64_t& deserializeTime, - int64_t& decompressTime) - : streamReader_(streamReader), - schema_(schema), - codec_(codec), - rowType_(rowType), - batchSize_(batchSize), - readerBufferSize_(readerBufferSize), - memoryManager_(memoryManager), - isValidityBuffer_(isValidityBuffer), - hasComplexType_(hasComplexType), - deserializeTime_(deserializeTime), - decompressTime_(decompressTime) {} - -bool GpuHashShuffleReaderDeserializer::resolveNextBlockType() { - GLUTEN_ASSIGN_OR_THROW(auto blockType, readBlockType(in_.get())); - switch (blockType) { - case BlockType::kEndOfStream: - return false; - case BlockType::kPlainPayload: - return true; - default: - throw GlutenException(fmt::format("Unsupported block type: {}", static_cast(blockType))); - } - return true; -} - -void GpuHashShuffleReaderDeserializer::loadNextStream() { - if (reachedEos_) { - return; - } - - auto in = streamReader_->readNextStream(memoryManager_->defaultArrowMemoryPool()); - if (in == nullptr) { - reachedEos_ = true; - return; - } - - GLUTEN_ASSIGN_OR_THROW( - in_, - arrow::io::BufferedInputStream::Create( - readerBufferSize_, memoryManager_->defaultArrowMemoryPool(), std::move(in))); -} - -std::shared_ptr GpuHashShuffleReaderDeserializer::next() { - if (in_ == nullptr) { - loadNextStream(); - - if (reachedEos_) { - return nullptr; - } - } - - while (!resolveNextBlockType()) { - loadNextStream(); - - if (reachedEos_) { - return nullptr; - } - } - - uint32_t numRows = 0; - GLUTEN_ASSIGN_OR_THROW( - auto arrowBuffers, - BlockPayload::deserialize( - in_.get(), codec_, memoryManager_->defaultArrowMemoryPool(), numRows, deserializeTime_, decompressTime_)); - - auto batch = makeColumnarBatch( - rowType_, - numRows, - std::move(arrowBuffers), - dictionaryFields_, - dictionaries_, - memoryManager_->getLeafMemoryPool().get(), - deserializeTime_); - - lockGpu(); - - return batch; -} - -} diff --git a/cpp/velox/shuffle/VeloxGpuShuffleReader.cc b/cpp/velox/shuffle/VeloxGpuShuffleReader.cc new file mode 100644 index 00000000000..441cb0d0667 --- /dev/null +++ b/cpp/velox/shuffle/VeloxGpuShuffleReader.cc @@ -0,0 +1,124 @@ +/* + * 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. + */ + +#include "VeloxGpuShuffleReader.h" + +#include +#include + +#include "memory/VeloxColumnarBatch.h" +#include "shuffle/Payload.h" +#include "shuffle/Utils.h" +#include "utils/Common.h" +#include "utils/Macros.h" +#include "utils/Timer.h" +#include "memory/GpuBufferColumnarBatch.h" + +#include + +using namespace facebook::velox; + +namespace gluten { +namespace { + +arrow::Result readBlockType(arrow::io::InputStream* inputStream) { + BlockType type; + ARROW_ASSIGN_OR_RAISE(auto bytes, inputStream->Read(sizeof(BlockType), &type)); + if (bytes == 0) { + // Reach EOS. + return BlockType::kEndOfStream; + } + return type; +} + +} // namespace + +VeloxGpuHashShuffleReaderDeserializer::VeloxGpuHashShuffleReaderDeserializer( + const std::shared_ptr& streamReader, + const std::shared_ptr& schema, + const std::shared_ptr& codec, + const facebook::velox::RowTypePtr& rowType, + int64_t readerBufferSize, + VeloxMemoryManager* memoryManager, + int64_t& deserializeTime, + int64_t& decompressTime) + : streamReader_(streamReader), + schema_(schema), + codec_(codec), + rowType_(rowType), + readerBufferSize_(readerBufferSize), + memoryManager_(memoryManager), + deserializeTime_(deserializeTime), + decompressTime_(decompressTime) {} + +bool VeloxGpuHashShuffleReaderDeserializer::resolveNextBlockType() { + GLUTEN_ASSIGN_OR_THROW(auto blockType, readBlockType(in_.get())); + switch (blockType) { + case BlockType::kEndOfStream: + return false; + case BlockType::kPlainPayload: + return true; + default: + throw GlutenException(fmt::format("Unsupported block type: {}", static_cast(blockType))); + } + return true; +} + +void VeloxGpuHashShuffleReaderDeserializer::loadNextStream() { + if (reachedEos_) { + return; + } + + auto in = streamReader_->readNextStream(memoryManager_->defaultArrowMemoryPool()); + if (in == nullptr) { + reachedEos_ = true; + return; + } + + GLUTEN_ASSIGN_OR_THROW( + in_, + arrow::io::BufferedInputStream::Create( + readerBufferSize_, memoryManager_->defaultArrowMemoryPool(), std::move(in))); +} + +std::shared_ptr VeloxGpuHashShuffleReaderDeserializer::next() { + if (in_ == nullptr) { + loadNextStream(); + + if (reachedEos_) { + return nullptr; + } + } + + while (!resolveNextBlockType()) { + loadNextStream(); + + if (reachedEos_) { + return nullptr; + } + } + + uint32_t numRows = 0; + GLUTEN_ASSIGN_OR_THROW( + auto arrowBuffers, + BlockPayload::deserialize( + in_.get(), codec_, memoryManager_->defaultArrowMemoryPool(), numRows, deserializeTime_, decompressTime_)); + + return std::make_shared(rowType_, std::move(arrowBuffers), static_cast(numRows)); +} + +} // namespace gluten diff --git a/cpp/velox/shuffle/GpuShuffleReader.h b/cpp/velox/shuffle/VeloxGpuShuffleReader.h similarity index 84% rename from cpp/velox/shuffle/GpuShuffleReader.h rename to cpp/velox/shuffle/VeloxGpuShuffleReader.h index 4617b0c2b61..498b28b6225 100644 --- a/cpp/velox/shuffle/GpuShuffleReader.h +++ b/cpp/velox/shuffle/VeloxGpuShuffleReader.h @@ -30,18 +30,15 @@ namespace gluten { /// Convert the buffers to cudf table. /// Add a lock after reader produces the Vector, relase the lock after the thread processes all the batches. /// After move the shuffle read operation to gpu, move the lock to start read. -class GpuHashShuffleReaderDeserializer final : public ColumnarBatchIterator { +class VeloxGpuHashShuffleReaderDeserializer final : public ColumnarBatchIterator { public: - GpuHashShuffleReaderDeserializer( + VeloxGpuHashShuffleReaderDeserializer( const std::shared_ptr& streamReader, const std::shared_ptr& schema, const std::shared_ptr& codec, const facebook::velox::RowTypePtr& rowType, - int32_t batchSize, int64_t readerBufferSize, VeloxMemoryManager* memoryManager, - std::vector* isValidityBuffer, - bool hasComplexType, int64_t& deserializeTime, int64_t& decompressTime); @@ -56,13 +53,9 @@ class GpuHashShuffleReaderDeserializer final : public ColumnarBatchIterator { std::shared_ptr schema_; std::shared_ptr codec_; facebook::velox::RowTypePtr rowType_; - int32_t batchSize_; int64_t readerBufferSize_; VeloxMemoryManager* memoryManager_; - std::vector* isValidityBuffer_; - bool hasComplexType_; - int64_t& deserializeTime_; int64_t& decompressTime_; @@ -70,9 +63,5 @@ class GpuHashShuffleReaderDeserializer final : public ColumnarBatchIterator { bool reachedEos_{false}; bool blockTypeResolved_{false}; - - // Not used. - std::vector dictionaryFields_{}; - std::vector dictionaries_{}; }; } // namespace gluten diff --git a/cpp/velox/shuffle/VeloxGpuShuffleWriter.cc b/cpp/velox/shuffle/VeloxGpuShuffleWriter.cc new file mode 100644 index 00000000000..f7233e7c1a4 --- /dev/null +++ b/cpp/velox/shuffle/VeloxGpuShuffleWriter.cc @@ -0,0 +1,72 @@ +/* + * 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. + */ + +#include "VeloxGpuShuffleWriter.h" + +namespace gluten { + +// Split bool bit to bytes. +void VeloxGpuHashShuffleWriter::splitBoolValueType(const uint8_t* srcAddr, const std::vector& dstAddrs) { + for (auto& pid : partitionUsed_) { + auto dstaddr = dstAddrs[pid]; + if (dstaddr == nullptr) { + continue; + } + auto dstPidBase = (uint8_t*)(dstaddr + partitionBufferBase_[pid] * sizeof(uint8_t)); + auto pos = partition2RowOffsetBase_[pid]; + auto end = partition2RowOffsetBase_[pid + 1]; + for (; pos < end; ++pos) { + auto rowId = rowOffset2RowId_[pos]; + // === Extract one bool value from the bit-packed source === + uint8_t byte = srcAddr[rowId >> 3]; // find the source byte + uint8_t bit = (byte >> (rowId & 7)) & 0x01; // extract 0 or 1 + *dstPidBase++ = bit; // copy + } + } +} + +// Split timestamp from int128_t to int64_t, both of them represents the timestamp nanoseconds. +arrow::Status VeloxGpuHashShuffleWriter::splitTimestamp(const uint8_t* srcAddr, const std::vector& dstAddrs) { + for (auto& pid : partitionUsed_) { + auto dstPidBase = (int64_t*)(dstAddrs[pid] + partitionBufferBase_[pid] * sizeof(int64_t)); + auto pos = partition2RowOffsetBase_[pid]; + auto end = partition2RowOffsetBase_[pid + 1]; + for (; pos < end; ++pos) { + auto rowId = rowOffset2RowId_[pos]; + auto* src = reinterpret_cast(srcAddr) + rowId * 2; + // src[0] is seconds, src[1] is nanoseconds in Timestamp. + *dstPidBase++ = src[0] * 1'000'000'000LL + src[1]; + } + } + return arrow::Status::OK(); +} + +arrow::Result> VeloxGpuHashShuffleWriter::create( + uint32_t numPartitions, + const std::shared_ptr& partitionWriter, + const std::shared_ptr& options, + MemoryManager* memoryManager) { + if (auto hashOptions = std::dynamic_pointer_cast(options)) { + std::shared_ptr res = + std::make_shared(numPartitions, partitionWriter, hashOptions, memoryManager); + RETURN_NOT_OK(res->init()); + return res; + } + return arrow::Status::Invalid("Error casting ShuffleWriterOptions to GpuHashShuffleWriterOptions. "); +} + +} // namespace gluten diff --git a/cpp/velox/shuffle/VeloxGpuShuffleWriter.h b/cpp/velox/shuffle/VeloxGpuShuffleWriter.h new file mode 100644 index 00000000000..e6f0e1a3586 --- /dev/null +++ b/cpp/velox/shuffle/VeloxGpuShuffleWriter.h @@ -0,0 +1,57 @@ +/* + * 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. + */ + +#pragma once + +#include "VeloxHashShuffleWriter.h" + +namespace gluten { + +class VeloxGpuHashShuffleWriter : public VeloxHashShuffleWriter { + public: + static arrow::Result> create( + uint32_t numPartitions, + const std::shared_ptr& partitionWriter, + const std::shared_ptr& options, + MemoryManager* memoryManager); + + VeloxGpuHashShuffleWriter( + uint32_t numPartitions, + const std::shared_ptr& partitionWriter, + const std::shared_ptr& options, + MemoryManager* memoryManager) + : VeloxHashShuffleWriter(numPartitions, partitionWriter, options, memoryManager) {} + + private: + // Split the bool to byte. + void splitBoolValueType(const uint8_t* srcAddr, const std::vector& dstAddrs) override; + + uint64_t valueBufferSizeForBool(uint32_t newSize) override { + return newSize; + } + + bool boolIsBit() override { + return false; + } + + arrow::Status splitTimestamp(const uint8_t* srcAddr, const std::vector& dstAddrs) override; + + uint64_t valueBufferSizeForTimestamp(uint32_t newSize) override { + return sizeof(int64_t) * newSize; + } +}; +} // namespace gluten diff --git a/cpp/velox/shuffle/VeloxHashShuffleWriter.cc b/cpp/velox/shuffle/VeloxHashShuffleWriter.cc index 5dc3377e995..dadee313337 100644 --- a/cpp/velox/shuffle/VeloxHashShuffleWriter.cc +++ b/cpp/velox/shuffle/VeloxHashShuffleWriter.cc @@ -460,7 +460,7 @@ arrow::Status VeloxHashShuffleWriter::splitFixedWidthValueBuffer(const facebook: // No value buffer created for NullType. break; case 1: // arrow::BooleanType::type_id: - RETURN_NOT_OK(splitBoolType(srcAddr, dstAddrs)); + splitBoolValueType(srcAddr, dstAddrs); break; case 8: RETURN_NOT_OK(splitFixedType(srcAddr, dstAddrs)); @@ -473,7 +473,7 @@ arrow::Status VeloxHashShuffleWriter::splitFixedWidthValueBuffer(const facebook: break; case 64: { if (column->type()->kind() == facebook::velox::TypeKind::TIMESTAMP) { - RETURN_NOT_OK(splitFixedType(srcAddr, dstAddrs)); + RETURN_NOT_OK(splitTimestamp(srcAddr, dstAddrs)); } else { RETURN_NOT_OK(splitFixedType(srcAddr, dstAddrs)); } @@ -502,7 +502,7 @@ arrow::Status VeloxHashShuffleWriter::splitFixedWidthValueBuffer(const facebook: return arrow::Status::OK(); } -arrow::Status VeloxHashShuffleWriter::splitBoolType(const uint8_t* srcAddr, const std::vector& dstAddrs) { +void VeloxHashShuffleWriter::splitBoolType(const uint8_t* srcAddr, const std::vector& dstAddrs) { // assume batch size = 32k; reducer# = 4K; row/reducer = 8 for (auto& pid : partitionUsed_) { // set the last byte @@ -588,7 +588,6 @@ arrow::Status VeloxHashShuffleWriter::splitBoolType(const uint8_t* srcAddr, cons dstaddr[dstOffset >> 3] = dst; } } - return arrow::Status::OK(); } arrow::Status VeloxHashShuffleWriter::splitValidityBuffer(const facebook::velox::RowVector& rv) { @@ -611,7 +610,7 @@ arrow::Status VeloxHashShuffleWriter::splitValidityBuffer(const facebook::velox: } auto srcAddr = (const uint8_t*)(column->mutableRawNulls()); - RETURN_NOT_OK(splitBoolType(srcAddr, dstAddrs)); + splitBoolType(srcAddr, dstAddrs); } else { VsPrintLF(colIdx, " column hasn't null"); } @@ -768,7 +767,7 @@ arrow::Status VeloxHashShuffleWriter::initColumnTypes(const facebook::velox::Row case arrow::BooleanType::type_id: { simpleColumnIndices_.push_back(i); isValidityBuffer_.push_back(true); - isValidityBuffer_.push_back(true); + isValidityBuffer_.push_back(boolIsBit()); } break; case arrow::NullType::type_id: break; @@ -1073,11 +1072,11 @@ arrow::Result>> VeloxHashShuffleWrite auto& valueBuffer = buffers[kFixedWidthValueBufferIndex]; ARROW_RETURN_IF(!valueBuffer, arrow::Status::Invalid("Value buffer of fixed-width array is null.")); if (arrowColumnTypes_[i]->id() == arrow::BooleanType::type_id) { - valueBufferSize = arrow::bit_util::BytesForBits(numRows); + valueBufferSize = valueBufferSizeForBool(numRows); } else if (veloxColumnTypes_[i]->isShortDecimal()) { valueBufferSize = numRows * (arrow::bit_width(arrow::Int64Type::type_id) >> 3); } else if (veloxColumnTypes_[i]->kind() == facebook::velox::TypeKind::TIMESTAMP) { - valueBufferSize = facebook::velox::BaseVector::byteSize(numRows); + valueBufferSize = valueBufferSizeForTimestamp(numRows); } else { valueBufferSize = numRows * (arrow::bit_width(arrowColumnTypes_[i]->id()) >> 3); } @@ -1273,15 +1272,23 @@ uint64_t VeloxHashShuffleWriter::valueBufferSizeForBinaryArray(uint32_t binaryId return (binaryArrayTotalSizeBytes_[binaryIdx] + totalInputNumRows_ - 1) / totalInputNumRows_ * newSize + 1024; } +uint64_t VeloxHashShuffleWriter::valueBufferSizeForBool(uint32_t newSize) { + return arrow::bit_util::BytesForBits(newSize); +} + +uint64_t VeloxHashShuffleWriter::valueBufferSizeForTimestamp(uint32_t newSize) { + return facebook::velox::BaseVector::byteSize(newSize); +} + uint64_t VeloxHashShuffleWriter::valueBufferSizeForFixedWidthArray(uint32_t fixedWidthIndex, uint32_t newSize) { uint64_t valueBufferSize = 0; auto columnIdx = simpleColumnIndices_[fixedWidthIndex]; if (arrowColumnTypes_[columnIdx]->id() == arrow::BooleanType::type_id) { - valueBufferSize = arrow::bit_util::BytesForBits(newSize); + valueBufferSize = valueBufferSizeForBool(newSize); } else if (veloxColumnTypes_[columnIdx]->isShortDecimal()) { valueBufferSize = newSize * (arrow::bit_width(arrow::Int64Type::type_id) >> 3); } else if (veloxColumnTypes_[columnIdx]->kind() == facebook::velox::TypeKind::TIMESTAMP) { - valueBufferSize = facebook::velox::BaseVector::byteSize(newSize); + valueBufferSize = valueBufferSizeForTimestamp(newSize); } else { valueBufferSize = newSize * (arrow::bit_width(arrowColumnTypes_[columnIdx]->id()) >> 3); } diff --git a/cpp/velox/shuffle/VeloxHashShuffleWriter.h b/cpp/velox/shuffle/VeloxHashShuffleWriter.h index fa2f3ea040a..da67413eb25 100644 --- a/cpp/velox/shuffle/VeloxHashShuffleWriter.h +++ b/cpp/velox/shuffle/VeloxHashShuffleWriter.h @@ -188,7 +188,25 @@ class VeloxHashShuffleWriter : public VeloxShuffleWriter { VS_PRINT_CONTAINER(input_has_null_); } - private: + protected: + virtual uint64_t valueBufferSizeForBool(uint32_t newSize); + + virtual void splitBoolValueType(const uint8_t* srcAddr, const std::vector& dstAddrs) { + splitBoolType(srcAddr, dstAddrs); + } + + virtual bool boolIsBit() { + return true; + } + + virtual uint64_t valueBufferSizeForTimestamp(uint32_t newSize); + + virtual arrow::Status splitTimestamp(const uint8_t* srcAddr, const std::vector& dstAddrs) { + return splitFixedType(srcAddr, dstAddrs); + } + + arrow::Status init(); + VeloxHashShuffleWriter( uint32_t numPartitions, const std::shared_ptr& partitionWriter, @@ -200,8 +218,7 @@ class VeloxHashShuffleWriter : public VeloxShuffleWriter { arenas_.resize(numPartitions); } - arrow::Status init(); - + private: arrow::Status initPartitions(); arrow::Status initColumnTypes(const facebook::velox::RowVector& rv); @@ -233,8 +250,6 @@ class VeloxHashShuffleWriter : public VeloxShuffleWriter { arrow::Status splitFixedWidthValueBuffer(const facebook::velox::RowVector& rv); - arrow::Status splitBoolType(const uint8_t* srcAddr, const std::vector& dstAddrs); - arrow::Status splitValidityBuffer(const facebook::velox::RowVector& rv); arrow::Status splitBinaryArray(const facebook::velox::RowVector& rv); @@ -249,6 +264,9 @@ class VeloxHashShuffleWriter : public VeloxShuffleWriter { arrow::Result>> assembleBuffers(uint32_t partitionId, bool reuseBuffers); + // Split bool value type and validity buffer. + void splitBoolType(const uint8_t* srcAddr, const std::vector& dstAddrs); + template arrow::Status splitFixedType(const uint8_t* srcAddr, const std::vector& dstAddrs) { for (auto& pid : partitionUsed_) { @@ -305,6 +323,7 @@ class VeloxHashShuffleWriter : public VeloxShuffleWriter { arrow::Status partitioningAndDoSplit(facebook::velox::RowVectorPtr rv, int64_t memLimit); + protected: int32_t splitBufferSize_; double splitBufferReallocThreshold_; diff --git a/cpp/velox/shuffle/VeloxShuffleReader.cc b/cpp/velox/shuffle/VeloxShuffleReader.cc index 2dac73224ab..a352e63fb09 100644 --- a/cpp/velox/shuffle/VeloxShuffleReader.cc +++ b/cpp/velox/shuffle/VeloxShuffleReader.cc @@ -38,7 +38,7 @@ #include #ifdef GLUTEN_ENABLE_GPU -#include "GpuShuffleReader.h" +#include "VeloxGpuShuffleReader.h" #endif using namespace facebook::velox; @@ -821,16 +821,14 @@ std::unique_ptr VeloxShuffleReaderDeserializerFactory::cr switch (shuffleWriterType_) { case ShuffleWriterType::kGpuHashShuffle: #ifdef GLUTEN_ENABLE_GPU - return std::make_unique( + VELOX_CHECK(!hasComplexType_); + return std::make_unique( streamReader, schema_, codec_, rowType_, - batchSize_, readerBufferSize_, memoryManager_, - &isValidityBuffer_, - hasComplexType_, deserializeTime_, decompressTime_); #endif diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.cc b/cpp/velox/shuffle/VeloxShuffleWriter.cc index c45d9419722..93e17310c70 100644 --- a/cpp/velox/shuffle/VeloxShuffleWriter.cc +++ b/cpp/velox/shuffle/VeloxShuffleWriter.cc @@ -20,6 +20,10 @@ #include "shuffle/VeloxRssSortShuffleWriter.h" #include "shuffle/VeloxSortShuffleWriter.h" +#ifdef GLUTEN_ENABLE_GPU +#include "VeloxGpuShuffleWriter.h" +#endif + namespace gluten { arrow::Result> VeloxShuffleWriter::create( @@ -36,7 +40,11 @@ arrow::Result> VeloxShuffleWriter::create( return VeloxSortShuffleWriter::create(numPartitions, std::move(partitionWriter), options, memoryManager); case ShuffleWriterType::kRssSortShuffle: return VeloxRssSortShuffleWriter::create(numPartitions, std::move(partitionWriter), options, memoryManager); - default: +#ifdef GLUTEN_ENABLE_GPU + case ShuffleWriterType::kGpuHashShuffle: + return VeloxGpuHashShuffleWriter::create(numPartitions, std::move(partitionWriter), options, memoryManager); +#endif + default: return arrow::Status::Invalid("Unsupported shuffle writer type: ", typeToString(type)); } } diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.cc b/cpp/velox/substrait/SubstraitToVeloxPlan.cc index 073888a12fe..6f1134edcc4 100644 --- a/cpp/velox/substrait/SubstraitToVeloxPlan.cc +++ b/cpp/velox/substrait/SubstraitToVeloxPlan.cc @@ -32,11 +32,11 @@ #include "config/VeloxConfig.h" #ifdef GLUTEN_ENABLE_GPU +#include "operators/plannodes/CudfVectorStream.h" #include "velox/experimental/cudf/connectors/hive/CudfHiveDataSink.h" #include "velox/experimental/cudf/connectors/hive/CudfHiveTableHandle.h" #include "velox/experimental/cudf/exec/ToCudf.h" #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" - using namespace cudf_velox::connector::hive; #endif @@ -848,7 +848,11 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait:: // Child should be either ProjectNode or ValueStreamNode in case of project fallback. VELOX_CHECK( (std::dynamic_pointer_cast(childNode) != nullptr || - std::dynamic_pointer_cast(childNode) != nullptr) && + std::dynamic_pointer_cast(childNode) != nullptr) +#ifdef GLUTEN_ENABLE_GPU + || std::dynamic_pointer_cast(childNode) != nullptr +#endif + && childNode->outputType()->size() > requiredChildOutput.size(), "injectedProject is true, but the ProjectNode or ValueStreamNode (in case of projection fallback)" " is missing or does not have the corresponding projection field"); @@ -1224,6 +1228,7 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait:: nextPlanNodeId(), sortingKeys, sortingOrders, static_cast(topNRel.n()), false /*isPartial*/, childNode); } +template core::PlanNodePtr SubstraitToVeloxPlanConverter::constructValueStreamNode( const ::substrait::ReadRel& readRel, int32_t streamIdx) { @@ -1251,7 +1256,7 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::constructValueStreamNode( VELOX_CHECK_LT(streamIdx, inputIters_.size(), "Could not find stream index {} in input iterator list.", streamIdx); iterator = inputIters_[streamIdx]; } - auto node = std::make_shared(nextPlanNodeId(), outputType, std::move(iterator)); + auto node = std::make_shared(nextPlanNodeId(), outputType, std::move(iterator)); auto splitInfo = std::make_shared(); splitInfo->isStream = true; @@ -1288,12 +1293,16 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait:: // Check if the ReadRel specifies an input of stream. If yes, build ValueStreamNode as the data source. auto streamIdx = getStreamIndex(readRel); if (streamIdx >= 0) { - // Only used in benchmark enable query trace, replace ValueStreamNode to ValuesNode to support serialization. - if (LIKELY(!veloxCfg_->get(kQueryTraceEnabled, false))) { - return constructValueStreamNode(readRel, streamIdx); - } else { - return constructValuesNode(readRel, streamIdx); +#ifdef GLUTEN_ENABLE_GPU + if (veloxCfg_->get(kCudfEnabled, kCudfEnabledDefault)) { + return constructValueStreamNode(readRel, streamIdx); } +#endif + if (!veloxCfg_->get(kQueryTraceEnabled, false)) { + return constructValueStreamNode(readRel, streamIdx); + } + // Only used in benchmark enable query trace, replace ValueStreamNode to ValuesNode to support serialization. + return constructValuesNode(readRel, streamIdx); } // Otherwise, will create TableScan node for ReadRel. @@ -1589,4 +1598,10 @@ bool SubstraitToVeloxPlanConverter::checkTypeExtension(const ::substrait::Plan& return true; } +#ifdef GLUTEN_ENABLE_GPU +template core::PlanNodePtr SubstraitToVeloxPlanConverter::constructValueStreamNode(const ::substrait::ReadRel& sRead, int32_t streamIdx); +#endif + +template core::PlanNodePtr SubstraitToVeloxPlanConverter::constructValueStreamNode(const ::substrait::ReadRel& sRead, int32_t streamIdx); + } // namespace gluten diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.h b/cpp/velox/substrait/SubstraitToVeloxPlan.h index 545bdae0179..8c0639d051e 100644 --- a/cpp/velox/substrait/SubstraitToVeloxPlan.h +++ b/cpp/velox/substrait/SubstraitToVeloxPlan.h @@ -133,6 +133,7 @@ class SubstraitToVeloxPlanConverter { /// FileProperties: the file sizes and modification times of the files to be scanned. core::PlanNodePtr toVeloxPlan(const ::substrait::ReadRel& sRead); + template core::PlanNodePtr constructValueStreamNode(const ::substrait::ReadRel& sRead, int32_t streamIdx); // This is only used in benchmark and enable query trace, which will load all the data to ValuesNode. diff --git a/cpp/velox/tests/CMakeLists.txt b/cpp/velox/tests/CMakeLists.txt index 4fdc95910ac..5ad7d8c7317 100644 --- a/cpp/velox/tests/CMakeLists.txt +++ b/cpp/velox/tests/CMakeLists.txt @@ -32,6 +32,84 @@ function(add_velox_test TEST_EXEC) gtest_discover_tests(${TEST_EXEC} DISCOVERY_MODE PRE_TEST) endfunction() +if(ENABLE_GPU) + function(import_library TARGET_NAME LIB_PATH) + if(NOT EXISTS ${LIB_PATH}) + message(FATAL_ERROR "Library does not exist: ${LIB_PATH}") + endif() + add_library(${TARGET_NAME} STATIC IMPORTED) + set_target_properties(${TARGET_NAME} PROPERTIES IMPORTED_LOCATION + ${LIB_PATH}) + endfunction() + + function(add_velox_gpu_test TEST_EXEC) + set(options) + set(one_value_args) + set(multi_value_args SOURCES) + cmake_parse_arguments(ARG "${options}" "${one_value_args}" + "${multi_value_args}" ${ARGN}) + + if(ARG_SOURCES) + set(SOURCES ${ARG_SOURCES}) + else() + message(FATAL_ERROR "No sources specified for test ${TEST_NAME}") + endif() + add_executable(${TEST_EXEC} ${SOURCES} ${VELOX_TEST_COMMON_SRCS}) + target_include_directories(${TEST_EXEC} PRIVATE ${CMAKE_SOURCE_DIR}/velox + ${CMAKE_SOURCE_DIR}/src) + import_library( + facebook::velox::velox_cudf_expression + ${VELOX_BUILD_PATH}/velox/experimental/cudf/expression/libvelox_cudf_expression.a + ) + import_library( + facebook::velox::velox_cudf_exec + ${VELOX_BUILD_PATH}/velox/experimental/cudf/exec/libvelox_cudf_exec.a) + import_library( + facebook::velox::velox_cudf_vector + ${VELOX_BUILD_PATH}/velox/experimental/cudf/vector/libvelox_cudf_vector.a) + import_library( + facebook::velox::velox_cudf_hive_connector + ${VELOX_BUILD_PATH}/velox/experimental/cudf/connectors/hive/libvelox_cudf_hive_connector.a + ) + target_include_directories( + ${TEST_EXEC} + PRIVATE ${VELOX_BUILD_PATH}/_deps/cudf-src/cpp/include + ${VELOX_BUILD_PATH}/_deps/rmm-src/cpp/include + ${VELOX_BUILD_PATH}/_deps/kvikio-src/cpp/include + ${VELOX_BUILD_PATH}/_deps/nvtx3-src/c/include + ${VELOX_BUILD_PATH}/_deps/nvcomp_proprietary_binary-src/include + ${VELOX_BUILD_PATH}/_deps/rapids_logger-src/include + /usr/local/cuda/include) + target_compile_definitions( + ${TEST_EXEC} PUBLIC LIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE) + target_link_libraries( + ${TEST_EXEC} + PRIVATE + ${VELOX_BUILD_PATH}/_deps/cudf-build/libcudf.so + ${VELOX_BUILD_PATH}/_deps/rmm-build/librmm.so + ${VELOX_BUILD_PATH}/_deps/kvikio-build/libkvikio.so + ${VELOX_BUILD_PATH}/_deps/nvcomp_proprietary_binary-src/lib64/libnvcomp.so + ${VELOX_BUILD_PATH}/_deps/nvcomp_proprietary_binary-src/lib64/libnvcomp_cpu.so + ${VELOX_BUILD_PATH}/_deps/rapids_logger-build/librapids_logger.so + /usr/local/cuda-12.8/lib64/libcudart.so.12) + target_link_libraries( + ${TEST_EXEC} + PUBLIC facebook::velox::velox_cudf_expression + facebook::velox::velox_cudf_exec + facebook::velox::velox_cudf_vector + facebook::velox::velox_cudf_hive_connector) + target_link_libraries(${TEST_EXEC} PRIVATE velox GTest::gtest + GTest::gtest_main) + gtest_discover_tests( + ${TEST_EXEC} DISCOVERY_MODE PRE_TEST + PROPERTIES LABELS "cuda_driver" + TIMEOUT 3000) + endfunction() + + add_velox_gpu_test(velox_gpu_shuffle_writer_test SOURCES + VeloxGpuShuffleWriterTest.cc) +endif() + set(VELOX_TEST_COMMON_SRCS JsonToProtoConverter.cc FilePathGenerator.cc) add_velox_test(velox_shuffle_writer_test SOURCES VeloxShuffleWriterTest.cc) diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc new file mode 100644 index 00000000000..3405548c517 --- /dev/null +++ b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc @@ -0,0 +1,741 @@ +/* + * 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. + */ + +#include +#include + +#include "config/GlutenConfig.h" +#include "shuffle/VeloxHashShuffleWriter.h" +#include "tests/VeloxShuffleWriterTestBase.h" +#include "tests/utils/TestAllocationListener.h" +#include "tests/utils/TestStreamReader.h" +#include "tests/utils/TestUtils.h" +#include "utils/VeloxArrowUtils.h" +#include "memory/GpuBufferColumnarBatch.h" +#include "utils/GpuBufferBatchResizer.h" + +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" +#include "velox/experimental/cudf/vector/CudfVector.h" +#include "velox/vector/tests/utils/VectorTestBase.h" + +using namespace facebook::velox; + +namespace gluten { + +namespace { + +class ColumnarBatchArray : public ColumnarBatchIterator { + public: + explicit ColumnarBatchArray(const std::vector> batches) + : batches_(std::move(batches)) {} + + std::shared_ptr next() override { + if (cursor_ >= batches_.size()) { + return nullptr; + } + return batches_[cursor_++]; + } + + private: + const std::vector> batches_; + int32_t cursor_ = 0; +}; + +struct GpuShuffleTestParams { + ShuffleWriterType shuffleWriterType; + PartitionWriterType partitionWriterType; + arrow::Compression::type compressionType; + int32_t compressionThreshold{0}; + int32_t mergeBufferSize{0}; + int32_t diskWriteBufferSize{0}; + bool enableDictionary{false}; + int64_t deserializerBufferSize{0}; + + std::string toString() const { + std::ostringstream out; + out << "shuffleWriterType = " << ShuffleWriter::typeToString(shuffleWriterType) + << ", partitionWriterType = " << PartitionWriter::typeToString(partitionWriterType) + << ", compressionType = " << arrow::util::Codec::GetCodecAsString(compressionType) + << ", compressionThreshold = " << compressionThreshold << ", mergeBufferSize = " << mergeBufferSize + << ", compressionBufferSize = " << diskWriteBufferSize + << ", deserializerBufferSize = " << deserializerBufferSize; + return out.str(); + } +}; + +std::vector takeRows( + const std::vector& sources, + const std::vector>& indices) { + std::vector result; + + for (size_t i = 0; i < sources.size(); ++i) { + if (indices[i].empty()) { + result.push_back(sources[i]); + continue; + } + + auto copy = RowVector::createEmpty(sources[0]->type(), sources[0]->pool()); + for (const auto row : indices[i]) { + GLUTEN_CHECK(row < sources[i]->size(), fmt::format("Index out of bound: {}", row)); + copy->append(sources[i]->slice(row, 1).get()); + } + result.push_back(std::move(copy)); + } + + return result; +} + +RowVectorPtr mergeRowVectors(const std::vector& sources) { + RowVectorPtr result = RowVector::createEmpty(sources[0]->type(), sources[0]->pool()); + + for (const auto& source : sources) { + result->append(source.get()); + } + + return result; +} + +RowVectorPtr mergeBufferColumnarBatches(std::vector>& bufferBatches) { + GpuBufferBatchResizer resizer( + getDefaultMemoryManager()->defaultArrowMemoryPool(), + getDefaultMemoryManager()->getLeafMemoryPool().get(), + 1200, // output one batch + std::make_unique(bufferBatches)); + auto cb = resizer.next(); + auto batch = std::dynamic_pointer_cast(cb); + VELOX_CHECK_NOT_NULL(batch); + auto vector = std::dynamic_pointer_cast(batch->getRowVector()); + VELOX_CHECK_NOT_NULL(vector); + + auto tableView = vector->getTableView(); + + // Convert back to Velox + return cudf_velox::with_arrow::toVeloxColumn( + tableView, getDefaultMemoryManager()->getLeafMemoryPool().get(), "", vector->stream()); +} + +std::vector getTestParams() { + static std::vector params{}; + + if (!params.empty()) { + return params; + } + + const std::vector compressions = { + arrow::Compression::UNCOMPRESSED, arrow::Compression::LZ4_FRAME, arrow::Compression::ZSTD}; + const std::vector compressionThresholds = {-1, 0, 3, 4, 10, 4096}; + const std::vector mergeBufferSizes = {0, 3, 4, 10, 4096}; + + for (const auto& compression : compressions) { + // Hash-based shuffle. + for (const auto compressionThreshold : compressionThresholds) { + // Local. + for (const auto mergeBufferSize : mergeBufferSizes) { + params.push_back( + GpuShuffleTestParams{ + .shuffleWriterType = ShuffleWriterType::kGpuHashShuffle, + .partitionWriterType = PartitionWriterType::kLocal, + .compressionType = compression, + .compressionThreshold = compressionThreshold, + .mergeBufferSize = mergeBufferSize}); + } + + // Rss. + params.push_back( + GpuShuffleTestParams{ + .shuffleWriterType = ShuffleWriterType::kGpuHashShuffle, + .partitionWriterType = PartitionWriterType::kRss, + .compressionType = compression, + .compressionThreshold = compressionThreshold}); + } + } + + return params; +} + +std::shared_ptr createPartitionWriter( + PartitionWriterType partitionWriterType, + uint32_t numPartitions, + const std::string& dataFile, + const std::vector& localDirs, + arrow::Compression::type compressionType, + int32_t mergeBufferSize, + int32_t compressionThreshold, + bool enableDictionary) { + GLUTEN_ASSIGN_OR_THROW(auto codec, arrow::util::Codec::Create(compressionType)); + switch (partitionWriterType) { + case PartitionWriterType::kLocal: { + auto options = std::make_shared(); + options->mergeBufferSize = mergeBufferSize; + options->compressionThreshold = compressionThreshold; + options->enableDictionary = enableDictionary; + return std::make_shared( + numPartitions, std::move(codec), getDefaultMemoryManager(), options, dataFile, std::move(localDirs)); + } + case PartitionWriterType::kRss: { + auto options = std::make_shared(); + auto rssClient = std::make_unique(dataFile); + return std::make_shared( + numPartitions, std::move(codec), getDefaultMemoryManager(), options, std::move(rssClient)); + } + default: + throw GlutenException("Unsupported partition writer type."); + } +} + +} // namespace + +class GpuVeloxShuffleWriterTestEnvironment : public ::testing::Environment { + public: + void SetUp() override { + VeloxShuffleWriterTestBase::setUpVeloxBackend(); + } + + void TearDown() override { + VeloxShuffleWriterTestBase::tearDownVeloxBackend(); + } +}; + +class GpuVeloxShuffleWriterTest : public ::testing::TestWithParam, + public VeloxShuffleWriterTestBase { + protected: + static std::shared_ptr createShuffleWriterOptions( + Partitioning partitioning, + int32_t splitBufferSize) { + std::shared_ptr options; + const auto& params = GetParam(); + switch (params.shuffleWriterType) { + case ShuffleWriterType::kGpuHashShuffle: { + auto hashOptions = std::make_shared(); + hashOptions->splitBufferSize = splitBufferSize; + options = hashOptions; + } break; + default: + throw GlutenException("Unreachable"); + } + options->partitioning = partitioning; + + return options; + } + + virtual std::shared_ptr defaultShuffleWriterOptions() = 0; + + std::shared_ptr createShuffleWriter( + uint32_t numPartitions, + std::shared_ptr shuffleWriterOptions = nullptr) { + if (shuffleWriterOptions == nullptr) { + shuffleWriterOptions = defaultShuffleWriterOptions(); + } + + const auto& params = GetParam(); + const auto partitionWriter = createPartitionWriter( + params.partitionWriterType, + numPartitions, + dataFile_, + localDirs_, + params.compressionType, + params.mergeBufferSize, + params.compressionThreshold, + params.enableDictionary); + + GLUTEN_ASSIGN_OR_THROW( + auto shuffleWriter, + VeloxShuffleWriter::create( + params.shuffleWriterType, numPartitions, partitionWriter, shuffleWriterOptions, getDefaultMemoryManager())); + + return shuffleWriter; + } + + void SetUp() override { + std::cout << "Running test with param: " << GetParam().toString() << std::endl; + VeloxShuffleWriterTestBase::setUpTestData(); + // Remove the UNKNOWN type because cudf does not support. + children1_.pop_back(); + children2_.pop_back(); + inputVector1_ = makeRowVector(children1_); + inputVector2_ = makeRowVector(children2_); + } + + void TearDown() override { + if (file_ != nullptr && !file_->closed()) { + GLUTEN_THROW_NOT_OK(file_->Close()); + } + } + + static void checkFileExists(const std::string& fileName) { + ASSERT_EQ(*arrow::internal::FileExists(*arrow::internal::PlatformFilename::FromString(fileName)), true); + } + + std::shared_ptr getArrowSchema(const facebook::velox::RowVectorPtr& rowVector) const { + return toArrowSchema(rowVector->type(), getDefaultMemoryManager()->getLeafMemoryPool().get()); + } + + void setReadableFile(const std::string& fileName) { + if (file_ != nullptr && !file_->closed()) { + GLUTEN_THROW_NOT_OK(file_->Close()); + } + GLUTEN_ASSIGN_OR_THROW(file_, arrow::io::ReadableFile::Open(fileName)) + } + + void getBufferColumnarBatches( + arrow::Compression::type compressionType, + const RowTypePtr& rowType, + std::vector>& bufferBatches, + std::shared_ptr in) { + const auto veloxCompressionType = arrowCompressionTypeToVelox(compressionType); + const auto schema = toArrowSchema(rowType, getDefaultMemoryManager()->getLeafMemoryPool().get()); + auto codec = createCompressionCodec(compressionType, CodecBackend::NONE); + + auto deserializerFactory = std::make_unique( + schema, + std::move(codec), + veloxCompressionType, + rowType, + kDefaultBatchSize, + kDefaultReadBufferSize, + GetParam().deserializerBufferSize, + getDefaultMemoryManager(), + GetParam().shuffleWriterType); + + const auto reader = std::make_shared(std::move(deserializerFactory)); + const auto iter = reader->read(std::make_shared(std::move(in))); + + while (iter->hasNext()) { + auto cb = std::dynamic_pointer_cast(iter->next()); + VELOX_CHECK_NOT_NULL(cb); + bufferBatches.emplace_back(cb); + } + + } + + void shuffleWriteReadMultiBlocks( + VeloxShuffleWriter& shuffleWriter, + int32_t expectPartitionLength, + const std::vector>& expectedVectors) { + ASSERT_NOT_OK(shuffleWriter.stop()); + checkFileExists(dataFile_); + const auto& lengths = shuffleWriter.partitionLengths(); + ASSERT_EQ(lengths.size(), expectPartitionLength); + + // Compute total size and verify file size + int64_t lengthSum = std::accumulate(lengths.begin(), lengths.end(), 0L); + + setReadableFile(dataFile_); + int64_t fileSize = *file_->GetSize(); + ASSERT_EQ(fileSize, lengthSum); + + int64_t offset = 0; + for (int32_t i = 0; i < expectPartitionLength; i++) { + if (expectedVectors[i].empty()) { + ASSERT_EQ(lengths[i], 0); + } else { + std::vector> deserializedVectors; + + // Compute byte range for this partition + int64_t start = offset; + int64_t size = lengths[i]; + offset += size; + + GLUTEN_ASSIGN_OR_THROW(auto in, arrow::io::RandomAccessFile::GetStream(file_, start, size)); + + getBufferColumnarBatches( + GetParam().compressionType, asRowType(expectedVectors[i][0]->type()), deserializedVectors, in); + + auto expectedVector = mergeRowVectors(expectedVectors[i]); + auto deserializedVector = mergeBufferColumnarBatches(deserializedVectors); + + facebook::velox::test::assertEqualVectors(expectedVector, deserializedVector); + } + } + + } + + void testShuffleRoundTrip( + VeloxShuffleWriter& shuffleWriter, + const std::vector>& batches, + int32_t expectPartitionLength, + const std::vector>& expectedVectors) { + for (const auto& batch : batches) { + ASSERT_NOT_OK(shuffleWriter.write(batch, ShuffleWriter::kMinMemLimit)); + } + shuffleWriteReadMultiBlocks(shuffleWriter, expectPartitionLength, expectedVectors); + } + + void testShuffleRoundTrip( + VeloxShuffleWriter& shuffleWriter, + std::vector inputs, + int32_t expectPartitionLength, + std::vector> expectedVectors) { + std::vector> batches; + for (const auto& input : inputs) { + batches.emplace_back(std::make_shared(input)); + } + testShuffleRoundTrip(shuffleWriter, batches, expectPartitionLength, expectedVectors); + } + + inline static TestAllocationListener* listener_{nullptr}; + + std::shared_ptr file_; +}; + +class GpuSinglePartitioningShuffleWriterTest : public GpuVeloxShuffleWriterTest { + protected: + std::shared_ptr defaultShuffleWriterOptions() override { + return createShuffleWriterOptions(Partitioning::kSingle, 10); + } +}; + +class GpuHashPartitioningShuffleWriterTest : public GpuVeloxShuffleWriterTest { + protected: + void SetUp() override { + GpuVeloxShuffleWriterTest::SetUp(); + + children1_.insert((children1_.begin()), makeFlatVector({1, 2, 2, 2, 2, 1, 1, 1, 2, 1})); + hashInputVector1_ = makeRowVector(children1_); + children2_.insert((children2_.begin()), makeFlatVector({2, 2})); + hashInputVector2_ = makeRowVector(children2_); + } + + std::shared_ptr defaultShuffleWriterOptions() override { + return createShuffleWriterOptions(Partitioning::kHash, 4); + } + + std::vector hashPartitionIds_{1, 2}; + + RowVectorPtr hashInputVector1_; + RowVectorPtr hashInputVector2_; +}; + +class GpuRangePartitioningShuffleWriterTest : public GpuVeloxShuffleWriterTest { + protected: + void SetUp() override { + GpuVeloxShuffleWriterTest::SetUp(); + + auto pid1 = makeRowVector({makeFlatVector({0, 1, 0, 1, 0, 1, 0, 1, 0, 1})}); + auto rangeVector1 = makeRowVector(inputVector1_->children()); + compositeBatch1_ = VeloxColumnarBatch::compose( + pool(), {std::make_shared(pid1), std::make_shared(rangeVector1)}); + + auto pid2 = makeRowVector({makeFlatVector({0, 1})}); + auto rangeVector2 = makeRowVector(inputVector2_->children()); + compositeBatch2_ = VeloxColumnarBatch::compose( + pool(), {std::make_shared(pid2), std::make_shared(rangeVector2)}); + } + + std::shared_ptr defaultShuffleWriterOptions() override { + return createShuffleWriterOptions(Partitioning::kRange, 4); + } + + std::shared_ptr compositeBatch1_; + std::shared_ptr compositeBatch2_; +}; + +class GpuRoundRobinPartitioningShuffleWriterTest : public GpuVeloxShuffleWriterTest { + protected: + std::shared_ptr defaultShuffleWriterOptions() override { + return createShuffleWriterOptions(Partitioning::kRoundRobin, 4096); + } +}; + +TEST_P(GpuSinglePartitioningShuffleWriterTest, single) { + if (GetParam().shuffleWriterType != ShuffleWriterType::kHashShuffle) { + return; + } + + // Split 1 RowVector. + { + auto shuffleWriter = createShuffleWriter(1); + testShuffleRoundTrip(*shuffleWriter, {inputVector1_}, 1, {{inputVector1_}}); + } + // Split > 1 RowVector. + { + auto shuffleWriter = createShuffleWriter(1); + testShuffleRoundTrip( + *shuffleWriter, + {inputVector1_, inputVector2_, inputVector1_}, + 1, + {{inputVector1_, inputVector2_, inputVector1_}}); + } + // Split null RowVector. + { + auto shuffleWriter = createShuffleWriter(1); + auto vector = makeRowVector({ + makeNullableFlatVector({std::nullopt}), + makeNullableFlatVector({std::nullopt}), + }); + testShuffleRoundTrip(*shuffleWriter, {vector}, 1, {{vector}}); + } +} + +TEST_P(GpuHashPartitioningShuffleWriterTest, hashPart1Vector) { + // Fixed-length input. + { + auto shuffleWriter = createShuffleWriter(2); + + // Remove the decimal type because cudf to velox is not supported. + std::vector data = { + makeNullableFlatVector({1, 2, 3, std::nullopt}), + makeNullableFlatVector({1, 2, 3, 4}), + makeFlatVector({1, 2, 3, 4}), + makeFlatVector({232, 34567235, 1212, 4567}), + 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))}; + + const auto vector = makeRowVector(data); + + const auto blocksPid0 = takeRows({vector}, {{1, 3}}); + const auto blocksPid1 = takeRows({vector}, {{0, 2}}); + + // Add partition id as the first column. + data.insert(data.begin(), makeFlatVector({1, 2, 1, 2})); + const auto input = makeRowVector(data); + + testShuffleRoundTrip(*shuffleWriter, {input}, 2, {blocksPid0, blocksPid1}); + } + + // Variable-length input. + { + auto shuffleWriter = createShuffleWriter(2); + + std::vector data = { + makeFlatVector({"nn", "", "fr", "juiu"}), + makeNullableFlatVector({1, 2, 3, std::nullopt}), + makeNullableFlatVector({std::nullopt, "de", "10 I'm not inline string", "de"})}; + + const auto vector = makeRowVector(data); + + const auto blocksPid0 = takeRows({vector}, {{0, 1, 3}}); + const auto blocksPid1 = takeRows({vector}, {{2}}); + + // Add partition id as the first column. + data.insert(data.begin(), makeFlatVector({2, 2, 1, 2})); + const auto input = makeRowVector(data); + + testShuffleRoundTrip(*shuffleWriter, {input}, 2, {blocksPid0, blocksPid1}); + } +} + +TEST_P(GpuHashPartitioningShuffleWriterTest, hashPart3Vectors) { + auto shuffleWriter = createShuffleWriter(2); + + auto blockPid2 = takeRows({inputVector1_, inputVector2_, inputVector1_}, {{1, 2, 3, 4, 8}, {0, 1}, {1, 2, 3, 4, 8}}); + auto blockPid1 = takeRows({inputVector1_, inputVector1_}, {{0, 5, 6, 7, 9}, {0, 5, 6, 7, 9}}); + + testShuffleRoundTrip( + *shuffleWriter, {hashInputVector1_, hashInputVector2_, hashInputVector1_}, 2, {blockPid2, blockPid1}); +} + +TEST_P(GpuHashPartitioningShuffleWriterTest, hashLargeVectors) { + const int32_t expectedMaxBatchSize = 8; + auto shuffleWriter = createShuffleWriter(2); + + // calculate maxBatchSize_ + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, hashInputVector1_)); + if (GetParam().shuffleWriterType == ShuffleWriterType::kHashShuffle) { + VELOX_CHECK_EQ(shuffleWriter->maxBatchSize(), expectedMaxBatchSize); + } + + auto blockPid2 = takeRows({inputVector1_, inputVector2_, inputVector1_}, {{1, 2, 3, 4, 8}, {}, {1, 2, 3, 4, 8}}); + auto blockPid1 = takeRows({inputVector1_, inputVector1_}, {{0, 5, 6, 7, 9}, {0, 5, 6, 7, 9}}); + + VELOX_CHECK(hashInputVector1_->size() > expectedMaxBatchSize); + testShuffleRoundTrip(*shuffleWriter, {hashInputVector2_, hashInputVector1_}, 2, {blockPid2, blockPid1}); +} + +TEST_P(GpuRangePartitioningShuffleWriterTest, range) { + auto shuffleWriter = createShuffleWriter(2); + + auto blockPid1 = takeRows({inputVector1_, inputVector2_, inputVector1_}, {{0, 2, 4, 6, 8}, {0}, {0, 2, 4, 6, 8}}); + auto blockPid2 = takeRows({inputVector1_, inputVector2_, inputVector1_}, {{1, 3, 5, 7, 9}, {1}, {1, 3, 5, 7, 9}}); + + testShuffleRoundTrip( + *shuffleWriter, {compositeBatch1_, compositeBatch2_, compositeBatch1_}, 2, {blockPid1, blockPid2}); +} + +TEST_P(GpuRoundRobinPartitioningShuffleWriterTest, roundRobin) { + auto shuffleWriter = createShuffleWriter(2); + + auto blockPid1 = takeRows({inputVector1_, inputVector2_, inputVector1_}, {{0, 2, 4, 6, 8}, {0}, {0, 2, 4, 6, 8}}); + auto blockPid2 = takeRows({inputVector1_, inputVector2_, inputVector1_}, {{1, 3, 5, 7, 9}, {1}, {1, 3, 5, 7, 9}}); + + testShuffleRoundTrip(*shuffleWriter, {inputVector1_, inputVector2_, inputVector1_}, 2, {blockPid1, blockPid2}); +} + +TEST_P(GpuRoundRobinPartitioningShuffleWriterTest, preAllocForceRealloc) { + if (GetParam().shuffleWriterType != ShuffleWriterType::kHashShuffle) { + return; + } + + auto shuffleWriterOptions = std::make_shared(); + shuffleWriterOptions->splitBufferSize = 4096; + shuffleWriterOptions->splitBufferReallocThreshold = 0; // Force re-alloc on buffer size changed. + auto shuffleWriter = createShuffleWriter(2, shuffleWriterOptions); + + // First spilt no null. + auto inputNoNull = inputVectorNoNull_; + + // Second split has null. Continue filling current partition buffers. + std::vector intHasNull = { + makeNullableFlatVector({std::nullopt, 1}), + makeNullableFlatVector({std::nullopt, -1}), + makeNullableFlatVector({std::nullopt, 100}), + makeNullableFlatVector({0, 1}), + makeNullableFlatVector({0, 0.142857}), + makeNullableFlatVector({false, true}), + makeNullableFlatVector({"", "alice"}), + makeNullableFlatVector({"alice", ""}), + }; + + auto inputHasNull = makeRowVector(intHasNull); + // Split first input no null. + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputNoNull)); + // Split second input, continue filling but update null. + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputHasNull)); + + // Split first input again. + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputNoNull)); + // Check when buffer is full, evict current buffers and reuse. + auto cachedPayloadSize = shuffleWriter->cachedPayloadSize(); + auto partitionBufferBeforeEvict = shuffleWriter->partitionBufferSize(); + int64_t evicted; + ASSERT_NOT_OK(shuffleWriter->reclaimFixedSize(cachedPayloadSize, &evicted)); + // Check only cached data being spilled. + ASSERT_EQ(evicted, cachedPayloadSize); + VELOX_CHECK_EQ(shuffleWriter->partitionBufferSize(), partitionBufferBeforeEvict); + + // Split more data with null. New buffer size is larger. + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_)); + + // Split more data with null. New buffer size is smaller. + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector2_)); + + // Split more data with null. New buffer size is larger and current data is preserved. + // Evict cached data first. + ASSERT_NOT_OK(shuffleWriter->reclaimFixedSize(shuffleWriter->cachedPayloadSize(), &evicted)); + // Set a large buffer size. + shuffleWriter->setPartitionBufferSize(100); + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_)); + // No data got evicted so the cached size is 0. + ASSERT_EQ(shuffleWriter->cachedPayloadSize(), 0); + + ASSERT_NOT_OK(shuffleWriter->stop()); +} + +TEST_P(GpuRoundRobinPartitioningShuffleWriterTest, preAllocForceReuse) { + if (GetParam().shuffleWriterType != ShuffleWriterType::kHashShuffle) { + return; + } + auto shuffleWriterOptions = std::make_shared(); + shuffleWriterOptions->splitBufferSize = 4096; + shuffleWriterOptions->splitBufferReallocThreshold = 1; // Force reuse on buffer size changed. + auto shuffleWriter = createShuffleWriter(2, shuffleWriterOptions); + + // First spilt no null. + auto inputNoNull = inputVectorNoNull_; + // Second split has null int, null string and non-null string, + auto inputFixedWidthHasNull = inputVector1_; + // Third split has null string. + std::vector stringHasNull = { + makeNullableFlatVector({0, 1}), + makeNullableFlatVector({0, -1}), + makeNullableFlatVector({0, 100}), + makeNullableFlatVector({0, 1}), + makeNullableFlatVector({0, 0.142857}), + makeNullableFlatVector({false, true}), + makeNullableFlatVector({std::nullopt, std::nullopt}), + makeNullableFlatVector({std::nullopt, std::nullopt}), + }; + auto inputStringHasNull = makeRowVector(stringHasNull); + + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputNoNull)); + // Split more data with null. Already filled + to be filled > buffer size, Buffer is resized larger. + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputFixedWidthHasNull)); + // Split more data with null. Already filled + to be filled > buffer size, newSize is smaller so buffer is not + // resized. + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputStringHasNull)); + + ASSERT_NOT_OK(shuffleWriter->stop()); +} + +TEST_P(GpuRoundRobinPartitioningShuffleWriterTest, spillVerifyResult) { + if (GetParam().shuffleWriterType != ShuffleWriterType::kHashShuffle) { + return; + } + + auto shuffleWriter = createShuffleWriter(2); + + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_)); + + // Clear buffers and evict payloads and cache. + for (auto pid : {0, 1}) { + ASSERT_NOT_OK(shuffleWriter->evictPartitionBuffers(pid, true)); + } + + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_)); + + // Evict all payloads and spill. + int64_t evicted; + auto cachedPayloadSize = shuffleWriter->cachedPayloadSize(); + auto partitionBufferSize = shuffleWriter->partitionBufferSize(); + ASSERT_NOT_OK(shuffleWriter->reclaimFixedSize(cachedPayloadSize + partitionBufferSize, &evicted)); + + ASSERT_EQ(evicted, cachedPayloadSize + partitionBufferSize); + + // No more cached payloads after spill. + ASSERT_EQ(shuffleWriter->cachedPayloadSize(), 0); + ASSERT_EQ(shuffleWriter->partitionBufferSize(), 0); + + ASSERT_NOT_OK(splitRowVector(*shuffleWriter, inputVector1_)); + + auto blockPid1 = + takeRows({inputVector1_, inputVector1_, inputVector1_}, {{0, 2, 4, 6, 8}, {0, 2, 4, 6, 8}, {0, 2, 4, 6, 8}}); + auto blockPid2 = + takeRows({inputVector1_, inputVector1_, inputVector1_}, {{1, 3, 5, 7, 9}, {1, 3, 5, 7, 9}, {1, 3, 5, 7, 9}}); + + // Stop and verify. + shuffleWriteReadMultiBlocks(*shuffleWriter, 2, {blockPid1, blockPid2}); +} + +INSTANTIATE_TEST_SUITE_P( + SinglePartitioningShuffleWriterGroup, + GpuSinglePartitioningShuffleWriterTest, + ::testing::ValuesIn(getTestParams())); + +INSTANTIATE_TEST_SUITE_P( + RoundRobinPartitioningShuffleWriterGroup, + GpuRoundRobinPartitioningShuffleWriterTest, + ::testing::ValuesIn(getTestParams())); + +INSTANTIATE_TEST_SUITE_P( + HashPartitioningShuffleWriterGroup, + GpuHashPartitioningShuffleWriterTest, + ::testing::ValuesIn(getTestParams())); + +INSTANTIATE_TEST_SUITE_P( + RangePartitioningShuffleWriterGroup, + GpuRangePartitioningShuffleWriterTest, + ::testing::ValuesIn(getTestParams())); + +} // namespace gluten + +int main(int argc, char** argv) { + testing::InitGoogleTest(&argc, argv); + ::testing::AddGlobalTestEnvironment(new gluten::GpuVeloxShuffleWriterTestEnvironment); + return RUN_ALL_TESTS(); +} diff --git a/cpp/velox/utils/GpuBufferBatchResizer.cc b/cpp/velox/utils/GpuBufferBatchResizer.cc new file mode 100644 index 00000000000..ad98fd47017 --- /dev/null +++ b/cpp/velox/utils/GpuBufferBatchResizer.cc @@ -0,0 +1,214 @@ +/* + * 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. + */ + +#include "GpuBufferBatchResizer.h" +#include "cudf/GpuLock.h" +#include "memory/GpuBufferColumnarBatch.h" +#include "utils/Timer.h" +#include "velox/experimental/cudf/exec/Utilities.h" +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" +#include "velox/experimental/cudf/vector/CudfVector.h" +#include "velox/vector/FlatVector.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace facebook::velox; + +namespace gluten { + +namespace { + +struct DispatchColumn { + rmm::cuda_stream_view stream; + rmm::device_async_resource_ref mr; + const std::vector>& buffers; + const int32_t numRows; + int32_t bufferIdx = 0; + + std::unique_ptr getMaskBuffer(const std::shared_ptr& buffer) { + if (buffer == nullptr || buffer->size() == 0) { + return std::make_unique(0, stream, mr); + } + + auto mask = std::make_unique(buffer->size(), stream, mr); + CUDF_CUDA_TRY( + cudaMemcpyAsync(mask->data(), buffer->data(), buffer->size(), cudaMemcpyHostToDevice, stream.value())); + return mask; + } + + // For timestamp, it is cudf cudf::type_id::TIMESTAMP_NANOSECONDS, Velox uses int128_t while cudf uses int64_t to + // represent it. + template ::NativeType> + std::unique_ptr readFlatColumn(cudf::type_id typeId) { + // === Step 1: get CPU buffers === + auto nulls = buffers[bufferIdx++]; + auto values = buffers[bufferIdx++]; + + // === Step 2: allocate GPU device buffers and copy === + rmm::device_buffer dataBuf(values->size(), stream); + CUDF_CUDA_TRY( + cudaMemcpyAsync(dataBuf.data(), values->data(), values->size(), cudaMemcpyHostToDevice, stream.value())); + + auto nullBuf = getMaskBuffer(nulls); + + // === Step 3: create cudf::column === + cudf::data_type cudfType{typeId}; + size_t nullCount = nulls == nullptr || nulls->size() == 0 + ? 0 + : cudf::null_count(reinterpret_cast(nulls->data()), 0, numRows, stream); + return std::make_unique(cudfType, numRows, std::move(dataBuf), std::move(*nullBuf), nullCount); + } + + /// We can optimize it in shuffle writer side, returns the offset buffer instead of length buffer. + /// Then we don't need to recover the offsetBuf by rawLengths, also change the merge strategic, update the merge + /// buffer offset from last offset. + std::unique_ptr getOffsetsColumn(const std::shared_ptr& offsets) { + VELOX_CHECK_GT(numRows, 0); + // --- 2. Copy offsets to GPU --- + rmm::device_buffer offsetBuf(offsets->size(), stream, mr); + CUDF_CUDA_TRY( + cudaMemcpyAsync(offsetBuf.data(), offsets->data(), offsets->size(), cudaMemcpyHostToDevice, stream.value())); + + // --- 3. Empty null mask (no nulls in offset column) --- + rmm::device_buffer nullBuf(0, stream, mr); + + // --- 4. Create cudf::column --- + return std::make_unique( + cudf::data_type{cudf::type_id::INT32}, + static_cast(numRows + 1), + std::move(offsetBuf), + std::move(nullBuf), + 0); // null_count = 0 + } + + std::unique_ptr readFlatColumnStringView(cudf::type_id /*typeId*/) { + auto nulls = buffers[bufferIdx++]; + auto offsets = buffers[bufferIdx++]; + auto valueBuffer = buffers[bufferIdx++]; + + if (numRows == 0) { + return make_empty_column(cudf::type_id::STRING); + } + + auto mask = getMaskBuffer(nulls); + + // === Step 3: create cudf::column === + size_t nullCount = nulls == nullptr || nulls->size() == 0 + ? 0 + : cudf::null_count(reinterpret_cast(nulls->data()), 0, numRows, stream); + + auto offsetColumn = getOffsetsColumn(offsets); + + rmm::device_buffer chars(valueBuffer->size(), stream, mr); + CUDF_CUDA_TRY(cudaMemcpyAsync( + chars.data(), valueBuffer->data_as(), chars.size(), cudaMemcpyDefault, stream.value())); + return cudf::make_strings_column( + numRows, std::move(offsetColumn), std::move(chars), nullCount, std::move(*mask.release())); + } +}; + +template <> +std::unique_ptr DispatchColumn::readFlatColumn(cudf::type_id typeId) { + return readFlatColumnStringView(typeId); +} + +template <> +std::unique_ptr DispatchColumn::readFlatColumn(cudf::type_id typeId) { + return readFlatColumnStringView(typeId); +} + +std::shared_ptr makeCudfTable( + RowTypePtr type, + int32_t numRows, + const std::vector>& buffers, + memory::MemoryPool* pool) { + std::vector> cudfColumns; + cudfColumns.reserve(type->size()); + + auto stream = cudf_velox::cudfGlobalStreamPool().get_stream(); + DispatchColumn dispatch{stream, cudf::get_current_device_resource_ref(), buffers, numRows}; + for (const auto& colType : type->children()) { + auto res = VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH( + dispatch.readFlatColumn, colType->kind(), cudf_velox::veloxToCudfTypeId(colType)); + cudfColumns.emplace_back(std::move(res)); + } + auto cudfTable = std::make_unique(std::move(cudfColumns)); + stream.synchronize(); + return std::make_shared( + std::make_shared(pool, type, numRows, std::move(cudfTable), stream)); +} + +} // namespace + +GpuBufferBatchResizer::GpuBufferBatchResizer( + arrow::MemoryPool* arrowPool, + facebook::velox::memory::MemoryPool* pool, + int32_t minOutputBatchSize, + std::unique_ptr in) + : arrowPool_(arrowPool), + pool_(pool), + minOutputBatchSize_(minOutputBatchSize), + in_(std::move(in)) { + VELOX_CHECK_GT(minOutputBatchSize_, 0, "minOutputBatchSize should be larger than 0"); +} + +std::shared_ptr GpuBufferBatchResizer::next() { + std::vector> cachedBatches; + int32_t cachedRows = 0; + while (cachedRows < minOutputBatchSize_) { + auto nextCb = in_->next(); + if (!nextCb) { + // No more input. + break; + } + + auto nextBatch = std::dynamic_pointer_cast(nextCb); + VELOX_CHECK_NOT_NULL(nextBatch); + if (nextBatch->numRows() == 0) { + continue; + } + + cachedRows += nextBatch->numRows(); + cachedBatches.push_back(std::move(nextBatch)); + } + if (cachedRows == 0) { + return nullptr; + } + + // Compose all cached batches into one + auto batch = GpuBufferColumnarBatch::compose(arrowPool_, cachedBatches, cachedRows); + + lockGpu(); + + return makeCudfTable(batch->getRowType(), batch->numRows(), batch->buffers(), pool_); +} + +int64_t GpuBufferBatchResizer::spillFixedSize(int64_t size) { + return in_->spillFixedSize(size); +} + +} // namespace gluten diff --git a/cpp/velox/utils/GpuBufferBatchResizer.h b/cpp/velox/utils/GpuBufferBatchResizer.h new file mode 100644 index 00000000000..b3ed0866293 --- /dev/null +++ b/cpp/velox/utils/GpuBufferBatchResizer.h @@ -0,0 +1,44 @@ +/* + * 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. + */ + +#include "memory/ColumnarBatchIterator.h" +#include "memory/VeloxColumnarBatch.h" +#include "utils/Exception.h" +#include "velox/common/memory/MemoryPool.h" + +namespace gluten { + +class GpuBufferBatchResizer : public ColumnarBatchIterator { + public: + GpuBufferBatchResizer( + arrow::MemoryPool* arrowPool, + facebook::velox::memory::MemoryPool* pool, + int32_t minOutputBatchSize, + std::unique_ptr in); + + std::shared_ptr next() override; + + int64_t spillFixedSize(int64_t size) override; + + private: + arrow::MemoryPool* arrowPool_; + facebook::velox::memory::MemoryPool* pool_; + const int32_t minOutputBatchSize_; + std::unique_ptr in_; +}; + +} // namespace gluten diff --git a/dev/docker/cudf/Dockerfile b/dev/docker/cudf/Dockerfile index 5d2255fcb67..de92527a1e4 100644 --- a/dev/docker/cudf/Dockerfile +++ b/dev/docker/cudf/Dockerfile @@ -28,7 +28,7 @@ ENV CUDA_ARCHITECTURES=70 WORKDIR /opt/gluten -RUN bash ./dev/buildbundle-veloxbe.sh --run_setup_script=OFF --build_arrow=ON --spark_version=3.4 --enable_gpu=ON && rm -rf /opt/gluten +RUN bash ./dev/buildbundle-veloxbe.sh --run_setup_script=OFF --build_arrow=ON --spark_version=3.4 --build_tests=ON --build_benchmarks=ON --enable_gpu=ON && rm -rf /opt/gluten # You can try the data in folder backends-velox/src/test/resources/tpch-data-parquet diff --git a/docs/velox-configuration.md b/docs/velox-configuration.md index fd7cd5cdcaf..cacb8cf4805 100644 --- a/docs/velox-configuration.md +++ b/docs/velox-configuration.md @@ -22,6 +22,7 @@ nav_order: 16 | spark.gluten.sql.columnar.backend.velox.cacheEnabled | false | Enable Velox cache, default off. It's recommended to enablesoft-affinity as well when enable velox cache. | | spark.gluten.sql.columnar.backend.velox.cachePrefetchMinPct | 0 | Set prefetch cache min pct for velox file scan | | spark.gluten.sql.columnar.backend.velox.checkUsageLeak | true | Enable check memory usage leak. | +| spark.gluten.sql.columnar.backend.velox.cudf.batchSize | 2147483647 | Cudf input batch size after shuffle reader | | spark.gluten.sql.columnar.backend.velox.cudf.enableTableScan | false | Enable cudf table scan | | spark.gluten.sql.columnar.backend.velox.cudf.enableValidation | true | Heuristics you can apply to validate a cuDF/GPU plan and only offload when the entire stage can be fully and profitably executed on GPU | | spark.gluten.sql.columnar.backend.velox.cudf.memoryPercent | 50 | The initial percent of GPU memory to allocate for memory resource for one thread. | diff --git a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleWriterJniWrapper.java b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleWriterJniWrapper.java index 8e8b44347ff..a389da68603 100644 --- a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleWriterJniWrapper.java +++ b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleWriterJniWrapper.java @@ -63,6 +63,14 @@ public native long createRssSortShuffleWriter( String codec, long partitionWriterHandle); + public native long createGpuHashShuffleWriter( + int numPartitions, + String partitioningName, + int startPartitionId, + int splitBufferSize, + double splitBufferReallocThreshold, + long partitionWriterHandle); + /** * Reclaim memory from the shuffle writer instance. It will first try to shrink allocated memory, * and may trigger a spill if needed. 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 bdb9f92de44..bccd4eedf8f 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 @@ -16,11 +16,9 @@ */ package org.apache.spark.sql.execution -import org.apache.gluten.backendsapi.BackendsApiManager import org.apache.gluten.sql.shims.SparkShimLoader import org.apache.spark.internal.Logging -import org.apache.spark.serializer.Serializer import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.execution.exchange._ @@ -33,16 +31,10 @@ case class ColumnarShuffleExchangeExec( advisoryPartitionSize: Option[Long] = None) extends ColumnarShuffleExchangeExecBase(outputPartitioning, child, projectOutputAttributes) { - // super.stringArgs ++ Iterator(output.map(o => s"${o}#${o.dataType.simpleString}")) - val serializer: Serializer = BackendsApiManager.getSparkPlanExecApiInstance - .createColumnarBatchSerializer(schema, metrics, shuffleWriterType) - override def nodeName: String = "ColumnarExchange" protected def withNewChildInternal(newChild: SparkPlan): ColumnarShuffleExchangeExec = copy(child = newChild) - - override def getSerializer: Serializer = serializer } object ColumnarShuffleExchangeExec extends Logging { diff --git a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExecBase.scala b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExecBase.scala index af8ba753930..44f37e4ffbc 100644 --- a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExecBase.scala +++ b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExecBase.scala @@ -48,8 +48,11 @@ abstract class ColumnarShuffleExchangeExecBase( private[sql] lazy val readMetrics = SQLColumnarShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) - val shuffleWriterType: ShuffleWriterType = - BackendsApiManager.getSparkPlanExecApiInstance.getShuffleWriterType(outputPartitioning, output) + lazy val shuffleWriterType: ShuffleWriterType = getShuffleWriterType + + // super.stringArgs ++ Iterator(output.map(o => s"${o}#${o.dataType.simpleString}")) + lazy val serializer: Serializer = BackendsApiManager.getSparkPlanExecApiInstance + .createColumnarBatchSerializer(schema, metrics, shuffleWriterType) // Note: "metrics" is made transient to avoid sending driver-side metrics to tasks. @transient override lazy val metrics = @@ -81,7 +84,7 @@ abstract class ColumnarShuffleExchangeExecBase( child.output, projectOutputAttributes, outputPartitioning, - getSerializer, + serializer, writeMetrics, metrics, shuffleWriterType) @@ -110,7 +113,8 @@ abstract class ColumnarShuffleExchangeExecBase( Statistics(dataSize, Some(rowCount)) } - def getSerializer: Serializer + def getShuffleWriterType: ShuffleWriterType = + BackendsApiManager.getSparkPlanExecApiInstance.getShuffleWriterType(outputPartitioning, output) // Required for Spark 4.0 to implement a trait method. // The "override" keyword is omitted to maintain compatibility with earlier Spark versions. diff --git a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/GPUColumnarShuffleExchangeExec.scala b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/GPUColumnarShuffleExchangeExec.scala index 79a66b00b77..3564d6fa56b 100644 --- a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/GPUColumnarShuffleExchangeExec.scala +++ b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/GPUColumnarShuffleExchangeExec.scala @@ -17,13 +17,11 @@ package org.apache.spark.sql.execution import org.apache.gluten.backendsapi.BackendsApiManager -import org.apache.gluten.config.GpuHashShuffleWriterType -import org.apache.gluten.config.HashShuffleWriterType +import org.apache.gluten.config.{GpuHashShuffleWriterType, HashShuffleWriterType, ShuffleWriterType} import org.apache.gluten.execution.ValidationResult import org.apache.gluten.sql.shims.SparkShimLoader import org.apache.spark.internal.Logging -import org.apache.spark.serializer.Serializer import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.execution.exchange._ @@ -42,21 +40,21 @@ case class GPUColumnarShuffleExchangeExec( if (!validation.ok()) { return validation } + val shuffleWriterType = BackendsApiManager.getSparkPlanExecApiInstance.getShuffleWriterType( + outputPartitioning, + output) if (shuffleWriterType != HashShuffleWriterType) { return ValidationResult.failed("Only support hash partitioning") } ValidationResult.succeeded } - // super.stringArgs ++ Iterator(output.map(o => s"${o}#${o.dataType.simpleString}")) - val serializer: Serializer = BackendsApiManager.getSparkPlanExecApiInstance - .createColumnarBatchSerializer(schema, metrics, GpuHashShuffleWriterType) override def nodeName: String = "CudfColumnarExchange" + override def getShuffleWriterType: ShuffleWriterType = GpuHashShuffleWriterType + protected def withNewChildInternal(newChild: SparkPlan): GPUColumnarShuffleExchangeExec = copy(child = newChild) - - override def getSerializer: Serializer = serializer } object GPUColumnarShuffleExchangeExec extends Logging { From 51099e9f1e4bc87725832eb5aa8ac2c69a7d2222 Mon Sep 17 00:00:00 2001 From: Chengcheng Jin Date: Fri, 21 Nov 2025 13:58:40 +0000 Subject: [PATCH 2/3] fix --- cpp/velox/tests/VeloxGpuShuffleWriterTest.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc index 3405548c517..c53c24f24af 100644 --- a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc +++ b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc @@ -20,6 +20,7 @@ #include "config/GlutenConfig.h" #include "shuffle/VeloxHashShuffleWriter.h" +#include "shuffle/VeloxGpuShuffleWriter.h" #include "tests/VeloxShuffleWriterTestBase.h" #include "tests/utils/TestAllocationListener.h" #include "tests/utils/TestStreamReader.h" From 71f7e487b36f325571134f6fd48181b25777bf80 Mon Sep 17 00:00:00 2001 From: Chengcheng Jin Date: Fri, 21 Nov 2025 14:01:17 +0000 Subject: [PATCH 3/3] fix --- cpp/velox/tests/VeloxGpuShuffleWriterTest.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc index c53c24f24af..364d31e180c 100644 --- a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc +++ b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc @@ -19,8 +19,8 @@ #include #include "config/GlutenConfig.h" -#include "shuffle/VeloxHashShuffleWriter.h" #include "shuffle/VeloxGpuShuffleWriter.h" +#include "shuffle/VeloxHashShuffleWriter.h" #include "tests/VeloxShuffleWriterTestBase.h" #include "tests/utils/TestAllocationListener.h" #include "tests/utils/TestStreamReader.h" @@ -221,7 +221,7 @@ class GpuVeloxShuffleWriterTest : public ::testing::TestWithParam(); + auto hashOptions = std::make_shared(); hashOptions->splitBufferSize = splitBufferSize; options = hashOptions; } break;