Add retry support to RowToColumnarIterator

sql-plugin/src/main/java/com/nvidia/spark/rapids/GpuColumnVector.java

@@ -29,6 +29,7 @@
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
+import java.util.function.Function;
 
 /**
  * A GPU accelerated version of the Spark ColumnVector.
@@ -134,16 +135,18 @@ public static abstract class GpuColumnarBatchBuilderBase implements AutoCloseabl
 
     public abstract void copyColumnar(ColumnVector cv, int colNum, int rows);
 
-    protected abstract ColumnVector buildAndPutOnDevice(int builderIndex);
-    protected abstract int buildersLength();
+    protected abstract ai.rapids.cudf.ColumnVector buildAndPutOnDevice(int builderIndex);
 
     public ColumnarBatch build(int rows) {
-      int buildersLen = buildersLength();
-      ColumnVector[] vectors = new ColumnVector[buildersLen];
+      return build(rows, this::buildAndPutOnDevice);
+    }
+
+    protected ColumnarBatch build(int rows, Function<Integer, ai.rapids.cudf.ColumnVector> col) {
+      ColumnVector[] vectors = new ColumnVector[fields.length];
       boolean success = false;
       try {
-        for (int i = 0; i < buildersLen; i++) {
-          vectors[i] = buildAndPutOnDevice(i);
+        for (int i = 0; i < fields.length; i++) {
+          vectors[i] = new GpuColumnVector(fields[i].dataType(), col.apply(i));
         }
         ColumnarBatch ret = new ColumnarBatch(vectors, rows);
         success = true;
@@ -191,17 +194,11 @@ public GpuArrowColumnarBatchBuilder(StructType schema) {
     }
 
     @Override
-    protected int buildersLength() {
-      return builders.length;
-    }
-
-    @Override
-    protected ColumnVector buildAndPutOnDevice(int builderIndex) {
+    protected ai.rapids.cudf.ColumnVector buildAndPutOnDevice(int builderIndex) {
       ai.rapids.cudf.ColumnVector cv = builders[builderIndex].buildAndPutOnDevice();
-      GpuColumnVector gcv = new GpuColumnVector(fields[builderIndex].dataType(), cv);
       referenceHolders[builderIndex].releaseReferences();
       builders[builderIndex] = null;
-      return gcv;
+      return cv;
     }
 
     @Override
@@ -230,6 +227,7 @@ public void close() {
 
   public static final class GpuColumnarBatchBuilder extends GpuColumnarBatchBuilderBase {
     private final ai.rapids.cudf.HostColumnVector.ColumnBuilder[] builders;
+    private ai.rapids.cudf.HostColumnVector[] hostColumns;
 
     /**
      * A collection of builders for building up columnar data.
@@ -270,16 +268,10 @@ public ai.rapids.cudf.HostColumnVector.ColumnBuilder builder(int i) {
     }
 
     @Override
-    protected int buildersLength() {
-      return builders.length;
-    }
-
-    @Override
-    protected ColumnVector buildAndPutOnDevice(int builderIndex) {
+    protected ai.rapids.cudf.ColumnVector buildAndPutOnDevice(int builderIndex) {
       ai.rapids.cudf.ColumnVector cv = builders[builderIndex].buildAndPutOnDevice();
-      GpuColumnVector gcv = new GpuColumnVector(fields[builderIndex].dataType(), cv);
       builders[builderIndex] = null;
-      return gcv;
+      return cv;
     }
 
     public HostColumnVector[] buildHostColumns() {
@@ -303,12 +295,40 @@ public HostColumnVector[] buildHostColumns() {
       }
     }
 
+    /**
+     * Build a columnar batch without releasing the holding data on host.
+     * It is safe to call this multiple times, and data will be released
+     * after a call to `close`.
+     */
+    public ColumnarBatch tryBuild(int rows) {
+      if (hostColumns == null) {
+        hostColumns = buildHostColumns();
+      }
+      return build(rows, i -> hostColumns[i].copyToDevice());
+    }
+
     @Override
     public void close() {
       for (ai.rapids.cudf.HostColumnVector.ColumnBuilder b: builders) {
         if (b != null) {
-          b.close();
+          try {
+            b.close();
+          } catch (Throwable e) {
+            /* ignore the exception */
+          }
+        }
+      }
+      if (hostColumns != null) {
+        for (ai.rapids.cudf.HostColumnVector hcv: hostColumns) {
+          if (hcv != null) {
+            try {
+              hcv.close();
+            } catch (Throwable e) {
+              /* ignore the exception */
+            }
+          }
         }
+        hostColumns = null;
       }
     }
   }

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

@@ -653,8 +653,10 @@ class RowToColumnarIterator(
             .foreach(ctx => GpuSemaphore.acquireIfNecessary(ctx))
 
         val ret = withResource(new NvtxWithMetrics("RowToColumnar", NvtxColor.GREEN,
-          opTime)) { _ =>
-          builders.build(rowCount)
+            opTime)) { _ =>
+          RmmRapidsRetryIterator.withRetryNoSplit[ColumnarBatch] {
+            builders.tryBuild(rowCount)
+          }
         }
         numInputRows += rowCount
         numOutputRows += rowCount

tests/src/test/scala/com/nvidia/spark/rapids/RowToColumnarIteratorRetrySuite.scala

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed 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 com.nvidia.spark.rapids
+
+import com.nvidia.spark.rapids.jni.{RmmSpark, SplitAndRetryOOM}
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.types._
+
+class RowToColumnarIteratorRetrySuite extends RmmSparkRetrySuiteBase {
+  private val schema = StructType(Seq(StructField("a", IntegerType)))
+
+  test("test simple OOM retry") {
+    val rowIter: Iterator[InternalRow] = (1 to 10).map(InternalRow(_)).toIterator
+    val row2ColIter = new RowToColumnarIterator(
+      rowIter, schema, RequireSingleBatch, new GpuRowToColumnConverter(schema))
+    RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId)
+    Arm.withResource(row2ColIter.next()) { batch =>
+      assertResult(10)(batch.numRows())
+    }
+  }
+
+  test("test simple OOM split and retry") {
+    val rowIter: Iterator[InternalRow] = (1 to 10).map(InternalRow(_)).toIterator
+    val row2ColIter = new RowToColumnarIterator(
+      rowIter, schema, RequireSingleBatch, new GpuRowToColumnConverter(schema))
+    RmmSpark.forceSplitAndRetryOOM(RmmSpark.getCurrentThreadId)
+    assertThrows[SplitAndRetryOOM] {
+      row2ColIter.next()
+    }
+  }
+}