diff --git a/java/src/main/java/ai/rapids/cudf/Aggregation.java b/java/src/main/java/ai/rapids/cudf/Aggregation.java index c07a58ed8a5..42fa871abeb 100644 --- a/java/src/main/java/ai/rapids/cudf/Aggregation.java +++ b/java/src/main/java/ai/rapids/cudf/Aggregation.java @@ -396,7 +396,16 @@ private HostUDFAggregation(HostUDFWrapper wrapper) { @Override long createNativeInstance() { - return Aggregation.createHostUDFAgg(wrapper.udfNativeHandle); + long udf = 0; + try { + udf = wrapper.createUDFInstance(); + return Aggregation.createHostUDFAgg(udf); + } finally { + // a new UDF is cloned in `createHostUDFAgg`, here should close the UDF instance. + if (udf != 0) { + HostUDFWrapper.closeUDFInstance(udf); + } + } } @Override diff --git a/java/src/main/java/ai/rapids/cudf/HostUDFWrapper.java b/java/src/main/java/ai/rapids/cudf/HostUDFWrapper.java index 124f2c99188..0bd3ccad0d3 100644 --- a/java/src/main/java/ai/rapids/cudf/HostUDFWrapper.java +++ b/java/src/main/java/ai/rapids/cudf/HostUDFWrapper.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -19,18 +19,56 @@ /** * A wrapper around native host UDF aggregations. *
- * This class is used to store the native handle of a host UDF aggregation and is used as + * This class is used to create the native handle of a host UDF aggregation and is used as * a proxy object to compute hash code and compare two host UDF aggregations for equality. *
* A new host UDF aggregation implementation must extend this class and override the - * {@code hashCode} and {@code equals} methods for such purposes. - * In addition, since this class implements {@code AutoCloseable}, the {@code close} method must - * also be overridden to automatically delete the native UDF instance upon class destruction. + * {@code computeHashCode} and {@code isEqual} methods for such purposes. + * */ -public abstract class HostUDFWrapper implements AutoCloseable { - public final long udfNativeHandle; +public abstract class HostUDFWrapper { + + /** + * Create a derived host UDF native instance. + * The instance created by this function MUST be closed by `closeUDFInstance` + *
Typical usage, refer to Aggregation.java:
+ *+ * long udf = 0; + * try { + * udf = wrapper.createUDFInstance(); + * return Aggregation.createHostUDFAgg(udf); + * } finally { + * // a new UDF is cloned in `createHostUDFAgg`, here should close the UDF instance. + * if (udf != 0) { + * HostUDFWrapper.closeUDFInstance(udf); + * } + * } + *+ * + */ + public abstract long createUDFInstance(); + + /** + * Close the derived UDF instance created by `createUDFInstance`. + * @param hostUDFInstance the UDF instance + */ + public static void closeUDFInstance(long hostUDFInstance) { + close(hostUDFInstance); + } + + public abstract int computeHashCode(); - public HostUDFWrapper(long udfNativeHandle) { - this.udfNativeHandle = udfNativeHandle; + @Override + public int hashCode() { + return computeHashCode(); } + + public abstract boolean isEqual(Object obj); + + @Override + public boolean equals(Object obj) { + return isEqual(obj); + } + + static native void close(long hostUDFInstance); } diff --git a/java/src/main/native/CMakeLists.txt b/java/src/main/native/CMakeLists.txt index bd1714aa476..3923d8b45e3 100644 --- a/java/src/main/native/CMakeLists.txt +++ b/java/src/main/native/CMakeLists.txt @@ -148,6 +148,7 @@ add_library( src/DataSourceHelperJni.cpp src/HashJoinJni.cpp src/HostMemoryBufferNativeUtilsJni.cpp + src/HostUDFWrapperJni.cpp src/NvcompJni.cpp src/NvtxRangeJni.cpp src/NvtxUniqueRangeJni.cpp diff --git a/java/src/main/native/src/HostUDFWrapperJni.cpp b/java/src/main/native/src/HostUDFWrapperJni.cpp new file mode 100644 index 00000000000..1dd1b483fac --- /dev/null +++ b/java/src/main/native/src/HostUDFWrapperJni.cpp @@ -0,0 +1,34 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cudf_jni_apis.hpp" + +#include