diff --git a/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java b/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java index 032b075bab7..4614ce24024 100644 --- a/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java +++ b/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java @@ -1,6 +1,6 @@ /* * - * Copyright (c) 2019-2023, NVIDIA CORPORATION. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -261,6 +261,11 @@ static void register(ColumnVector vec, Cleaner cleaner) { all.put(cleaner.id, new CleanerWeakReference(vec, cleaner, collected, true)); } + static void register(Scalar s, Cleaner cleaner) { + // It is now registered... + all.put(cleaner.id, new CleanerWeakReference(s, cleaner, collected, true)); + } + static void register(HostColumnVectorCore vec, Cleaner cleaner) { // It is now registered... all.put(cleaner.id, new CleanerWeakReference(vec, cleaner, collected, false)); diff --git a/java/src/main/java/ai/rapids/cudf/Scalar.java b/java/src/main/java/ai/rapids/cudf/Scalar.java index 70538ab082f..286b5c208c9 100644 --- a/java/src/main/java/ai/rapids/cudf/Scalar.java +++ b/java/src/main/java/ai/rapids/cudf/Scalar.java @@ -1,6 +1,6 @@ /* * - * Copyright (c) 2019-2022, NVIDIA CORPORATION. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -524,6 +524,7 @@ private static ColumnVector buildNullColumnVector(HostColumnVector.DataType host Scalar(DType type, long scalarHandle) { this.type = type; this.offHeap = new OffHeapState(scalarHandle); + MemoryCleaner.register(this, offHeap); incRefCount(); } @@ -536,6 +537,7 @@ public synchronized Scalar incRefCount() { offHeap.logRefCountDebug("INC AFTER CLOSE " + this); throw new IllegalStateException("Scalar is already closed"); } + offHeap.addRef(); ++refCount; return this; }