Optimize date/time truncation for scalar format (#2687)

Signed-off-by: Nghia Truong <[email protected]>
NVIDIA · Dec 13, 2024 · 8805168 · 8805168
1 parent daf34f1
commit 8805168
Show file tree

Hide file tree

Showing 4 changed files with 216 additions and 90 deletions.
diff --git a/src/main/cpp/src/DateTimeUtilsJni.cpp b/src/main/cpp/src/DateTimeUtilsJni.cpp
@@ -47,20 +47,34 @@ JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_DateTimeUtils_rebaseJul
   CATCH_STD(env, 0);
 }
 
-JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_DateTimeUtils_truncate(JNIEnv* env,
-                                                                                jclass,
-                                                                                jlong datetime,
-                                                                                jlong format)
+JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_DateTimeUtils_truncateWithColumnFormat(
+  JNIEnv* env, jclass, jlong datetime, jlong format)
 {
-  JNI_NULL_CHECK(env, datetime, "datetime column is null", 0);
-  JNI_NULL_CHECK(env, format, "format column is null", 0);
+  JNI_NULL_CHECK(env, datetime, "input datetime is null", 0);
+  JNI_NULL_CHECK(env, format, "input format is null", 0);
 
   try {
     cudf::jni::auto_set_device(env);
+
     auto const datetime_cv = reinterpret_cast<cudf::column_view const*>(datetime);
     auto const format_cv   = reinterpret_cast<cudf::column_view const*>(format);
-    auto output            = spark_rapids_jni::truncate(*datetime_cv, *format_cv);
-    return reinterpret_cast<jlong>(output.release());
+    return reinterpret_cast<jlong>(spark_rapids_jni::truncate(*datetime_cv, *format_cv).release());
+  }
+  CATCH_STD(env, 0);
+}
+
+JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_DateTimeUtils_truncateWithScalarFormat(
+  JNIEnv* env, jclass, jlong datetime, jstring format)
+{
+  JNI_NULL_CHECK(env, datetime, "input datetime is null", 0);
+
+  try {
+    cudf::jni::auto_set_device(env);
+
+    auto const datetime_cv = reinterpret_cast<cudf::column_view const*>(datetime);
+    auto const format_jstr = cudf::jni::native_jstring(env, format);
+    auto const format      = std::string(format_jstr.get(), format_jstr.size_bytes());
+    return reinterpret_cast<jlong>(spark_rapids_jni::truncate(*datetime_cv, format).release());
   }
   CATCH_STD(env, 0);
 }