suxiaogang223 commented on code in PR #65381: URL: https://github.com/apache/doris/pull/65381#discussion_r3611909789
########## be/src/exec/sink/writer/paimon/jni_paimon_write_backend.cpp: ########## @@ -0,0 +1,363 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "exec/sink/writer/paimon/jni_paimon_write_backend.h" + +#include <arrow/buffer.h> +#include <arrow/io/memory.h> +#include <arrow/ipc/writer.h> +#include <arrow/record_batch.h> +#include <arrow/type.h> + +#include <map> + +#include "common/check.h" +#include "common/logging.h" +#include "format/arrow/arrow_block_convertor.h" +#include "format/arrow/arrow_row_batch.h" +#include "runtime/runtime_state.h" +#include "util/jni-util.h" + +namespace doris { + +// ──────────────────────────────────────────────────────────── +// JNI helpers — JVM attachment and class loading +// ──────────────────────────────────────────────────────────── + +static constexpr const char* PAIMON_JNI_WRITER_CLASS = "org/apache/doris/paimon/PaimonJniWriter"; +static constexpr const char* SCANNER_LOADER_CLASS = + "org/apache/doris/common/classloader/ScannerLoader"; + +/// Attach the current native thread to the JVM if not already attached, +/// and return a valid JNIEnv pointer. +static Status _get_jni_env(JNIEnv** env) { + JavaVM* jvm = nullptr; + jsize n_vms = 0; + jint result = JNI_GetCreatedJavaVMs(&jvm, 1, &n_vms); + if (result != JNI_OK || n_vms == 0) { + return Status::InternalError("Failed to get created JavaVM"); + } + result = jvm->GetEnv(reinterpret_cast<void**>(env), JNI_VERSION_1_8); + if (result == JNI_EDETACHED) { + result = jvm->AttachCurrentThread(reinterpret_cast<void**>(env), nullptr); + if (result != JNI_OK) { + return Status::InternalError("Failed to attach current thread to JVM"); + } + } else if (result != JNI_OK) { + return Status::InternalError("Failed to get JNIEnv"); + } + return Status::OK(); +} + +JniPaimonWriteBackend::~JniPaimonWriteBackend() { + JNIEnv* env = nullptr; + if (_get_jni_env(&env).ok()) { + if (_jni_writer_obj != nullptr) { + DCHECK(_close_id != nullptr); + env->CallVoidMethod(_jni_writer_obj, _close_id); + Status close_status = _check_jni_exception(env, "close PaimonJniWriter"); + if (!close_status.ok()) { + LOG(WARNING) << "Failed to close PaimonJniWriter: " << close_status.to_string(); + } + env->DeleteGlobalRef(_jni_writer_obj); + _jni_writer_obj = nullptr; + } + if (_jni_writer_cls != nullptr) { + env->DeleteGlobalRef(_jni_writer_cls); + _jni_writer_cls = nullptr; + } + } + _opened = false; +} + +Status JniPaimonWriteBackend::_check_jni_exception(JNIEnv* env, const std::string& method_name) { + if (env->ExceptionCheck()) { + Status st = + Jni::Env::GetJniExceptionMsg(env, true, "JNI exception in " + method_name + ": "); + LOG(WARNING) << st.to_string(); + return st; + } + return Status::OK(); +} + +Status JniPaimonWriteBackend::_load_writer_class(JNIEnv* env, jclass* writer_class) { + jclass loader_class = env->FindClass(SCANNER_LOADER_CLASS); + RETURN_IF_ERROR(_check_jni_exception(env, "find ScannerLoader")); + + jmethodID loader_constructor = env->GetMethodID(loader_class, "<init>", "()V"); + jmethodID get_loaded_class = env->GetMethodID(loader_class, "getLoadedClass", + "(Ljava/lang/String;)Ljava/lang/Class;"); + RETURN_IF_ERROR(_check_jni_exception(env, "resolve ScannerLoader methods")); + + jobject loader = env->NewObject(loader_class, loader_constructor); + jstring class_name = env->NewStringUTF(PAIMON_JNI_WRITER_CLASS); + auto* loaded_class = + static_cast<jclass>(env->CallObjectMethod(loader, get_loaded_class, class_name)); + RETURN_IF_ERROR(_check_jni_exception(env, "load PaimonJniWriter")); + + *writer_class = loaded_class; + env->DeleteLocalRef(class_name); + env->DeleteLocalRef(loader); + env->DeleteLocalRef(loader_class); + return Status::OK(); +} + +static jobject _to_java_options(JNIEnv* env, const std::map<std::string, std::string>& options) { + jclass map_cls = env->FindClass("java/util/HashMap"); + jmethodID map_ctor = env->GetMethodID(map_cls, "<init>", "()V"); + jmethodID put_method = env->GetMethodID( + map_cls, "put", "(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;"); + + jobject map_obj = env->NewObject(map_cls, map_ctor); + for (const auto& kv : options) { + jstring key = env->NewStringUTF(kv.first.c_str()); + jstring val = env->NewStringUTF(kv.second.c_str()); + env->CallObjectMethod(map_obj, put_method, key, val); + env->DeleteLocalRef(key); + env->DeleteLocalRef(val); + } + env->DeleteLocalRef(map_cls); + return map_obj; +} + +Status JniPaimonWriteBackend::open(const TPaimonTableSink& sink, RuntimeState* state) { + _sink = sink; + DORIS_CHECK(sink.__isset.column_names); + DORIS_CHECK(sink.__isset.write_mode); + DORIS_CHECK(sink.__isset.serialized_table); + DORIS_CHECK(!sink.serialized_table.empty()); + DORIS_CHECK(sink.__isset.transaction_id); + DORIS_CHECK(sink.transaction_id > 0); + DORIS_CHECK(sink.__isset.commit_user); + DORIS_CHECK(!sink.commit_user.empty()); + + JNIEnv* env = nullptr; + RETURN_IF_ERROR(_get_jni_env(&env)); + + // Step 1: Load PaimonJniWriter class through ScannerLoader (Paimon jars are + // not on the default application classpath, so FindClass won't work). + jclass local_cls = nullptr; + RETURN_IF_ERROR(_load_writer_class(env, &local_cls)); + _jni_writer_cls = static_cast<jclass>(env->NewGlobalRef(local_cls)); + env->DeleteLocalRef(local_cls); + + // Step 2: Cache JNI method IDs for write, prepareCommit, abort, close. + jmethodID open_id = env->GetMethodID( + _jni_writer_cls, "open", + "(Ljava/lang/String;Ljava/util/Map;[Ljava/lang/String;JLjava/lang/String;Z)V"); + _write_id = env->GetMethodID(_jni_writer_cls, "write", "(Ljava/nio/ByteBuffer;)V"); + _prepare_commit_id = env->GetMethodID(_jni_writer_cls, "prepareCommit", "()[[B"); + _abort_id = env->GetMethodID(_jni_writer_cls, "abort", "()V"); + _close_id = env->GetMethodID(_jni_writer_cls, "close", "()V"); + RETURN_IF_ERROR(_check_jni_exception(env, "GetMethodID")); + + // Step 3: Create the Java PaimonJniWriter instance. + jmethodID ctor_id = env->GetMethodID(_jni_writer_cls, "<init>", "()V"); + jobject local_obj = env->NewObject(_jni_writer_cls, ctor_id); + RETURN_IF_ERROR(_check_jni_exception(env, "NewObject")); + _jni_writer_obj = env->NewGlobalRef(local_obj); + env->DeleteLocalRef(local_obj); + + // Step 4: Build Java arguments and call PaimonJniWriter.open(). + const std::map<std::string, std::string> empty_config; + jstring j_serialized_table = env->NewStringUTF(sink.serialized_table.c_str()); + jobject j_hadoop_config = + _to_java_options(env, sink.__isset.hadoop_config ? sink.hadoop_config : empty_config); + jstring j_commit_user = env->NewStringUTF(sink.commit_user.c_str()); + + jclass string_cls = env->FindClass("java/lang/String"); + jobjectArray j_cols = + env->NewObjectArray(static_cast<jsize>(sink.column_names.size()), string_cls, nullptr); + for (size_t i = 0; i < sink.column_names.size(); ++i) { + jstring str = env->NewStringUTF(sink.column_names[i].c_str()); + env->SetObjectArrayElement(j_cols, static_cast<jsize>(i), str); + env->DeleteLocalRef(str); + } + + env->CallVoidMethod(_jni_writer_obj, open_id, j_serialized_table, j_hadoop_config, j_cols, + static_cast<jlong>(sink.transaction_id), j_commit_user, + static_cast<jboolean>(sink.write_mode == TPaimonWriteMode::OVERWRITE)); + Status st = _check_jni_exception(env, "open"); + + env->DeleteLocalRef(j_serialized_table); + env->DeleteLocalRef(j_hadoop_config); + env->DeleteLocalRef(j_commit_user); + env->DeleteLocalRef(j_cols); + env->DeleteLocalRef(string_cls); + + if (st.ok()) { + _opened = true; + } + return st; +} + +// Writer creation stays non-const because the backend interface also supports future stateful FFI +// implementations. +Status JniPaimonWriteBackend::create_writer( // NOLINT(readability-make-member-function-const) + std::unique_ptr<IPaimonWriter>* writer) { + DORIS_CHECK(_opened); + *writer = std::make_unique<JniPaimonWriter>(_jni_writer_obj, _write_id, _prepare_commit_id, + _abort_id, std::make_unique<ArrowMemoryPool<>>(), + _sink); + return Status::OK(); +} + +JniPaimonWriter::JniPaimonWriter(jobject jni_writer_obj, jmethodID write_id, + jmethodID prepare_commit_id, jmethodID abort_id, + std::unique_ptr<ArrowMemoryPool<>> arrow_pool, + TPaimonTableSink sink) + : _jni_writer_obj(jni_writer_obj), + _write_id(write_id), + _prepare_commit_id(prepare_commit_id), + _abort_id(abort_id), + _arrow_pool(std::move(arrow_pool)), + _sink(std::move(sink)) {} + +Status JniPaimonWriter::_write_projected_block(RuntimeState* state, Block& block) { + if (block.rows() == 0) { + return Status::OK(); + } + + // Use Thrift column_names as the authoritative schema source for both + // Arrow schema construction and Java-side write type derivation. + DORIS_CHECK(_sink.__isset.column_names); + DORIS_CHECK_EQ(_sink.column_names.size(), block.columns()); + for (size_t i = 0; i < _sink.column_names.size(); ++i) { + block.get_by_position(i).name = _sink.column_names[i]; + } + + // Pipeline: Doris Block → Arrow Schema → Arrow RecordBatch → IPC Stream → JNI direct buffer + // + // Step 1: Build Arrow schema from the projected Block. + std::shared_ptr<arrow::Schema> arrow_schema; + RETURN_IF_ERROR(get_arrow_schema_from_block(block, &arrow_schema, state->timezone())); + + // Step 2: Convert Doris Block columns to an Arrow RecordBatch. + std::shared_ptr<arrow::RecordBatch> record_batch; + RETURN_IF_ERROR(convert_to_arrow_batch(block, arrow_schema, _arrow_pool.get(), &record_batch, + state->timezone_obj())); + + // Step 3: Serialize the RecordBatch to Arrow IPC Stream format in memory. + auto out_stream_res = arrow::io::BufferOutputStream::Create(); Review Comment: Fixed in bc98f94c3e8. The IPC `BufferOutputStream` is now created with `_arrow_pool`, so its allocations go through the same task/query-tracked Arrow memory pool as the record batch. ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindSink.java: ########## @@ -878,6 +884,81 @@ private void validateStaticPartition(UnboundIcebergTableSink<?> sink, IcebergExt } } + private Plan bindPaimonTableSink(MatchingContext<UnboundPaimonTableSink<Plan>> ctx) { + UnboundPaimonTableSink<?> sink = ctx.root; + Pair<PaimonExternalDatabase, PaimonExternalTable> pair = bind(ctx.cascadesContext, sink); + PaimonExternalDatabase database = pair.first; + PaimonExternalTable table = pair.second; + LogicalPlan child = ((LogicalPlan) sink.child()); + + Map<String, Expression> staticPartitions = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); + staticPartitions.putAll(sink.getStaticPartitionKeyValues()); + Set<String> staticPartitionColNames = staticPartitions.keySet(); + if (!staticPartitionColNames.isEmpty()) { + Set<String> partitionColumnNames = Sets.newTreeSet(String.CASE_INSENSITIVE_ORDER); + partitionColumnNames.addAll(table.getPartitionColumnNames(Optional.empty())); + for (String columnName : staticPartitionColNames) { + if (!partitionColumnNames.contains(columnName)) { + throw new AnalysisException(String.format( + "Column '%s' is not a partition column of Paimon table '%s'", + columnName, table.getName())); + } + Expression partitionValue = staticPartitions.get(columnName); + if (!(partitionValue instanceof Literal)) { + throw new AnalysisException(String.format( + "Partition value for column '%s' must be a literal, but got: %s", + columnName, partitionValue)); + } + } + } + + List<Column> bindColumns; + if (sink.getColNames().isEmpty()) { + bindColumns = table.getBaseSchema(true).stream() + .filter(Column::isVisible) + .filter(column -> !staticPartitionColNames.contains(column.getName())) + .collect(ImmutableList.toImmutableList()); + } else { + bindColumns = sink.getColNames().stream().map(cn -> { + if (staticPartitionColNames.contains(cn)) { + throw new AnalysisException(String.format( + "Static partition column '%s' must not appear in the insert column list", cn)); + } + Column column = table.getColumn(cn); + if (column == null) { + throw new AnalysisException(String.format( + "column %s is not found in table %s", cn, table.getName())); + } + return column; + }).collect(ImmutableList.toImmutableList()); + } + + if (bindColumns.size() != child.getOutput().size()) { + throw new AnalysisException("insert into cols should be corresponding to the query output"); + } + Map<String, NamedExpression> columnToOutput = getSpecifiedColumnToOutput(bindColumns, child); Review Comment: Fixed in bc98f94c3e8. `BindSink` now rejects duplicate target names case-insensitively before building the mapping, and `PaimonWriteSchema` independently validates unique sink column names. Unit and regression coverage were added. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
