Skip to content

Commit c3da023

Browse files
Hyper-FFmergify[bot]
Hyper-FF
authored and
mergify[bot]
committed
[Enhancement] support DATE and DATETIME types in Java UDF (#72337)
Co-authored-by: Claude <noreply@anthropic.com> (cherry picked from commit add5dee) # Conflicts: # be/src/udf/java/java_data_converter.cpp # be/src/udf/java/java_udf.cpp # be/src/udf/java/java_udf.h # be/test/udf/java/java_data_converter_test.cpp # be/test/udf/java/java_native_method_test.cpp # be/test/udf/java/java_udf_test.cpp # docs/en/sql-reference/sql-functions/JAVA_UDF.md # docs/ja/sql-reference/sql-functions/JAVA_UDF.md # docs/zh/sql-reference/sql-functions/JAVA_UDF.md # fe/fe-core/src/main/java/com/starrocks/sql/analyzer/CreateFunctionAnalyzer.java # fe/fe-core/src/test/java/com/starrocks/sql/analyzer/CreateFunctionStmtAnalyzerTest.java # java-extensions/udf-extensions/src/main/java/com/starrocks/udf/UDFHelper.java
1 parent d9f4c90 commit c3da023

17 files changed

Lines changed: 2235 additions & 0 deletions

File tree

‎be/src/udf/java/java_data_converter.cpp‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,17 @@
2929
#include "column/vectorized_fwd.h"
3030
#include "common/compiler_util.h"
3131
#include "common/status.h"
32+
<<<<<<< HEAD
3233
#include "runtime/types.h"
3334
#include "types/logical_type.h"
35+
=======
36+
#include "types/date_value.h"
37+
#include "types/datum.h"
38+
#include "types/decimalv3.h"
39+
#include "types/logical_type.h"
40+
#include "types/timestamp_value.h"
41+
#include "types/type_descriptor.h"
42+
>>>>>>> add5deeffc ([Enhancement] support DATE and DATETIME types in Java UDF (#72337))
3443
#include "udf/java/java_udf.h"
3544
#include "udf/java/type_traits.h"
3645
#include "util/defer_op.h"
@@ -179,6 +188,8 @@ DEFINE_CAST_TO_JVALUE(TYPE_BIGINT, helper.newLong(data_value));
179188
DEFINE_CAST_TO_JVALUE(TYPE_FLOAT, helper.newFloat(data_value));
180189
DEFINE_CAST_TO_JVALUE(TYPE_DOUBLE, helper.newDouble(data_value));
181190
DEFINE_CAST_TO_JVALUE(TYPE_VARCHAR, helper.newString(data_value.get_data(), data_value.get_size()));
191+
DEFINE_CAST_TO_JVALUE(TYPE_DATE, helper.newLocalDate(data_value._julian));
192+
DEFINE_CAST_TO_JVALUE(TYPE_DATETIME, helper.newLocalDateTime(data_value.timestamp()));
182193

183194
void release_jvalue(bool is_box, jvalue val) {
184195
if (is_box && val.l) {
@@ -243,6 +254,26 @@ StatusOr<jvalue> cast_to_jvalue(const TypeDescriptor& type_desc, bool is_boxed,
243254
v.l = helper.newString(slice.get_data(), slice.get_size());
244255
break;
245256
}
257+
<<<<<<< HEAD
258+
=======
259+
case TYPE_DATE: {
260+
auto spec_col = down_cast<const RunTimeColumnType<TYPE_DATE>*>(col);
261+
v.l = helper.newLocalDate(spec_col->immutable_data()[row_num]._julian);
262+
break;
263+
}
264+
case TYPE_DATETIME: {
265+
auto spec_col = down_cast<const RunTimeColumnType<TYPE_DATETIME>*>(col);
266+
v.l = helper.newLocalDateTime(spec_col->immutable_data()[row_num].timestamp());
267+
break;
268+
}
269+
case TYPE_DECIMAL32:
270+
case TYPE_DECIMAL64:
271+
case TYPE_DECIMAL128:
272+
case TYPE_DECIMAL256: {
273+
v.l = decimal_cell_to_bigdecimal(type_desc, col, row_num, helper);
274+
break;
275+
}
276+
>>>>>>> add5deeffc ([Enhancement] support DATE and DATETIME types in Java UDF (#72337))
246277
case TYPE_ARRAY: {
247278
auto spec_col = down_cast<const ArrayColumn*>(col);
248279
auto [offset, size] = spec_col->get_element_offset_size(row_num);
@@ -324,6 +355,18 @@ void assign_jvalue(MethodTypeDescriptor method_type_desc, Column* col, int row_n
324355
ASSIGN_BOX_TYPE(TYPE_BIGINT, int64_t)
325356
ASSIGN_BOX_TYPE(TYPE_FLOAT, float)
326357
ASSIGN_BOX_TYPE(TYPE_DOUBLE, double)
358+
case TYPE_DATE: {
359+
DateValue dv;
360+
dv._julian = helper.valLocalDate(val.l);
361+
down_cast<RunTimeColumnType<TYPE_DATE>*>(data_col)->get_data()[row_num] = dv;
362+
break;
363+
}
364+
case TYPE_DATETIME: {
365+
TimestampValue tv;
366+
tv.set_timestamp(helper.valLocalDateTime(val.l));
367+
down_cast<RunTimeColumnType<TYPE_DATETIME>*>(data_col)->get_data()[row_num] = tv;
368+
break;
369+
}
327370
case TYPE_VARCHAR: {
328371
if (val.l == nullptr) {
329372
col->append_nulls(1);
@@ -375,6 +418,18 @@ Status append_jvalue(const TypeDescriptor& type_desc, bool is_box, Column* col,
375418
APPEND_BOX_TYPE(TYPE_FLOAT, float)
376419
APPEND_BOX_TYPE(TYPE_DOUBLE, double)
377420

421+
case TYPE_DATE: {
422+
DateValue dv;
423+
dv._julian = helper.valLocalDate(val.l);
424+
col->append_datum(Datum(dv));
425+
break;
426+
}
427+
case TYPE_DATETIME: {
428+
TimestampValue tv;
429+
tv.set_timestamp(helper.valLocalDateTime(val.l));
430+
col->append_datum(Datum(tv));
431+
break;
432+
}
378433
case TYPE_VARCHAR: {
379434
std::string buffer;
380435
auto slice = helper.sliceVal((jstring)val.l, &buffer);
@@ -473,6 +528,39 @@ Status check_type_matched(const TypeDescriptor& type_desc, jobject val) {
473528
}
474529
break;
475530
}
531+
<<<<<<< HEAD
532+
=======
533+
case TYPE_DECIMAL32:
534+
case TYPE_DECIMAL64:
535+
case TYPE_DECIMAL128:
536+
case TYPE_DECIMAL256: {
537+
if (!env->IsInstanceOf(val, helper.big_decimal_class())) {
538+
auto clazz = env->GetObjectClass(val);
539+
LOCAL_REF_GUARD(clazz);
540+
return Status::InternalError(
541+
fmt::format("Type not matched, expect java.math.BigDecimal, but got {}", helper.to_string(clazz)));
542+
}
543+
break;
544+
}
545+
case TYPE_DATE: {
546+
if (!env->IsInstanceOf(val, helper.local_date_class())) {
547+
auto clazz = env->GetObjectClass(val);
548+
LOCAL_REF_GUARD(clazz);
549+
return Status::InternalError(
550+
fmt::format("Type not matched, expect java.time.LocalDate, but got {}", helper.to_string(clazz)));
551+
}
552+
break;
553+
}
554+
case TYPE_DATETIME: {
555+
if (!env->IsInstanceOf(val, helper.local_datetime_class())) {
556+
auto clazz = env->GetObjectClass(val);
557+
LOCAL_REF_GUARD(clazz);
558+
return Status::InternalError(fmt::format("Type not matched, expect java.time.LocalDateTime, but got {}",
559+
helper.to_string(clazz)));
560+
}
561+
break;
562+
}
563+
>>>>>>> add5deeffc ([Enhancement] support DATE and DATETIME types in Java UDF (#72337))
476564
case TYPE_ARRAY: {
477565
if (!env->IsInstanceOf(val, helper.list_meta().list_class->clazz())) {
478566
auto clazz = env->GetObjectClass(val);

‎be/src/udf/java/java_native_method.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,9 @@
2525
#include "column/nullable_column.h"
2626
#include "column/vectorized_fwd.h"
2727
#include "gutil/casts.h"
28+
#include "types/date_value.h"
2829
#include "types/logical_type.h"
30+
#include "types/timestamp_value.h"
2931

3032
namespace starrocks {
3133

@@ -116,6 +118,12 @@ class GetColumnLogicalTypeVistor : public ColumnVisitorAdapter<GetColumnLogicalT
116118
} else if constexpr (std::is_same_v<T, double>) {
117119
*_result = TYPE_DOUBLE;
118120
return Status::OK();
121+
} else if constexpr (std::is_same_v<T, DateValue>) {
122+
*_result = TYPE_DATE;
123+
return Status::OK();
124+
} else if constexpr (std::is_same_v<T, TimestampValue>) {
125+
*_result = TYPE_DATETIME;
126+
return Status::OK();
119127
}
120128
return Status::NotSupported("unsupported UDF type");
121129
}

‎be/src/udf/java/java_udf.cpp‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,9 @@
2525
#include "exprs/function_context.h"
2626
#include "fmt/core.h"
2727
#include "jni.h"
28+
#include "types/date_value.h"
2829
#include "types/logical_type.h"
30+
#include "types/timestamp_value.h"
2931
#include "udf/java/java_native_method.h"
3032
#include "udf/java/type_traits.h"
3133
#include "udf/java/utils.h"
@@ -184,6 +186,16 @@ void JVMFunctionHelper::_init() {
184186
CHECK(_jarrays_class);
185187
CHECK(_exception_util_class);
186188

189+
_local_date_class = JNI_FIND_CLASS("java/time/LocalDate");
190+
CHECK(_local_date_class);
191+
_local_date_of_epoch_day = _env->GetStaticMethodID(_local_date_class, "ofEpochDay", "(J)Ljava/time/LocalDate;");
192+
CHECK(_local_date_of_epoch_day);
193+
_local_date_to_epoch_day = _env->GetMethodID(_local_date_class, "toEpochDay", "()J");
194+
CHECK(_local_date_to_epoch_day);
195+
196+
_local_datetime_class = JNI_FIND_CLASS("java/time/LocalDateTime");
197+
CHECK(_local_datetime_class);
198+
187199
ADD_NUMBERIC_CLASS(boolean, Boolean, Z);
188200
ADD_NUMBERIC_CLASS(byte, Byte, B);
189201
ADD_NUMBERIC_CLASS(short, Short, S);
@@ -251,6 +263,15 @@ void JVMFunctionHelper::_init() {
251263
_method_map.emplace(TYPE_ARRAY_METHOD_ID, tmp);
252264
INIT_HELPER_METHOD(tmp, "createBoxedMapArray", CREATE_BOXED_MAP_SIGNATURE);
253265
_method_map.emplace(TYPE_MAP_METHOD_ID, tmp);
266+
INIT_HELPER_METHOD(tmp, "createBoxedLocalDateArray", CREATE_BOXED_PRIMI_SIGNATURE);
267+
_method_map.emplace(JNIPrimTypeId<DateValue>::id, tmp);
268+
INIT_HELPER_METHOD(tmp, "createBoxedLocalDateTimeArray", CREATE_BOXED_PRIMI_SIGNATURE);
269+
_method_map.emplace(JNIPrimTypeId<TimestampValue>::id, tmp);
270+
271+
// LocalDateTime <-> packed int64 conversion stays in UDFHelper because the
272+
// packing is StarRocks-specific.
273+
INIT_HELPER_METHOD(_local_datetime_from_packed, "localDateTimeFromPackedTimestamp", "(J)Ljava/time/LocalDateTime;");
274+
INIT_HELPER_METHOD(_local_datetime_to_packed, "packedTimestampFromLocalDateTime", "(Ljava/time/LocalDateTime;)J");
254275

255276
// init bytebuffer
256277
_direct_buffer_class = JNI_FIND_CLASS("java/nio/ByteBuffer");
@@ -537,6 +558,65 @@ jobject JVMFunctionHelper::newString(const char* data, size_t size) {
537558
return nstr;
538559
}
539560

561+
<<<<<<< HEAD
562+
=======
563+
jobject JVMFunctionHelper::newBigDecimal(const std::string& s) {
564+
jobject jstr = newString(s.data(), s.size());
565+
if (jstr == nullptr) {
566+
return nullptr;
567+
}
568+
LOCAL_REF_GUARD(jstr);
569+
jobject bd = _env->NewObject(_big_decimal_class, _big_decimal_ctor_string, jstr);
570+
RETURN_IF_JNI_EXCEPTION(_env, "newBigDecimal: NewObject failed", nullptr);
571+
return bd;
572+
}
573+
574+
jobject JVMFunctionHelper::newBigDecimal(int64_t unscaled, int scale) {
575+
jobject bd = _env->CallStaticObjectMethod(_big_decimal_class, _big_decimal_value_of_ll,
576+
static_cast<jlong>(unscaled), static_cast<jint>(scale));
577+
RETURN_IF_JNI_EXCEPTION(_env, "newBigDecimal(long, int): CallStatic failed", nullptr);
578+
return bd;
579+
}
580+
581+
jlong JVMFunctionHelper::unscaled_long(jobject big_decimal, int precision, int scale) {
582+
return _env->CallStaticLongMethod(_udf_helper_class, _bd_unscaled_long, big_decimal, static_cast<jint>(precision),
583+
static_cast<jint>(scale));
584+
}
585+
586+
jbyteArray JVMFunctionHelper::unscaled_le_bytes(jobject big_decimal, int precision, int scale, int byte_width) {
587+
return (jbyteArray)_env->CallStaticObjectMethod(_udf_helper_class, _bd_unscaled_le_bytes, big_decimal,
588+
static_cast<jint>(precision), static_cast<jint>(scale),
589+
static_cast<jint>(byte_width));
590+
}
591+
592+
// DateValue stores days as Julian day; LocalDate exposes them as days-since-1970-01-01.
593+
// 2440588 is the Julian day number of the Unix epoch.
594+
static constexpr jlong UNIX_EPOCH_JULIAN_DAYS = 2440588;
595+
596+
jobject JVMFunctionHelper::newLocalDate(int32_t julian) {
597+
jobject ld = _env->CallStaticObjectMethod(_local_date_class, _local_date_of_epoch_day,
598+
static_cast<jlong>(julian) - UNIX_EPOCH_JULIAN_DAYS);
599+
RETURN_IF_JNI_EXCEPTION(_env, "newLocalDate: ofEpochDay failed", nullptr);
600+
return ld;
601+
}
602+
603+
int32_t JVMFunctionHelper::valLocalDate(jobject obj) {
604+
jlong epoch_day = _env->CallLongMethod(obj, _local_date_to_epoch_day);
605+
return static_cast<int32_t>(epoch_day + UNIX_EPOCH_JULIAN_DAYS);
606+
}
607+
608+
jobject JVMFunctionHelper::newLocalDateTime(int64_t packed_timestamp) {
609+
jobject ldt = _env->CallStaticObjectMethod(_udf_helper_class, _local_datetime_from_packed,
610+
static_cast<jlong>(packed_timestamp));
611+
RETURN_IF_JNI_EXCEPTION(_env, "newLocalDateTime: localDateTimeFromPackedTimestamp failed", nullptr);
612+
return ldt;
613+
}
614+
615+
int64_t JVMFunctionHelper::valLocalDateTime(jobject obj) {
616+
return _env->CallStaticLongMethod(_udf_helper_class, _local_datetime_to_packed, obj);
617+
}
618+
619+
>>>>>>> add5deeffc ([Enhancement] support DATE and DATETIME types in Java UDF (#72337))
540620
Slice JVMFunctionHelper::sliceVal(jstring jstr, std::string* buffer) {
541621
const size_t utf_length = _env->GetStringUTFLength(jstr);
542622
buffer->resize(utf_length);
@@ -974,6 +1054,23 @@ Status ClassAnalyzer::get_udaf_method_desc(const std::string& sign, std::vector<
9741054
desc->emplace_back(MethodTypeDescriptor{TYPE_ARRAY, true});
9751055
} else if (type == "java/util/Map") {
9761056
desc->emplace_back(MethodTypeDescriptor{TYPE_MAP, true});
1057+
<<<<<<< HEAD
1058+
=======
1059+
} else if (type == "java/math/BigDecimal") {
1060+
// Structural placeholder only: `.type` is not dispatched on at runtime for
1061+
// BigDecimal params. The actual DECIMAL precision/scale is resolved from
1062+
// ctx->get_arg_type() / ctx->get_return_type() at the call site.
1063+
desc->emplace_back(MethodTypeDescriptor{TYPE_DECIMAL128, true});
1064+
} else if (type == "java/time/LocalDate") {
1065+
desc->emplace_back(MethodTypeDescriptor{TYPE_DATE, true});
1066+
} else if (type == "java/time/LocalDateTime") {
1067+
desc->emplace_back(MethodTypeDescriptor{TYPE_DATETIME, true});
1068+
} else {
1069+
// Unrecognized object class. Surface as TYPE_UNKNOWN so get_method_desc()
1070+
// validation produces a clear error instead of leaving method_desc short
1071+
// and crashing on a later method_desc[0].is_box access.
1072+
desc->emplace_back(MethodTypeDescriptor{TYPE_UNKNOWN, true});
1073+
>>>>>>> add5deeffc ([Enhancement] support DATE and DATETIME types in Java UDF (#72337))
9771074
}
9781075
continue;
9791076
}

‎be/src/udf/java/java_udf.h‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,22 @@ class JVMFunctionHelper {
162162

163163
Slice sliceVal(jstring jstr, std::string* buffer);
164164
jclass string_clazz() { return _string_class; }
165+
<<<<<<< HEAD
166+
=======
167+
jclass big_decimal_class() { return _big_decimal_class; }
168+
jclass local_date_class() { return _local_date_class; }
169+
jclass local_datetime_class() { return _local_datetime_class; }
170+
171+
// Box/unbox a StarRocks DateValue (int32 Julian day) as a java.time.LocalDate.
172+
// Round-trips StarRocks's internal Julian-day encoding.
173+
jobject newLocalDate(int32_t julian);
174+
int32_t valLocalDate(jobject obj);
175+
176+
// Box/unbox a StarRocks TimestampValue (packed int64: julian << 40 | usOfDay)
177+
// as a java.time.LocalDateTime.
178+
jobject newLocalDateTime(int64_t packed_timestamp);
179+
int64_t valLocalDateTime(jobject obj);
180+
>>>>>>> add5deeffc ([Enhancement] support DATE and DATETIME types in Java UDF (#72337))
165181
// replace '.' as '/'
166182
// eg: java.lang.Integer -> java/lang/Integer
167183
static std::string to_jni_class_name(const std::string& name);
@@ -197,8 +213,24 @@ class JVMFunctionHelper {
197213
jclass _string_class;
198214
jclass _jarrays_class;
199215
jclass _exception_util_class;
216+
<<<<<<< HEAD
200217

201218
jmethodID _string_construct_with_bytes;
219+
=======
220+
jclass _big_decimal_class;
221+
jclass _local_date_class;
222+
jclass _local_datetime_class;
223+
224+
jmethodID _string_construct_with_bytes;
225+
jmethodID _big_decimal_ctor_string;
226+
jmethodID _big_decimal_value_of_ll;
227+
// java.time.LocalDate.ofEpochDay(long) / java.time.LocalDate.toEpochDay()
228+
jmethodID _local_date_of_epoch_day;
229+
jmethodID _local_date_to_epoch_day;
230+
// UDFHelper.localDateTimeFromPackedTimestamp(long) / packedTimestampFromLocalDateTime(LocalDateTime)
231+
jmethodID _local_datetime_from_packed;
232+
jmethodID _local_datetime_to_packed;
233+
>>>>>>> add5deeffc ([Enhancement] support DATE and DATETIME types in Java UDF (#72337))
202234

203235
ListMeta _list_meta;
204236
MapMeta _map_meta;

‎be/src/udf/java/type_traits.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919
#include "util/slice.h"
2020

2121
namespace starrocks {
22+
class DateValue;
23+
class TimestampValue;
24+
2225
template <class Type>
2326
struct JNIPrimTypeId {
2427
static constexpr bool supported = false;
@@ -42,5 +45,9 @@ DEFINE_TYPE(double, 8);
4245
DEFINE_TYPE(Slice, 9);
4346
static constexpr int TYPE_ARRAY_METHOD_ID = 10;
4447
static constexpr int TYPE_MAP_METHOD_ID = 11;
48+
// DateValue/TimestampValue are PODs whose storage IS the wire format
49+
// (int32 Julian day / packed int64). Java UDFHelper decodes inline.
50+
DEFINE_TYPE(DateValue, 12);
51+
DEFINE_TYPE(TimestampValue, 13);
4552

4653
} // namespace starrocks

0 commit comments

Comments
 (0)