[Enhancement] support DATE and DATETIME types in Java UDF (backport #72337)

be/src/udf/java/java_data_converter.cpp

@@ -269,6 +269,8 @@ DEFINE_CAST_TO_JVALUE(TYPE_BIGINT, helper.newLong(data_value));
 DEFINE_CAST_TO_JVALUE(TYPE_FLOAT, helper.newFloat(data_value));
 DEFINE_CAST_TO_JVALUE(TYPE_DOUBLE, helper.newDouble(data_value));
 DEFINE_CAST_TO_JVALUE(TYPE_VARCHAR, helper.newString(data_value.get_data(), data_value.get_size()));
+DEFINE_CAST_TO_JVALUE(TYPE_DATE, helper.newLocalDate(data_value._julian));
+DEFINE_CAST_TO_JVALUE(TYPE_DATETIME, helper.newLocalDateTime(data_value.timestamp()));
 
 void release_jvalue(bool is_box, jvalue val) {
     if (is_box && val.l) {
@@ -333,6 +335,16 @@ StatusOr<jvalue> cast_to_jvalue(const TypeDescriptor& type_desc, bool is_boxed,
         v.l = helper.newString(slice.get_data(), slice.get_size());
         break;
     }
+    case TYPE_DATE: {
+        auto spec_col = down_cast<const RunTimeColumnType<TYPE_DATE>*>(col);
+        v.l = helper.newLocalDate(spec_col->immutable_data()[row_num]._julian);
+        break;
+    }
+    case TYPE_DATETIME: {
+        auto spec_col = down_cast<const RunTimeColumnType<TYPE_DATETIME>*>(col);
+        v.l = helper.newLocalDateTime(spec_col->immutable_data()[row_num].timestamp());
+        break;
+    }
     case TYPE_DECIMAL32:
     case TYPE_DECIMAL64:
     case TYPE_DECIMAL128:
@@ -446,6 +458,18 @@ Status assign_jvalue(const TypeDescriptor& type_desc, bool is_box, Column* col,
         ASSIGN_BOX_TYPE(TYPE_BIGINT, int64_t)
         ASSIGN_BOX_TYPE(TYPE_FLOAT, float)
         ASSIGN_BOX_TYPE(TYPE_DOUBLE, double)
+    case TYPE_DATE: {
+        DateValue dv;
+        dv._julian = helper.valLocalDate(val.l);
+        down_cast<RunTimeColumnType<TYPE_DATE>*>(data_col)->get_data()[row_num] = dv;
+        break;
+    }
+    case TYPE_DATETIME: {
+        TimestampValue tv;
+        tv.set_timestamp(helper.valLocalDateTime(val.l));
+        down_cast<RunTimeColumnType<TYPE_DATETIME>*>(data_col)->get_data()[row_num] = tv;
+        break;
+    }
     case TYPE_VARCHAR: {
         if (val.l == nullptr) {
             col->append_nulls(1);
@@ -542,6 +566,18 @@ Status append_jvalue(const TypeDescriptor& type_desc, bool is_box, Column* col,
             APPEND_BOX_TYPE(TYPE_FLOAT, float)
             APPEND_BOX_TYPE(TYPE_DOUBLE, double)
 
+        case TYPE_DATE: {
+            DateValue dv;
+            dv._julian = helper.valLocalDate(val.l);
+            col->append_datum(Datum(dv));
+            break;
+        }
+        case TYPE_DATETIME: {
+            TimestampValue tv;
+            tv.set_timestamp(helper.valLocalDateTime(val.l));
+            col->append_datum(Datum(tv));
+            break;
+        }
         case TYPE_VARCHAR: {
             std::string buffer;
             auto slice = helper.sliceVal((jstring)val.l, &buffer);
@@ -660,6 +696,24 @@ Status check_type_matched(const TypeDescriptor& type_desc, jobject val) {
         }
         break;
     }
+    case TYPE_DATE: {
+        if (!env->IsInstanceOf(val, helper.local_date_class())) {
+            auto clazz = env->GetObjectClass(val);
+            LOCAL_REF_GUARD(clazz);
+            return Status::InternalError(
+                    fmt::format("Type not matched, expect java.time.LocalDate, but got {}", helper.to_string(clazz)));
+        }
+        break;
+    }
+    case TYPE_DATETIME: {
+        if (!env->IsInstanceOf(val, helper.local_datetime_class())) {
+            auto clazz = env->GetObjectClass(val);
+            LOCAL_REF_GUARD(clazz);
+            return Status::InternalError(fmt::format("Type not matched, expect java.time.LocalDateTime, but got {}",
+                                                     helper.to_string(clazz)));
+        }
+        break;
+    }
     case TYPE_ARRAY: {
         if (!env->IsInstanceOf(val, helper.list_meta().list_class->clazz())) {
             auto clazz = env->GetObjectClass(val);

be/src/udf/java/java_native_method.cpp

@@ -26,7 +26,9 @@
 #include "column/nullable_column.h"
 #include "column/vectorized_fwd.h"
 #include "gutil/casts.h"
+#include "types/date_value.h"
 #include "types/logical_type.h"
+#include "types/timestamp_value.h"
 
 namespace starrocks {
 
@@ -138,6 +140,12 @@ class GetColumnLogicalTypeVistor : public ColumnVisitorAdapter<GetColumnLogicalT
         } else if constexpr (std::is_same_v<T, double>) {
             *_result = TYPE_DOUBLE;
             return Status::OK();
+        } else if constexpr (std::is_same_v<T, DateValue>) {
+            *_result = TYPE_DATE;
+            return Status::OK();
+        } else if constexpr (std::is_same_v<T, TimestampValue>) {
+            *_result = TYPE_DATETIME;
+            return Status::OK();
         }
         return Status::NotSupported("unsupported UDF type");
     }

be/src/udf/java/java_udf.cpp

@@ -25,7 +25,9 @@
 #include "exprs/function_context.h"
 #include "fmt/core.h"
 #include "jni.h"
+#include "types/date_value.h"
 #include "types/logical_type.h"
+#include "types/timestamp_value.h"
 #include "udf/java/java_native_method.h"
 #include "udf/java/type_traits.h"
 #include "udf/java/utils.h"
@@ -190,6 +192,16 @@ void JVMFunctionHelper::_init() {
     _big_decimal_value_of_ll = _env->GetStaticMethodID(_big_decimal_class, "valueOf", "(JI)Ljava/math/BigDecimal;");
     CHECK(_big_decimal_value_of_ll);
 
+    _local_date_class = JNI_FIND_CLASS("java/time/LocalDate");
+    CHECK(_local_date_class);
+    _local_date_of_epoch_day = _env->GetStaticMethodID(_local_date_class, "ofEpochDay", "(J)Ljava/time/LocalDate;");
+    CHECK(_local_date_of_epoch_day);
+    _local_date_to_epoch_day = _env->GetMethodID(_local_date_class, "toEpochDay", "()J");
+    CHECK(_local_date_to_epoch_day);
+
+    _local_datetime_class = JNI_FIND_CLASS("java/time/LocalDateTime");
+    CHECK(_local_datetime_class);
+
     ADD_NUMBERIC_CLASS(boolean, Boolean, Z);
     ADD_NUMBERIC_CLASS(byte, Byte, B);
     ADD_NUMBERIC_CLASS(short, Short, S);
@@ -262,6 +274,15 @@ void JVMFunctionHelper::_init() {
     _method_map.emplace(TYPE_ARRAY_METHOD_ID, tmp);
     INIT_HELPER_METHOD(tmp, "createBoxedMapArray", CREATE_BOXED_MAP_SIGNATURE);
     _method_map.emplace(TYPE_MAP_METHOD_ID, tmp);
+    INIT_HELPER_METHOD(tmp, "createBoxedLocalDateArray", CREATE_BOXED_PRIMI_SIGNATURE);
+    _method_map.emplace(JNIPrimTypeId<DateValue>::id, tmp);
+    INIT_HELPER_METHOD(tmp, "createBoxedLocalDateTimeArray", CREATE_BOXED_PRIMI_SIGNATURE);
+    _method_map.emplace(JNIPrimTypeId<TimestampValue>::id, tmp);
+
+    // LocalDateTime <-> packed int64 conversion stays in UDFHelper because the
+    // packing is StarRocks-specific.
+    INIT_HELPER_METHOD(_local_datetime_from_packed, "localDateTimeFromPackedTimestamp", "(J)Ljava/time/LocalDateTime;");
+    INIT_HELPER_METHOD(_local_datetime_to_packed, "packedTimestampFromLocalDateTime", "(Ljava/time/LocalDateTime;)J");
 
     // init bytebuffer
     _direct_buffer_class = JNI_FIND_CLASS("java/nio/ByteBuffer");
@@ -594,6 +615,33 @@ jbyteArray JVMFunctionHelper::unscaled_le_bytes(jobject big_decimal, int precisi
                                                     static_cast<jint>(byte_width));
 }
 
+// DateValue stores days as Julian day; LocalDate exposes them as days-since-1970-01-01.
+// 2440588 is the Julian day number of the Unix epoch.
+static constexpr jlong UNIX_EPOCH_JULIAN_DAYS = 2440588;
+
+jobject JVMFunctionHelper::newLocalDate(int32_t julian) {
+    jobject ld = _env->CallStaticObjectMethod(_local_date_class, _local_date_of_epoch_day,
+                                              static_cast<jlong>(julian) - UNIX_EPOCH_JULIAN_DAYS);
+    RETURN_IF_JNI_EXCEPTION(_env, "newLocalDate: ofEpochDay failed", nullptr);
+    return ld;
+}
+
+int32_t JVMFunctionHelper::valLocalDate(jobject obj) {
+    jlong epoch_day = _env->CallLongMethod(obj, _local_date_to_epoch_day);
+    return static_cast<int32_t>(epoch_day + UNIX_EPOCH_JULIAN_DAYS);
+}
+
+jobject JVMFunctionHelper::newLocalDateTime(int64_t packed_timestamp) {
+    jobject ldt = _env->CallStaticObjectMethod(_udf_helper_class, _local_datetime_from_packed,
+                                               static_cast<jlong>(packed_timestamp));
+    RETURN_IF_JNI_EXCEPTION(_env, "newLocalDateTime: localDateTimeFromPackedTimestamp failed", nullptr);
+    return ldt;
+}
+
+int64_t JVMFunctionHelper::valLocalDateTime(jobject obj) {
+    return _env->CallStaticLongMethod(_udf_helper_class, _local_datetime_to_packed, obj);
+}
+
 Slice JVMFunctionHelper::sliceVal(jstring jstr, std::string* buffer) {
     const size_t utf_length = _env->GetStringUTFLength(jstr);
     buffer->resize(utf_length);
@@ -1094,6 +1142,15 @@ Status ClassAnalyzer::get_udaf_method_desc(const std::string& sign, std::vector<
                 // BigDecimal params. The actual DECIMAL precision/scale is resolved from
                 // ctx->get_arg_type() / ctx->get_return_type() at the call site.
                 desc->emplace_back(MethodTypeDescriptor{TYPE_DECIMAL128, true});
+            } else if (type == "java/time/LocalDate") {
+                desc->emplace_back(MethodTypeDescriptor{TYPE_DATE, true});
+            } else if (type == "java/time/LocalDateTime") {
+                desc->emplace_back(MethodTypeDescriptor{TYPE_DATETIME, true});
+            } else {
+                // Unrecognized object class. Surface as TYPE_UNKNOWN so get_method_desc()
+                // validation produces a clear error instead of leaving method_desc short
+                // and crashing on a later method_desc[0].is_box access.
+                desc->emplace_back(MethodTypeDescriptor{TYPE_UNKNOWN, true});
             }
             continue;
         }

be/src/udf/java/java_udf.h

@@ -196,6 +196,18 @@ class JVMFunctionHelper {
     Slice sliceVal(jstring jstr, std::string* buffer);
     jclass string_clazz() { return _string_class; }
     jclass big_decimal_class() { return _big_decimal_class; }
+    jclass local_date_class() { return _local_date_class; }
+    jclass local_datetime_class() { return _local_datetime_class; }
+
+    // Box/unbox a StarRocks DateValue (int32 Julian day) as a java.time.LocalDate.
+    // Round-trips StarRocks's internal Julian-day encoding.
+    jobject newLocalDate(int32_t julian);
+    int32_t valLocalDate(jobject obj);
+
+    // Box/unbox a StarRocks TimestampValue (packed int64: julian << 40 | usOfDay)
+    // as a java.time.LocalDateTime.
+    jobject newLocalDateTime(int64_t packed_timestamp);
+    int64_t valLocalDateTime(jobject obj);
     // replace '.' as '/'
     // eg: java.lang.Integer -> java/lang/Integer
     static std::string to_jni_class_name(const std::string& name);
@@ -232,10 +244,18 @@ class JVMFunctionHelper {
     jclass _jarrays_class;
     jclass _exception_util_class;
     jclass _big_decimal_class;
+    jclass _local_date_class;
+    jclass _local_datetime_class;
 
     jmethodID _string_construct_with_bytes;
     jmethodID _big_decimal_ctor_string;
     jmethodID _big_decimal_value_of_ll;
+    // java.time.LocalDate.ofEpochDay(long) / java.time.LocalDate.toEpochDay()
+    jmethodID _local_date_of_epoch_day;
+    jmethodID _local_date_to_epoch_day;
+    // UDFHelper.localDateTimeFromPackedTimestamp(long) / packedTimestampFromLocalDateTime(LocalDateTime)
+    jmethodID _local_datetime_from_packed;
+    jmethodID _local_datetime_to_packed;
 
     ListMeta _list_meta;
     MapMeta _map_meta;

be/src/udf/java/type_traits.h

@@ -19,6 +19,9 @@
 #include "util/slice.h"
 
 namespace starrocks {
+class DateValue;
+class TimestampValue;
+
 template <class Type>
 struct JNIPrimTypeId {
     static constexpr bool supported = false;
@@ -42,5 +45,9 @@ DEFINE_TYPE(double, 8);
 DEFINE_TYPE(Slice, 9);
 static constexpr int TYPE_ARRAY_METHOD_ID = 10;
 static constexpr int TYPE_MAP_METHOD_ID = 11;
+// DateValue/TimestampValue are PODs whose storage IS the wire format
+// (int32 Julian day / packed int64). Java UDFHelper decodes inline.
+DEFINE_TYPE(DateValue, 12);
+DEFINE_TYPE(TimestampValue, 13);
 
 } // namespace starrocks