Skip to content

Commit

Permalink
Fix JNI leak on copy to device (#10229)
Browse files Browse the repository at this point in the history
We were not closing intermediate column view handles when copying nested data from the device to the host. I think we could completely replace `NestedColumnVector` with `ColumnView` in the future, but it would require some optimizations in `ColumnView` to make it happen cleanly, so I didn't do any of that here.

I also renamed one of the methods that had a leak so it would be more clear what it was doing, and I marked it as private so it would be less likely to leak out and be used by other code.

Authors:
  - Robert (Bobby) Evans (https://github.com/revans2)

Approvers:
  - Jason Lowe (https://github.com/jlowe)

URL: #10229
  • Loading branch information
revans2 authored Feb 7, 2022
1 parent 2e458b9 commit 7987791
Showing 1 changed file with 40 additions and 20 deletions.
60 changes: 40 additions & 20 deletions java/src/main/java/ai/rapids/cudf/ColumnView.java
Original file line number Diff line number Diff line change
Expand Up @@ -4012,20 +4012,29 @@ static ColumnVector createColumnVector(DType type, int rows, HostMemoryBuffer da
}
if (mainColOffsets != null) {
// The offset buffer has (no. of rows + 1) entries, where each entry is INT32.sizeInBytes
long offsetsLen = OFFSET_SIZE * (mainColRows + 1);
long offsetsLen = OFFSET_SIZE * (((long)mainColRows) + 1);
mainOffsetsDevBuff = DeviceMemoryBuffer.allocate(offsetsLen);
mainOffsetsDevBuff.copyFromHostBuffer(mainColOffsets, 0, offsetsLen);
}
List<DeviceMemoryBuffer> toClose = new ArrayList<>();
long[] childHandles = new long[devChildren.size()];
for (ColumnView.NestedColumnVector ncv : devChildren) {
toClose.addAll(ncv.getBuffersToClose());
}
for (int i = 0; i < devChildren.size(); i++) {
childHandles[i] = devChildren.get(i).getViewHandle();
try {
for (ColumnView.NestedColumnVector ncv : devChildren) {
toClose.addAll(ncv.getBuffersToClose());
}
for (int i = 0; i < devChildren.size(); i++) {
childHandles[i] = devChildren.get(i).createViewHandle();
}
return new ColumnVector(mainColType, mainColRows, nullCount, mainDataDevBuff,
mainValidDevBuff, mainOffsetsDevBuff, toClose, childHandles);
} finally {
for (int i = 0; i < childHandles.length; i++) {
if (childHandles[i] != 0) {
ColumnView.deleteColumnView(childHandles[i]);
childHandles[i] = 0;
}
}
}
return new ColumnVector(mainColType, mainColRows, nullCount, mainDataDevBuff,
mainValidDevBuff, mainOffsetsDevBuff, toClose, childHandles);
}

private static NestedColumnVector createNewNestedColumnVector(
Expand All @@ -4048,21 +4057,32 @@ private static NestedColumnVector createNewNestedColumnVector(
children);
}

long getViewHandle() {
private long createViewHandle() {
long[] childrenColViews = null;
if (children != null) {
childrenColViews = new long[children.size()];
for (int i = 0; i < children.size(); i++) {
childrenColViews[i] = children.get(i).getViewHandle();
try {
if (children != null) {
childrenColViews = new long[children.size()];
for (int i = 0; i < children.size(); i++) {
childrenColViews[i] = children.get(i).createViewHandle();
}
}
long dataAddr = data == null ? 0 : data.address;
long dataLen = data == null ? 0 : data.length;
long offsetAddr = offsets == null ? 0 : offsets.address;
long validAddr = valid == null ? 0 : valid.address;
int nc = nullCount.orElse(ColumnVector.OffHeapState.UNKNOWN_NULL_COUNT).intValue();
return makeCudfColumnView(dataType.typeId.getNativeId(), dataType.getScale(), dataAddr, dataLen,
offsetAddr, validAddr, nc, (int) rows, childrenColViews);
} finally {
if (childrenColViews != null) {
for (int i = 0; i < childrenColViews.length; i++) {
if (childrenColViews[i] != 0) {
deleteColumnView(childrenColViews[i]);
childrenColViews[i] = 0;
}
}
}
}
long dataAddr = data == null ? 0 : data.address;
long dataLen = data == null ? 0 : data.length;
long offsetAddr = offsets == null ? 0 : offsets.address;
long validAddr = valid == null ? 0 : valid.address;
int nc = nullCount.orElse(ColumnVector.OffHeapState.UNKNOWN_NULL_COUNT).intValue();
return makeCudfColumnView(dataType.typeId.getNativeId(), dataType.getScale() , dataAddr, dataLen,
offsetAddr, validAddr, nc, (int)rows, childrenColViews);
}

List<DeviceMemoryBuffer> getBuffersToClose() {
Expand Down

0 comments on commit 7987791

Please sign in to comment.