Skip to content

Commit

Permalink
Use UDF
Browse files Browse the repository at this point in the history
  • Loading branch information
Chong Gao committed Dec 17, 2024
1 parent 8aaf0f6 commit 5bfb544
Show file tree
Hide file tree
Showing 7 changed files with 1,461 additions and 1 deletion.
4 changes: 3 additions & 1 deletion src/main/cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,7 @@ add_library(
src/DateTimeRebaseJni.cpp
src/DecimalUtilsJni.cpp
src/GpuTimeZoneDBJni.cpp
src/HLLPPHostUDFJni.cpp
src/HashJni.cpp
src/HistogramJni.cpp
src/HostTableJni.cpp
Expand All @@ -205,7 +206,6 @@ add_library(
src/SparkResourceAdaptorJni.cpp
src/SubStringIndexJni.cpp
src/ZOrderJni.cpp
src/HLLPP.cu
src/bloom_filter.cu
src/case_when.cu
src/cast_decimal_to_string.cu
Expand All @@ -219,6 +219,8 @@ add_library(
src/from_json_to_structs.cu
src/get_json_object.cu
src/histogram.cu
src/hllpp_host_udf.cu
src/hllpp.cu
src/json_utils.cu
src/murmur_hash.cu
src/parse_uri.cu
Expand Down
66 changes: 66 additions & 0 deletions src/main/cpp/src/HLLPPHostUDFJni.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
/*
* Copyright (c) 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.
* 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 "hllpp.hpp"
#include "hllpp_host_udf.hpp"

extern "C" {

JNIEXPORT jlong JNICALL
Java_com_nvidia_spark_rapids_jni_HLLPPHostUDF_createHLLPPHostUDF(
JNIEnv *env, jclass, jint agg_type, int precision) {
try {
cudf::jni::auto_set_device(env);
auto udf_ptr = [&] {
// The value of agg_type must be sync with
// `HLLPPHostUDF.java#AggregationType`.
switch (agg_type) {
case 0:
return spark_rapids_jni::create_hllpp_reduction_host_udf(precision);
case 1:
return spark_rapids_jni::create_hllpp_reduction_merge_host_udf(
precision);
case 2:
return spark_rapids_jni::create_hllpp_groupby_host_udf(precision);
default:
return spark_rapids_jni::create_hllpp_groupby_merge_host_udf(precision);
}
}();
CUDF_EXPECTS(udf_ptr != nullptr,
"Invalid HyperLogLogPlusPlus(HLLPP) UDF instance.");

return reinterpret_cast<jlong>(udf_ptr.release());
}
CATCH_STD(env, 0);
}

JNIEXPORT jlong JNICALL
Java_com_nvidia_spark_rapids_jni_HLLPPHostUDF_estimateDistinctValueFromSketches(
JNIEnv *env, jclass, jlong sketches, jint precision) {
JNI_NULL_CHECK(env, sketches, "Sketch column is null", 0);
try {
cudf::jni::auto_set_device(env);
auto const sketch_view =
reinterpret_cast<cudf::column_view const *>(sketches);
return cudf::jni::ptr_as_jlong(
spark_rapids_jni::estimate_from_hll_sketches(*sketch_view, precision)
.release());
}
CATCH_STD(env, 0);
}

} // extern "C"
Loading

0 comments on commit 5bfb544

Please sign in to comment.