Clean up buffers in case AssertionError

java/src/main/java/ai/rapids/cudf/BitVectorHelper.java

@@ -71,7 +71,7 @@ private static void shiftSrcLeftAndWriteToDst(HostMemoryBuffer src, HostMemoryBu
   /**
    * This method returns the length in bytes needed to represent X number of rows
    * e.g. getValidityLengthInBytes(5) => 1 byte
-   * getLengthInBytes(7) => 1 byte
+   * getValidityLengthInBytes(7) => 1 byte
    * getValidityLengthInBytes(14) => 2 bytes
    */
   static long getValidityLengthInBytes(long rows) {

java/src/main/java/ai/rapids/cudf/ColumnVector.java

@@ -119,7 +119,7 @@ public ColumnVector(DType type, long rows, Optional<Long> nullCount,
     incRefCountInternal(true);
   }
 
-  private static OffHeapState makeOffHeap(DType type, long rows, Optional<Long> nullCount,
+  static OffHeapState makeOffHeap(DType type, long rows, Optional<Long> nullCount,
       DeviceMemoryBuffer dataBuffer, DeviceMemoryBuffer validityBuffer,
       DeviceMemoryBuffer offsetBuffer, List<DeviceMemoryBuffer> toClose, long[] childHandles) {
     long viewHandle = initViewHandle(type, (int)rows, nullCount.orElse(UNKNOWN_NULL_COUNT).intValue(),
@@ -141,7 +141,7 @@ private static OffHeapState makeOffHeap(DType type, long rows, Optional<Long> nu
    * @param offsetBuffer a host buffer required for strings and string categories. The column
    *                    vector takes ownership of the buffer. Do not use the buffer after calling
    *                    this.
-   * @param toClose  List of buffers to track adn close once done, usually in case of children
+   * @param toClose  List of buffers to track and close once done, usually in case of children
    * @param childHandles array of longs for child column view handles.
    */
   public ColumnVector(DType type, long rows, Optional<Long> nullCount,

java/src/main/java/ai/rapids/cudf/ColumnView.java

@@ -43,32 +43,50 @@ public class ColumnView implements AutoCloseable, BinaryOperable {
   protected final ColumnVector.OffHeapState offHeap;
 
   /**
-   * Constructs a Column View given a native view address
+   * Constructs a Column View given a native view address. This asserts that if the ColumnView is
+   * of nested-type it doesn't contain non-empty nulls
    * @param address the view handle
+   * @throws AssertionError if the address points to a nested-type view with non-empty nulls
    */
   ColumnView(long address) {
     this.viewHandle = address;
     this.type = DType.fromNative(ColumnView.getNativeTypeId(viewHandle), ColumnView.getNativeTypeScale(viewHandle));
     this.rows = ColumnView.getNativeRowCount(viewHandle);
     this.nullCount = ColumnView.getNativeNullCount(viewHandle);
     this.offHeap = null;
-    AssertEmptyNulls.assertNullsAreEmpty(this);
+    try {
+      AssertEmptyNulls.assertNullsAreEmpty(this);
+    } catch (AssertionError ae) {
+      // offHeap state is null, so there is nothing to clean in offHeap
+      // delete ColumnView to avoid memory leak
+      deleteColumnView(viewHandle);
+      viewHandle = 0;
+      throw ae;
+    }
   }
 
 
   /**
    * Intended to be called from ColumnVector when it is being constructed. Because state creates a
    * cudf::column_view instance and will close it in all cases, we don't want to have to double
-   * close it.
+   * close it. This asserts that if the offHeapState is of nested-type it doesn't contain non-empty nulls
    * @param state the state this view is based off of.
+   * @throws AssertionError if offHeapState points to a nested-type view with non-empty nulls
    */
   protected ColumnView(ColumnVector.OffHeapState state) {
     offHeap = state;
     viewHandle = state.getViewHandle();
     type = DType.fromNative(ColumnView.getNativeTypeId(viewHandle), ColumnView.getNativeTypeScale(viewHandle));
     rows = ColumnView.getNativeRowCount(viewHandle);
     nullCount = ColumnView.getNativeNullCount(viewHandle);
-    AssertEmptyNulls.assertNullsAreEmpty(this);
+    try {
+      AssertEmptyNulls.assertNullsAreEmpty(this);
+    } catch (AssertionError ae) {
+      // cleanup offHeap
+      offHeap.clean(false);
+      viewHandle = 0;
+      throw ae;
+    }
   }
 
   /**
@@ -649,8 +667,14 @@ public final ColumnVector ifElse(Scalar trueValue, Scalar falseValue) {
   public final ColumnVector[] slice(int... indices) {
     long[] nativeHandles = slice(this.getNativeView(), indices);
     ColumnVector[] columnVectors = new ColumnVector[nativeHandles.length];
-    for (int i = 0; i < nativeHandles.length; i++) {
-      columnVectors[i] = new ColumnVector(nativeHandles[i]);
+    try {
+      for (int i = 0; i < nativeHandles.length; i++) {
+        columnVectors[i] = new ColumnVector(nativeHandles[i]);
+        nativeHandles[i] = 0;
+      }
+    } catch (Throwable t) {
+      cleanupColumnViews(nativeHandles, columnVectors);
+      throw t;
     }
     return columnVectors;
   }
@@ -788,12 +812,31 @@ public final ColumnVector[] split(int... indices) {
   public ColumnView[] splitAsViews(int... indices) {
     long[] nativeHandles = split(this.getNativeView(), indices);
     ColumnView[] columnViews = new ColumnView[nativeHandles.length];
-    for (int i = 0; i < nativeHandles.length; i++) {
-      columnViews[i] = new ColumnView(nativeHandles[i]);
+    try {
+      for (int i = 0; i < nativeHandles.length; i++) {
+        columnViews[i] = new ColumnView(nativeHandles[i]);
+        nativeHandles[i] = 0;
+      }
+    } catch (Throwable t) {
+      cleanupColumnViews(nativeHandles, columnViews);
+      throw t;
     }
     return columnViews;
   }
 
+  static void cleanupColumnViews(long[] nativeHandles, ColumnView[] columnViews) {
+    for (ColumnView columnView: columnViews) {
+      if (columnView != null) {
+        columnView.close();
+      }
+    }
+    for (long nativeHandle: nativeHandles) {
+      if (nativeHandle != 0) {
+        deleteColumnView(nativeHandle);
+      }
+    }
+  }
+
   /**
    * Create a new vector of "normalized" values, where:
    *  1. All representations of NaN (and -NaN) are replaced with the normalized NaN value

java/src/main/java/ai/rapids/cudf/Table.java

@@ -87,6 +87,7 @@ public Table(long[] cudfColumns) {
     try {
       for (int i = 0; i < cudfColumns.length; i++) {
         this.columns[i] = new ColumnVector(cudfColumns[i]);
+        cudfColumns[i] = 0;
       }
       long[] views = new long[columns.length];
       for (int i = 0; i < columns.length; i++) {
@@ -95,13 +96,7 @@ public Table(long[] cudfColumns) {
       nativeHandle = createCudfTableView(views);
       this.rows = columns[0].getRowCount();
     } catch (Throwable t) {
-      for (int i = 0; i < cudfColumns.length; i++) {
-        if (this.columns[i] != null) {
-          this.columns[i].close();
-        } else {
-          ColumnVector.deleteCudfColumn(cudfColumns[i]);
-        }
-      }
+      ColumnView.cleanupColumnViews(cudfColumns, this.columns);
       throw t;
     }
   }
@@ -3396,8 +3391,14 @@ public static GatherMap mixedLeftAntiJoinGatherMap(Table leftKeys, Table rightKe
   public ColumnVector[] convertToRows() {
     long[] ptrs = convertToRows(nativeHandle);
     ColumnVector[] ret = new ColumnVector[ptrs.length];
-    for (int i = 0; i < ptrs.length; i++) {
-      ret[i] = new ColumnVector(ptrs[i]);
+    try {
+      for (int i = 0; i < ptrs.length; i++) {
+        ret[i] = new ColumnVector(ptrs[i]);
+        ptrs[i] = 0;
+      }
+    } catch (Throwable t) {
+      ColumnView.cleanupColumnViews(ptrs, ret);
+      throw t;
     }
     return ret;
   }
@@ -3479,8 +3480,14 @@ public ColumnVector[] convertToRows() {
   public ColumnVector[] convertToRowsFixedWidthOptimized() {
     long[] ptrs = convertToRowsFixedWidthOptimized(nativeHandle);
     ColumnVector[] ret = new ColumnVector[ptrs.length];
-    for (int i = 0; i < ptrs.length; i++) {
-      ret[i] = new ColumnVector(ptrs[i]);
+    try {
+      for (int i = 0; i < ptrs.length; i++) {
+        ret[i] = new ColumnVector(ptrs[i]);
+        ptrs[i] = 0;
+      }
+    } catch (Throwable t) {
+      ColumnView.cleanupColumnViews(ptrs, ret);
+      throw t;
     }
     return ret;
   }
@@ -3552,14 +3559,7 @@ public static Table fromPackedTable(ByteBuffer metadata, DeviceMemoryBuffer data
       }
       result = new Table(columns);
     } catch (Throwable t) {
-      for (int i = 0; i < columns.length; i++) {
-        if (columns[i] != null) {
-          columns[i].close();
-        }
-        if (columnViewAddresses[i] != 0) {
-          ColumnView.deleteColumnView(columnViewAddresses[i]);
-        }
-      }
+      ColumnView.cleanupColumnViews(columnViewAddresses, columns);
       throw t;
     }
 

java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java

@@ -32,6 +32,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -6677,6 +6678,36 @@ void testApplyBooleanMaskFromListOfStructure() {
     }
   }
 
+  @Test
+  void testColumnViewWithNonEmptyNullsIsCleared() {
+    List<Integer> list0 = Arrays.asList(1, 2, 3);
+    List<Integer> list1 = Arrays.asList(4, 5, null);
+    List<Integer> list2 = Arrays.asList(7, 8, 9);
+    List<Integer> list3 = null;
+    try (ColumnVector input = ColumnVectorTest.makeListsColumn(DType.INT32, list0, list1, list2, list3);
+         BaseDeviceMemoryBuffer baseValidityBuffer = input.getDeviceBufferFor(BufferType.VALIDITY);
+         BaseDeviceMemoryBuffer baseOffsetBuffer = input.getDeviceBufferFor(BufferType.OFFSET);
+         HostMemoryBuffer newValidity = HostMemoryBuffer.allocate(BitVectorHelper.getValidityAllocationSizeInBytes(3))) {
+
+      newValidity.copyFromDeviceBuffer(baseValidityBuffer);
+      BitVectorHelper.setNullAt(newValidity, 1);
+      DeviceMemoryBuffer validityBuffer = DeviceMemoryBuffer.allocate(BitVectorHelper.getValidityAllocationSizeInBytes(3));
+      validityBuffer.copyFromHostBuffer(newValidity);
+      DeviceMemoryBuffer offsetBuffer = DeviceMemoryBuffer.allocate(baseOffsetBuffer.getLength());
+      offsetBuffer.copyFromMemoryBuffer(0, baseOffsetBuffer, 0,
+          baseOffsetBuffer.length, Cuda.DEFAULT_STREAM);
+
+      ColumnVector.OffHeapState offHeapState = ColumnVector.makeOffHeap(input.type, input.rows, Optional.of(2L),
+          null, validityBuffer, offsetBuffer,
+          null, Arrays.stream(input.getChildColumnViews()).mapToLong((c) -> c.viewHandle).toArray());
+      try {
+        new ColumnView(offHeapState);
+      } catch (AssertionError ae) {
+        assert offHeapState.isClean();
+      }
+    }
+  }
+
   @Test
   public void testEventHandlerIsCalledForEachClose() {
     final AtomicInteger onClosedWasCalled = new AtomicInteger(0);
@@ -6700,5 +6731,4 @@ public void testEventHandlerIsNotCalledIfNotSet() {
     }
     assertEquals(0, onClosedWasCalled.get());
   }
-
 }