diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/HashJoinExecTransformer.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/HashJoinExecTransformer.scala index 2c1c976e914..705f938a201 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/execution/HashJoinExecTransformer.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/execution/HashJoinExecTransformer.scala @@ -221,7 +221,7 @@ case class BroadcastHashJoinExecTransformer( } else { logInfo(s"Using executor-side broadcast hash table build for $buildBroadcastTableId") } - VeloxBroadcastBuildSideRDD(sparkContext, broadcast, context) + VeloxBroadcastBuildSideRDD(sparkContext, broadcast, context, cudfEnabled = offloadCuda) case unsafe: UnsafeColumnarBuildSideRelation => joinParamsForMetrics.foreach(_.usesDriverSideSerializedHashTable = false) @@ -235,7 +235,7 @@ case class BroadcastHashJoinExecTransformer( } else { logInfo(s"Using executor-side broadcast hash table build for $buildBroadcastTableId") } - VeloxBroadcastBuildSideRDD(sparkContext, broadcast, context) + VeloxBroadcastBuildSideRDD(sparkContext, broadcast, context, cudfEnabled = offloadCuda) case other => joinParamsForMetrics.foreach(_.usesDriverSideSerializedHashTable = false) @@ -243,7 +243,7 @@ case class BroadcastHashJoinExecTransformer( logWarning( s"Unknown broadcast relation type: ${other.getClass.getName}, " + "using executor-side build") - VeloxBroadcastBuildSideRDD(sparkContext, broadcast, context) + VeloxBroadcastBuildSideRDD(sparkContext, broadcast, context, cudfEnabled = offloadCuda) } // FIXME: Do we have to make build side a RDD? diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideRDD.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideRDD.scala index 074eae17dfc..1c5fcbb70e7 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideRDD.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideRDD.scala @@ -29,7 +29,8 @@ case class VeloxBroadcastBuildSideRDD( @transient private val sc: SparkContext, broadcasted: broadcast.Broadcast[BuildSideRelation], broadcastContext: BroadcastHashJoinContext, - isBNL: Boolean = false) + isBNL: Boolean = false, + cudfEnabled: Boolean = false) extends BroadcastBuildSideRDD(sc, broadcasted) { override def genBroadcastBuildSideIterator(): Iterator[ColumnarBatch] = { @@ -48,8 +49,17 @@ case class VeloxBroadcastBuildSideRDD( // reusable table and a CPU-fallback join builds from this stream as usual. val output = if (isBNL || !offload || GlutenConfig.get.enableColumnarCudf) { val relation = broadcasted.value.asReadOnlyCopy() + // cudfEnabled is the consuming stage's own tag (TransformSupport#offloadCuda). + val batches = relation match { + case columnar: ColumnarBuildSideRelation => + columnar.deserialized(cudfEnabled) + case unsafe: UnsafeColumnarBuildSideRelation => + unsafe.deserialized(cudfEnabled) + case other => + other.deserialized + } Iterators - .wrap(relation.deserialized) + .wrap(batches) .recyclePayload(batch => batch.close()) .create() } else { diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastNestedLoopJoinExecTransformer.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastNestedLoopJoinExecTransformer.scala index 6e0aaa27c6d..7ba3edbd8ee 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastNestedLoopJoinExecTransformer.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastNestedLoopJoinExecTransformer.scala @@ -45,7 +45,13 @@ case class VeloxBroadcastNestedLoopJoinExecTransformer( override def columnarInputRDDs: Seq[RDD[ColumnarBatch]] = { val streamedRDD = getColumnarInputRDDs(streamedPlan) val broadcast = buildPlan.executeBroadcast[BuildSideRelation]() - val broadcastRDD = VeloxBroadcastBuildSideRDD(sparkContext, broadcast, null, true) + val broadcastRDD = + VeloxBroadcastBuildSideRDD( + sparkContext, + broadcast, + null, + isBNL = true, + cudfEnabled = offloadCuda) // FIXME: Do we have to make build side a RDD? streamedRDD :+ broadcastRDD } diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala b/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala index 7374cf0b04d..dc1be02d740 100644 --- a/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala +++ b/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.execution import org.apache.gluten.backendsapi.BackendsApiManager import org.apache.gluten.columnarbatch.ColumnarBatches +import org.apache.gluten.config.GlutenConfig import org.apache.gluten.execution.BroadcastHashJoinContext import org.apache.gluten.expression.ConverterUtils import org.apache.gluten.iterator.Iterators @@ -40,6 +41,8 @@ import org.apache.spark.util.KnownSizeEstimation import org.apache.arrow.c.ArrowSchema +import java.util.Collections + import scala.collection.JavaConverters._ import scala.collection.JavaConverters.asScalaIteratorConverter import scala.collection.mutable.{ArrayBuffer, Map} @@ -110,9 +113,20 @@ case class ColumnarBuildSideRelation( } } - override def deserialized: Iterator[ColumnarBatch] = { + /** Host-resident deserialization, for CPU consumers. */ + override def deserialized: Iterator[ColumnarBatch] = deserialized(cudfEnabled = false) + + /** + * Residency follows the consuming stage: a cuDF-offloaded stage sources from CudfValueStream and + * needs device batches, a non-offloaded stage from RowVectorStream and needs host batches. + */ + def deserialized(cudfEnabled: Boolean): Iterator[ColumnarBatch] = { val runtime = - Runtimes.contextInstance(BackendsApiManager.getBackendName, "BuildSideRelation#deserialized") + Runtimes.contextInstance( + BackendsApiManager.getBackendName, + "BuildSideRelation#deserialized", + Collections.singletonMap(GlutenConfig.COLUMNAR_CUDF_ENABLED.key, cudfEnabled.toString) + ) val jniWrapper = ColumnarBatchSerializerJniWrapper.create(runtime) val serializeHandle: Long = { val allocator = ArrowBufferAllocators.contextInstance() diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala index d0129e887cf..1c2adbc57f2 100644 --- a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala +++ b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.execution.unsafe import org.apache.gluten.backendsapi.BackendsApiManager import org.apache.gluten.columnarbatch.ColumnarBatches +import org.apache.gluten.config.GlutenConfig import org.apache.gluten.execution.BroadcastHashJoinContext import org.apache.gluten.expression.ConverterUtils import org.apache.gluten.iterator.Iterators @@ -45,6 +46,7 @@ import com.esotericsoftware.kryo.io.{Input, Output} import org.apache.arrow.c.ArrowSchema import java.io.{Externalizable, ObjectInput, ObjectOutput} +import java.util.Collections import scala.collection.JavaConverters._ import scala.collection.JavaConverters.asScalaIteratorConverter @@ -372,11 +374,20 @@ class UnsafeColumnarBuildSideRelation( } } - override def deserialized: Iterator[ColumnarBatch] = { + /** Host-resident deserialization, for CPU consumers. */ + override def deserialized: Iterator[ColumnarBatch] = deserialized(cudfEnabled = false) + + /** + * Residency follows the consuming stage: a cuDF-offloaded stage sources from CudfValueStream and + * needs device batches, a non-offloaded stage from RowVectorStream and needs host batches. + */ + def deserialized(cudfEnabled: Boolean): Iterator[ColumnarBatch] = { val runtime = Runtimes.contextInstance( BackendsApiManager.getBackendName, - "UnsafeBuildSideRelation#deserialize") + "UnsafeBuildSideRelation#deserialize", + Collections.singletonMap(GlutenConfig.COLUMNAR_CUDF_ENABLED.key, cudfEnabled.toString) + ) val jniWrapper = ColumnarBatchSerializerJniWrapper.create(runtime) val serializerHandle: Long = { val allocator = ArrowBufferAllocators.contextInstance() diff --git a/backends-velox/src/test/scala/org/apache/gluten/execution/CudfBroadcastJoinSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/execution/CudfBroadcastJoinSuite.scala index 62a9a64e3be..b8c74bfe138 100644 --- a/backends-velox/src/test/scala/org/apache/gluten/execution/CudfBroadcastJoinSuite.scala +++ b/backends-velox/src/test/scala/org/apache/gluten/execution/CudfBroadcastJoinSuite.scala @@ -21,7 +21,7 @@ import org.apache.gluten.tags.CudfTest import org.apache.spark.SparkConf /** - * Regression tests for GLUTEN-12471: broadcast hash joins on the cuDF (GPU) backend silently + * Regression tests for GLUTEN-12812: broadcast hash joins on the cuDF (GPU) backend silently * returned empty results because CudfHashJoin built its hash table from the empty build-side * iterator instead of the prebuilt CPU table. * @@ -53,7 +53,7 @@ class CudfBroadcastJoinSuite extends VeloxWholeStageTransformerSuite { createTPCHNotNullTables() } - test("GLUTEN-12471: cuDF broadcast hash join returns non-empty, correct results") { + test("GLUTEN-12812: cuDF broadcast hash join returns non-empty, correct results") { val query = """ |SELECT l.l_orderkey, o.o_orderdate, l.l_extendedprice @@ -72,7 +72,28 @@ class CudfBroadcastJoinSuite extends VeloxWholeStageTransformerSuite { // demoted or fallen back. val bhj = collect(plan) { case j: BroadcastHashJoinExecTransformer => j } assert(bhj.nonEmpty, s"expected an offloaded broadcast hash join, got:\n$plan") - assert(df.count() > 0, "broadcast join must not return empty results (GLUTEN-12471)") + assert(df.count() > 0, "broadcast join must not return empty results (GLUTEN-12812)") + } + } + + test("GLUTEN-12838: broadcast build side follows the consuming stage's cuDF tag") { + // We need a broadcast that lands in a CPU stage. NOT IN gives us one for free: it + // becomes a null-aware anti join, which Spark can only run as a broadcast, and the + // stage reading it has a table scan, so cuDF never claims it. + val query = + """ + |SELECT p_partkey, p_brand + |FROM part + |WHERE p_partkey NOT IN ( + | SELECT ps_partkey FROM partsupp WHERE ps_availqty < 10 + |) + |""".stripMargin + + runQueryAndCompare(query) { + df => + val plan = df.queryExecution.executedPlan + val bhj = collect(plan) { case j: BroadcastHashJoinExecTransformer => j } + assert(bhj.nonEmpty, s"expected a broadcast hash join for the anti join, got:\n$plan") } }