Skip to content

Commit

Permalink
Fix leak of spillable host buffer
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrennan333 committed Mar 27, 2024
1 parent a74adbe commit a5fe235
Showing 1 changed file with 6 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -272,25 +272,26 @@ public GpuColumnarBatchBuilder(StructType schema, int rows) {
* A collection of builders for building up columnar data.
* @param schema the schema of the batch.
* @param rows the maximum number of rows in this batch.
* @param sBuf single spillable host buffer to slice up among columns
* @param spillableHostBuf single spillable host buffer to slice up among columns
* @param bufferSizes an array of sizes for each column
*/
public GpuColumnarBatchBuilder(StructType schema, int rows,
SpillableHostBuffer sBuf, long[] bufferSizes) {
SpillableHostBuffer spillableHostBuf, long[] bufferSizes) {
fields = schema.fields();
int len = fields.length;
builders = new RapidsHostColumnBuilder[len];
boolean success = false;
long offset = 0;
try (HostMemoryBuffer hBuf =
try (SpillableHostBuffer sBuf = spillableHostBuf;
HostMemoryBuffer hBuf =
RmmRapidsRetryIterator.withRetryNoSplit(() -> sBuf.getHostBuffer());) {
long offset = 0;
for (int i = 0; i < len; i++) {
StructField field = fields[i];
try (HostMemoryBuffer columnBuffer = hBuf.slice(offset, bufferSizes[i]);) {
offset += bufferSizes[i];
builders[i] =
new RapidsHostColumnBuilder(convertFrom(field.dataType(), field.nullable()), rows)
.preAllocateBuffers(columnBuffer, 0);
.preAllocateBuffers(columnBuffer, 0);
}
}
success = true;
Expand Down

0 comments on commit a5fe235

Please sign in to comment.