[Improvement](javaudf) support different date argument for date/datetime type (#13920)

This commit is contained in:
Gabriel
2022-11-03 20:33:20 +08:00
committed by GitHub
parent 5d7b51dcc2
commit 0a228a68d6
15 changed files with 1057 additions and 149 deletions

View File

@ -112,34 +112,37 @@ Status JavaFunctionCall::execute(FunctionContext* context, Block& block,
JniContext* jni_ctx = reinterpret_cast<JniContext*>(
context->get_function_state(FunctionContext::THREAD_LOCAL));
int arg_idx = 0;
ColumnPtr args[arguments.size()];
ColumnPtr data_cols[arguments.size()];
ColumnPtr null_cols[arguments.size()];
for (size_t col_idx : arguments) {
ColumnWithTypeAndName& column = block.get_by_position(col_idx);
args[arg_idx] = column.column->convert_to_full_column_if_const();
data_cols[arg_idx] = column.column->convert_to_full_column_if_const();
if (!_argument_types[arg_idx]->equals(*column.type)) {
return Status::InvalidArgument(strings::Substitute(
"$0-th input column's type $1 does not equal to required type $2", arg_idx,
column.type->get_name(), _argument_types[arg_idx]->get_name()));
}
if (auto* nullable = check_and_get_column<const ColumnNullable>(*args[arg_idx])) {
args[arg_idx] = nullable->get_nested_column_ptr();
auto null_col =
check_and_get_column<ColumnVector<UInt8>>(nullable->get_null_map_column_ptr());
jni_ctx->input_nulls_buffer_ptr.get()[arg_idx] =
reinterpret_cast<int64_t>(null_col->get_data().data());
if (auto* nullable = check_and_get_column<const ColumnNullable>(*data_cols[arg_idx])) {
null_cols[arg_idx] = nullable->get_null_map_column_ptr();
jni_ctx->input_nulls_buffer_ptr.get()[arg_idx] = reinterpret_cast<int64_t>(
check_and_get_column<ColumnVector<UInt8>>(null_cols[arg_idx])
->get_data()
.data());
data_cols[arg_idx] = nullable->get_nested_column_ptr();
} else {
jni_ctx->input_nulls_buffer_ptr.get()[arg_idx] = -1;
}
if (args[arg_idx]->is_column_string()) {
const ColumnString* str_col = assert_cast<const ColumnString*>(args[arg_idx].get());
if (data_cols[arg_idx]->is_column_string()) {
const ColumnString* str_col =
assert_cast<const ColumnString*>(data_cols[arg_idx].get());
jni_ctx->input_values_buffer_ptr.get()[arg_idx] =
reinterpret_cast<int64_t>(str_col->get_chars().data());
jni_ctx->input_offsets_ptrs.get()[arg_idx] =
reinterpret_cast<int64_t>(str_col->get_offsets().data());
} else if (args[arg_idx]->is_numeric() || args[arg_idx]->is_column_decimal()) {
} else if (data_cols[arg_idx]->is_numeric() || data_cols[arg_idx]->is_column_decimal()) {
jni_ctx->input_values_buffer_ptr.get()[arg_idx] =
reinterpret_cast<int64_t>(args[arg_idx]->get_raw_data().data);
reinterpret_cast<int64_t>(data_cols[arg_idx]->get_raw_data().data);
} else {
return Status::InvalidArgument(
strings::Substitute("Java UDF doesn't support type $0 now !",