Add retry for CrossJoinIterator and ConditionalNestedLoopJoinIterator

sql-plugin/src/main/scala/com/nvidia/spark/rapids/AbstractGpuJoinIterator.scala

@@ -198,7 +198,7 @@ abstract class SplittableJoinIterator(
   // If the join explodes this holds batches from the stream side split into smaller pieces.
   private val pendingSplits = scala.collection.mutable.Queue[LazySpillableColumnarBatch]()
 
-  protected def computeNumJoinRows(cb: ColumnarBatch): Long
+  protected def computeNumJoinRows(cb: LazySpillableColumnarBatch): Long
 
   /**
    * Create a join gatherer.
@@ -225,7 +225,7 @@ abstract class SplittableJoinIterator(
       }
       opTime.ns {
         withResource(scb) { scb =>
-          val numJoinRows = computeNumJoinRows(scb.getBatch)
+          val numJoinRows = computeNumJoinRows(scb)
 
           // We want the gather maps size to be around the target size. There are two gather maps
           // that are made up of ints, so compute how many rows on the stream side will produce the

sql-plugin/src/main/scala/com/nvidia/spark/rapids/JoinGatherer.scala

@@ -19,6 +19,7 @@ package com.nvidia.spark.rapids
 import ai.rapids.cudf.{ColumnVector, ColumnView, DeviceMemoryBuffer, DType, GatherMap, NvtxColor, NvtxRange, OrderByArg, OutOfBoundsPolicy, Scalar, Table}
 import com.nvidia.spark.Retryable
 import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
+import com.nvidia.spark.rapids.RmmRapidsRetryIterator.withRetryNoSplit
 
 import org.apache.spark.TaskContext
 import org.apache.spark.sql.types._
@@ -419,7 +420,7 @@ abstract class BaseCrossJoinGatherMap(leftCount: Int, rightCount: Int)
     extends LazySpillableGatherMap {
   override val getRowCount: Long = leftCount.toLong * rightCount.toLong
 
-  override def toColumnView(startRow: Long, numRows: Int): ColumnView = {
+  override def toColumnView(startRow: Long, numRows: Int): ColumnView = withRetryNoSplit {
     withResource(GpuScalar.from(startRow, LongType)) { startScalar =>
       withResource(ai.rapids.cudf.ColumnVector.sequence(startScalar, numRows)) { rowNum =>
         compute(rowNum)

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExecBase.scala

@@ -178,20 +178,26 @@ class ConditionalNestedLoopJoinIterator(
     }
   }
 
-  override def computeNumJoinRows(cb: ColumnarBatch): Long = {
-    withResource(GpuColumnVector.from(builtBatch.getBatch)) { builtTable =>
-      withResource(GpuColumnVector.from(cb)) { streamTable =>
-        val (left, right) = buildSide match {
-          case GpuBuildLeft => (builtTable, streamTable)
-          case GpuBuildRight => (streamTable, builtTable)
-        }
-        joinType match {
-          case _: InnerLike =>left.conditionalInnerJoinRowCount(right, condition)
-          case LeftOuter => left.conditionalLeftJoinRowCount(right, condition)
-          case RightOuter => right.conditionalLeftJoinRowCount(left, condition)
-          case LeftSemi => left.conditionalLeftSemiJoinRowCount(right, condition)
-          case LeftAnti => left.conditionalLeftAntiJoinRowCount(right, condition)
-          case _ => throw new IllegalStateException(s"Unsupported join type $joinType")
+  override def computeNumJoinRows(scb: LazySpillableColumnarBatch): Long = {
+    scb.checkpoint()
+    builtBatch.checkpoint()
+    withRetryNoSplit {
+      withRestoreOnRetry(Seq(builtBatch, scb)) {
+        withResource(GpuColumnVector.from(builtBatch.getBatch)) { builtTable =>
+          withResource(GpuColumnVector.from(scb.getBatch)) { streamTable =>
+            val (left, right) = buildSide match {
+              case GpuBuildLeft => (builtTable, streamTable)
+              case GpuBuildRight => (streamTable, builtTable)
+            }
+            joinType match {
+              case _: InnerLike => left.conditionalInnerJoinRowCount(right, condition)
+              case LeftOuter => left.conditionalLeftJoinRowCount(right, condition)
+              case RightOuter => right.conditionalLeftJoinRowCount(left, condition)
+              case LeftSemi => left.conditionalLeftSemiJoinRowCount(right, condition)
+              case LeftAnti => left.conditionalLeftAntiJoinRowCount(right, condition)
+              case _ => throw new IllegalStateException(s"Unsupported join type $joinType")
+            }
+          }
         }
       }
     }

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuHashJoin.scala

@@ -292,14 +292,14 @@ abstract class BaseHashJoinIterator(
       1.0
   }
 
-  override def computeNumJoinRows(cb: ColumnarBatch): Long = {
+  override def computeNumJoinRows(cb: LazySpillableColumnarBatch): Long = {
     // TODO: Replace this estimate with exact join row counts using the corresponding cudf APIs
     //       being added in https://github.com/rapidsai/cudf/issues/9053.
     joinType match {
       // Full Outer join is implemented via LeftOuter/RightOuter, so use same estimate.
       case _: InnerLike | LeftOuter | RightOuter | FullOuter =>
-        Math.ceil(cb.numRows() * streamMagnificationFactor).toLong
-      case _ => cb.numRows()
+        Math.ceil(cb.numRows * streamMagnificationFactor).toLong
+      case _ => cb.numRows
     }
   }