rename max-row_size and rearrange reaching limit

Signed-off-by: Ahmed Hussein (amahussein) <a@ahussein.me>
NVIDIA · Jul 22, 2022 · 15c00a8 · 15c00a8
1 parent 0858a7d
commit 15c00a8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/sql-plugin/src/main/java/com/nvidia/spark/rapids/InternalRowToColumnarBatchIterator.java b/sql-plugin/src/main/java/com/nvidia/spark/rapids/InternalRowToColumnarBatchIterator.java
@@ -188,13 +188,12 @@ public ColumnarBatch next() {
         // Handle corner case when the dataLength is too small to copy a single row.
         // We retry after doubling the bufferSize.
         // dataLength can be considered a rough estimate of a single row.
-        long newRowSizeEst = dataLength << 1 ;
-        newRowSizeEst = Math.min(newRowSizeEst, JCudfUtil.JCUDF_MAX_DATA_BUFFER_LENGTH);
-        if (newRowSizeEst <= dataLength) { // We already reached the limit.
+        if (dataLength >= JCudfUtil.JCUDF_MAX_ROW_BUFFER_LENGTH) {
           // proceed with throwing exception.
           throw new RuntimeException(
               "JCudf row is too large to fit in MemoryBuffer", ex);
         }
+        long newRowSizeEst = Math.min(dataLength << 1, JCudfUtil.JCUDF_MAX_ROW_BUFFER_LENGTH);
         // Recalculate the dataLength based on the new size estimate
         setBuffersCapacities(newRowSizeEst, goalTargetSize);
         retryCount++;

diff --git a/sql-plugin/src/main/java/com/nvidia/spark/rapids/JCudfUtil.java b/sql-plugin/src/main/java/com/nvidia/spark/rapids/JCudfUtil.java
@@ -60,9 +60,9 @@ public final class JCudfUtil {
   public static final int JCUDF_OPTIMIZED_CUDF_KERNEL_MAX_BYTES = 1536;
 
   /**
-   * The maximum buffer size allocated to copy JCudf row.
+   * The maximum buffer size allocated to copy JCudf row (8-bytes aligned).
    */
-  public static final long JCUDF_MAX_DATA_BUFFER_LENGTH =
+  public static final long JCUDF_MAX_ROW_BUFFER_LENGTH =
       Integer.MAX_VALUE - (JCUDF_ROW_ALIGNMENT - 1);
 
   /**