Expose max nesting depth in hash function to plugin

Signed-off-by: Yan Feng <[email protected]>
NVIDIA · Dec 11, 2024 · e989b56 · e989b56
1 parent bd8b7e6
commit e989b56
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 2 deletions.
diff --git a/src/main/cpp/src/HashJni.cpp b/src/main/cpp/src/HashJni.cpp
@@ -21,6 +21,15 @@
 
 extern "C" {
 
+JNIEXPORT jint JNICALL Java_com_nvidia_spark_rapids_jni_Hash_getMaxNestedDepth(JNIEnv* env, jclass)
+{
+  try {
+    cudf::jni::auto_set_device(env);
+    return spark_rapids_jni::MAX_NESTED_DEPTH;
+  }
+  CATCH_STD(env, 0);
+}
+
 JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_Hash_murmurHash32(
   JNIEnv* env, jclass, jint seed, jlongArray column_handles)
 {

diff --git a/src/main/cpp/src/hash.hpp b/src/main/cpp/src/hash.hpp
@@ -25,6 +25,7 @@
 namespace spark_rapids_jni {
 
 constexpr int64_t DEFAULT_XXHASH64_SEED = 42;
+constexpr int MAX_NESTED_DEPTH = 8;
 
 /**
  * @brief Computes the murmur32 hash value of each row in the input set of columns.

diff --git a/src/main/cpp/src/hive_hash.cu b/src/main/cpp/src/hive_hash.cu
@@ -15,6 +15,7 @@
  */
 
 #include "hash.cuh"
+#include "hash.hpp"
 
 #include <cudf/column/column_factories.hpp>
 #include <cudf/detail/utilities/vector_factories.hpp>
@@ -37,8 +38,6 @@ using hive_hash_value_t = int32_t;
 constexpr hive_hash_value_t HIVE_HASH_FACTOR = 31;
 constexpr hive_hash_value_t HIVE_INIT_HASH   = 0;
 
-constexpr int MAX_NESTED_DEPTH = 8;
-
 hive_hash_value_t __device__ inline compute_int(int32_t key) { return key; }
 
 hive_hash_value_t __device__ inline compute_long(int64_t key)

diff --git a/src/main/java/com/nvidia/spark/rapids/jni/Hash.java b/src/main/java/com/nvidia/spark/rapids/jni/Hash.java
@@ -25,6 +25,8 @@ public class Hash {
   // there doesn't appear to be a useful constant in spark to reference. this could break.
   static final long DEFAULT_XXHASH64_SEED = 42;
 
+  public static final int MAX_NESTED_DEPTH = getMaxNestedDepth();
+
   static {
     NativeDepsLoader.loadNativeDeps();
   }
@@ -101,6 +103,8 @@ public static ColumnVector hiveHash(ColumnView columns[]) {
     return new ColumnVector(hiveHash(columnViews));
   }
 
+  private static native int getMaxNestedDepth();
+
   private static native long murmurHash32(int seed, long[] viewHandles) throws CudfException;
 
   private static native long xxhash64(long seed, long[] viewHandles) throws CudfException;